@Test
public void testExportBigQueryToGcsSingleShardAwaitCompletion() throws Exception {

  when(mockBigqueryTablesGet.execute()).thenReturn(fakeTable);
  final ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
  doAnswer(
          new Answer<Job>() {
            @Override
            public Job answer(InvocationOnMock invocationOnMock) throws Throwable {
              verify(mockBigqueryJobs, times(1)).insert(eq(jobProjectId), jobCaptor.capture());
              return jobCaptor.getValue();
            }
          })
      .when(mockBigqueryJobsInsert)
      .execute();
  when(mockBigqueryJobsGet.execute()).thenReturn(jobHandle);
  // Run exportBigQueryToGCS method.
  helper.exportBigQueryToGcs(jobProjectId, tableRef, ImmutableList.of("test-export-path"), true);
  // Verify correct calls to BigQuery are made.
  verify(mockBigquery, times(2)).jobs();
  // Verify correct calls to BigQuery.Jobs are made.
  verify(mockBigqueryJobs, times(1))
      .get(eq(jobProjectId), eq(jobCaptor.getValue().getJobReference().getJobId()));
  Job job = jobCaptor.getValue();
  assertThat(job.getConfiguration().getExtract().getDestinationUris().get(0))
      .isEqualTo("test-export-path");
  assertThat(job.getConfiguration().getExtract().getSourceTable()).isEqualTo(tableRef);
  assertThat(job.getJobReference().getLocation()).isEqualTo("test_location");
  // Verify we poll for job in correct location
  verify(mockBigqueryJobsGet).setLocation(eq("test_location"));
  // Verify correct calls to BigQuery.Jobs.Get are made.
  verify(mockBigqueryJobsGet, times(1)).execute();
  // Verify correct calls to BigQuery.Jobs.Insert are made.
  verify(mockBigqueryJobsInsert, times(1)).execute();
  verify(mockBigquery).tables();
  verify(mockBigqueryTables).get(eq(projectId), eq(datasetId), eq(tableId));
  verify(mockBigqueryTablesGet).execute();
}