DataflowPipelineJob implements PipelineResult { @Override public @Nullable State waitUntilFinish() { return waitUntilFinish(Duration.millis(-1)); } DataflowPipelineJob(
      DataflowClient dataflowClient,
      String jobId,
      DataflowPipelineOptions dataflowOptions,
      Map<AppliedPTransform<?, ?, ?>, String> transformStepNames); String getJobId(); String getProjectId(); DataflowPipelineOptions getDataflowOptions(); String getRegion(); DataflowPipelineJob getReplacedByJob(); @Override @Nullable State waitUntilFinish(); @Override @Nullable State waitUntilFinish(Duration duration); @Nullable @VisibleForTesting State waitUntilFinish(Duration duration, MonitoringUtil.JobMessagesHandler messageHandler); @Override State cancel(); @Override State getState(); @Override MetricResults metrics();  }
@Test public void testWaitToFinishMessagesFail() throws Exception { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + State.DONE.name()); when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); MonitoringUtil.JobMessagesHandler jobHandler = mock(MonitoringUtil.JobMessagesHandler.class); Dataflow.Projects.Locations.Jobs.Messages mockMessages = mock(Dataflow.Projects.Locations.Jobs.Messages.class); Messages.List listRequest = mock(Dataflow.Projects.Locations.Jobs.Messages.List.class); when(mockJobs.messages()).thenReturn(mockMessages); when(mockMessages.list(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(listRequest); when(listRequest.setPageToken(eq((String) null))).thenReturn(listRequest); when(listRequest.execute()).thenThrow(SocketTimeoutException.class); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of()); State state = job.waitUntilFinish(Duration.standardMinutes(5), jobHandler, fastClock, fastClock); assertEquals(null, state); }
@Test public void testWaitToFinishFail() throws Exception { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenThrow(IOException.class); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of()); long startTime = fastClock.nanoTime(); State state = job.waitUntilFinish(Duration.standardMinutes(5), null, fastClock, fastClock); assertEquals(null, state); long timeDiff = TimeUnit.NANOSECONDS.toMillis(fastClock.nanoTime() - startTime); checkValidInterval( DataflowPipelineJob.MESSAGES_POLLING_INTERVAL, DataflowPipelineJob.MESSAGES_POLLING_RETRIES, timeDiff); }
@Test public void testWaitToFinishTimeFail() throws Exception { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenThrow(IOException.class); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of()); long startTime = fastClock.nanoTime(); State state = job.waitUntilFinish(Duration.millis(4), null, fastClock, fastClock); assertEquals(null, state); long timeDiff = TimeUnit.NANOSECONDS.toMillis(fastClock.nanoTime() - startTime); assertEquals(4L, timeDiff); }
@Test public void testCumulativeTimeOverflow() throws Exception { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_RUNNING"); when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); FastNanoClockAndFuzzySleeper clock = new FastNanoClockAndFuzzySleeper(); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of()); long startTime = clock.nanoTime(); State state = job.waitUntilFinish(Duration.millis(4), null, clock, clock); assertEquals(null, state); long timeDiff = TimeUnit.NANOSECONDS.toMillis(clock.nanoTime() - startTime); assertThat(timeDiff, lessThanOrEqualTo(4L)); }
@Test public void testWaitUntilFinishNoRepeatedLogs() throws Exception { DataflowPipelineJob job = new DataflowPipelineJob(mockDataflowClient, JOB_ID, options, null); Sleeper sleeper = new ZeroSleeper(); NanoClock nanoClock = mock(NanoClock.class); Instant separatingTimestamp = new Instant(42L); JobMessage theMessage = infoMessage(separatingTimestamp, "nothing"); MonitoringUtil mockMonitor = mock(MonitoringUtil.class); when(mockMonitor.getJobMessages(anyString(), anyLong())) .thenReturn(ImmutableList.of(theMessage)); Job fakeJob = new Job(); fakeJob.setCurrentState("JOB_STATE_RUNNING"); when(mockDataflowClient.getJob(anyString())).thenReturn(fakeJob); when(nanoClock.nanoTime()).thenReturn(0L).thenReturn(2000000000L); job.waitUntilFinish(Duration.standardSeconds(1), mockHandler, sleeper, nanoClock, mockMonitor); verify(mockHandler).process(ImmutableList.of(theMessage)); when(nanoClock.nanoTime()).thenReturn(3000000000L).thenReturn(6000000000L); job.waitUntilFinish(Duration.standardSeconds(1), mockHandler, sleeper, nanoClock, mockMonitor); verify(mockMonitor).getJobMessages(anyString(), eq(separatingTimestamp.getMillis())); }