@Test
public void getSplits() throws IOException {

  JobContext jobContext = new JobContextImpl(config, new JobID());
  CreateReadSessionRequest request =
      CreateReadSessionRequest.newBuilder()
          .setTableReference(
              TableReferenceProto.TableReference.newBuilder()
                  .setProjectId("publicdata")
                  .setDatasetId("test_dataset")
                  .setTableId("test_table"))
          .setRequestedStreams(3) // request 3, but only get 2 back
          .setParent("projects/foo-project")
          .setReadOptions(
              TableReadOptions.newBuilder()
                  .addAllSelectedFields(ImmutableList.of("foo", "bar"))
                  .setRowRestriction("foo == 0")
                  .build())
          .setFormat(DataFormat.AVRO)
          .build();
  ReadSession session =
      ReadSession.newBuilder()
          .setAvroSchema(AvroSchema.newBuilder().setSchema("schema").build())
          .addAllStreams(
              ImmutableList.of(
                  Stream.newBuilder().setName("stream1").build(),
                  Stream.newBuilder().setName("stream2").build()))
          .build();
  ImmutableList<DirectBigQueryInputSplit> expected =
      ImmutableList.of(
          new DirectBigQueryInputSplit("stream1", "schema", 14),
          new DirectBigQueryInputSplit("stream2", "schema", 14));
  when(bqClient.createReadSession(any(CreateReadSessionRequest.class))).thenReturn(session);
  try {
    List<InputSplit> splits = input.getSplits(jobContext);
    assertThat(splits).containsExactlyElementsIn(expected);
  } catch (Exception e) {
    e.printStackTrace();
    System.exit(1);
  }
  verify(bqHelper).getTable(tableRef);
  verify(bqClient).createReadSession(request);
}