@VisibleForTesting protected DataflowRunner(DataflowPipelineOptions options) { this.options = options; this.dataflowClient = DataflowClient.create(options); this.translator = DataflowPipelineTranslator.fromOptions(options); this.pcollectionsRequiringIndexedFormat = new HashSet<>(); this.ptransformViewsWithNonDeterministicKeyCoders = new HashSet<>(); }
DataflowClient.create(options), jobResult.getId(), options,
@Test public void testCancelUnterminatedJobThatSucceeds() throws IOException { Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenReturn(new Job().setCurrentState("JOB_STATE_CANCELLED")); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); assertEquals(State.CANCELLED, job.cancel()); Job content = new Job(); content.setProjectId(PROJECT_ID); content.setId(JOB_ID); content.setRequestedState("JOB_STATE_CANCELLED"); verify(mockJobs).update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), eq(content)); verifyNoMoreInteractions(mockJobs); }
public State mockWaitToFinishInState(State state) throws Exception { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + state.name()); if (state == State.UPDATED) { statusResponse.setReplacedByJobId(REPLACEMENT_JOB_ID); } when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of()); return job.waitUntilFinish(Duration.standardMinutes(1), null, fastClock, fastClock); }
@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); // Should only have slept for the 4 ms allowed. assertEquals(4L, timeDiff); }
/** * Test that {@link DataflowPipelineJob#cancel} doesn't throw if the Dataflow service returns * non-terminal state even though the cancel API call failed, which can happen in practice. * * <p>TODO: delete this code if the API calls become consistent. */ @Test public void testCancelTerminatedJobWithStaleState() throws IOException { 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(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenThrow(new IOException("Job has terminated in state SUCCESS")); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); State returned = job.cancel(); assertThat(returned, equalTo(State.RUNNING)); expectedLogs.verifyWarn("Cancel failed because job is already terminated."); }
@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 testCancelUnterminatedJobThatFails() throws IOException { 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(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenThrow(new IOException("Some random IOException")); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); thrown.expect(IOException.class); thrown.expectMessage( "Failed to cancel job in state RUNNING, " + "please go to the Developers Console to cancel it manually:"); job.cancel(); }
@Test public void testCancelTerminatedJob() throws IOException { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_FAILED"); when(mockJobs.get(PROJECT_ID, REGION_ID, JOB_ID)).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); Dataflow.Projects.Locations.Jobs.Update update = mock(Dataflow.Projects.Locations.Jobs.Update.class); when(mockJobs.update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), any(Job.class))) .thenReturn(update); when(update.execute()).thenThrow(new IOException()); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, null); assertEquals(State.FAILED, job.cancel()); Job content = new Job(); content.setProjectId(PROJECT_ID); content.setId(JOB_ID); content.setRequestedState("JOB_STATE_CANCELLED"); verify(mockJobs).update(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID), eq(content)); verify(mockJobs).get(PROJECT_ID, REGION_ID, JOB_ID); verifyNoMoreInteractions(mockJobs); }
/** Constructs a runner from the provided options. */ public static TestDataflowRunner fromOptions(PipelineOptions options) { TestDataflowPipelineOptions dataflowOptions = options.as(TestDataflowPipelineOptions.class); String tempLocation = Joiner.on("/") .join(dataflowOptions.getTempRoot(), dataflowOptions.getJobName(), "output", "results"); dataflowOptions.setTempLocation(tempLocation); return new TestDataflowRunner( dataflowOptions, DataflowClient.create(options.as(DataflowPipelineOptions.class))); }
@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 testGetStateReturnsServiceState() throws Exception { Dataflow.Projects.Locations.Jobs.Get statusRequest = mock(Dataflow.Projects.Locations.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + State.RUNNING.name()); when(mockJobs.get(eq(PROJECT_ID), eq(REGION_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); DataflowPipelineJob job = new DataflowPipelineJob(DataflowClient.create(options), JOB_ID, options, ImmutableMap.of()); assertEquals( State.RUNNING, job.getStateWithRetries( BackOffAdapter.toGcpBackOff(DataflowPipelineJob.STATUS_BACKOFF_FACTORY.backoff()), fastClock)); }
@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); // Should only have slept for the 4 ms allowed. assertThat(timeDiff, lessThanOrEqualTo(4L)); }
@Test public void testGetStateWithExceptionReturnsUnknown() 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(); assertEquals( State.UNKNOWN, job.getStateWithRetries( BackOffAdapter.toGcpBackOff(DataflowPipelineJob.STATUS_BACKOFF_FACTORY.backoff()), fastClock)); long timeDiff = TimeUnit.NANOSECONDS.toMillis(fastClock.nanoTime() - startTime); checkValidInterval( DataflowPipelineJob.STATUS_POLLING_INTERVAL, DataflowPipelineJob.STATUS_POLLING_RETRIES, timeDiff); }