use of com.google.api.services.dataflow.model.Job in project beam by apache.
the class DataflowRunnerTest method testSettingFlexRS.
@Test
public void testSettingFlexRS() throws IOException {
DataflowPipelineOptions options = buildPipelineOptions();
options.setFlexRSGoal(DataflowPipelineOptions.FlexResourceSchedulingGoal.COST_OPTIMIZED);
Pipeline p = Pipeline.create(options);
p.run();
ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture());
assertEquals("FLEXRS_COST_OPTIMIZED", jobCaptor.getValue().getEnvironment().getFlexResourceSchedulingGoal());
}
use of com.google.api.services.dataflow.model.Job in project beam by apache.
the class DataflowRunnerTest method testNonExistentStagingLocation.
@Test
public void testNonExistentStagingLocation() throws IOException {
DataflowPipelineOptions options = buildPipelineOptions();
options.setStagingLocation(NON_EXISTENT_BUCKET);
thrown.expect(IllegalArgumentException.class);
thrown.expectMessage(containsString("Output path does not exist or is not writeable: " + NON_EXISTENT_BUCKET));
DataflowRunner.fromOptions(options);
ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture());
assertValidJob(jobCaptor.getValue());
}
use of com.google.api.services.dataflow.model.Job in project beam by apache.
the class DataflowRunnerTest method buildMockDataflow.
private static Dataflow buildMockDataflow(Dataflow.Projects.Locations.Jobs mockJobs) throws IOException {
Dataflow mockDataflowClient = mock(Dataflow.class);
Dataflow.Projects mockProjects = mock(Dataflow.Projects.class);
Dataflow.Projects.Locations mockLocations = mock(Dataflow.Projects.Locations.class);
Dataflow.Projects.Locations.Jobs.Create mockRequest = mock(Dataflow.Projects.Locations.Jobs.Create.class);
Dataflow.Projects.Locations.Jobs.List mockList = mock(Dataflow.Projects.Locations.Jobs.List.class);
when(mockDataflowClient.projects()).thenReturn(mockProjects);
when(mockProjects.locations()).thenReturn(mockLocations);
when(mockLocations.jobs()).thenReturn(mockJobs);
when(mockJobs.create(eq(PROJECT_ID), eq(REGION_ID), isA(Job.class))).thenReturn(mockRequest);
when(mockJobs.list(eq(PROJECT_ID), eq(REGION_ID))).thenReturn(mockList);
when(mockList.setPageToken(any())).thenReturn(mockList);
when(mockList.execute()).thenReturn(new ListJobsResponse().setJobs(Arrays.asList(new Job().setName("oldjobname").setId("oldJobId").setCurrentState("JOB_STATE_RUNNING"))));
Job resultJob = new Job();
resultJob.setId("newid");
when(mockRequest.execute()).thenReturn(resultJob);
return mockDataflowClient;
}
use of com.google.api.services.dataflow.model.Job in project beam by apache.
the class DataflowRunnerTest method testTransformTranslatorMissing.
@Test
public void testTransformTranslatorMissing() throws IOException {
DataflowPipelineOptions options = buildPipelineOptions();
Pipeline p = Pipeline.create(options);
p.apply(Create.of(Arrays.asList(1, 2, 3))).apply(new TestTransform());
thrown.expect(IllegalStateException.class);
thrown.expectMessage(containsString("no translator registered"));
SdkComponents sdkComponents = SdkComponents.create(options);
RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p, sdkComponents, true);
DataflowPipelineTranslator.fromOptions(options).translate(p, pipelineProto, sdkComponents, DataflowRunner.fromOptions(options), Collections.emptyList());
ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture());
assertValidJob(jobCaptor.getValue());
}
use of com.google.api.services.dataflow.model.Job in project beam by apache.
the class DataflowRunnerTest method testRun.
@Test
public void testRun() throws IOException {
DataflowPipelineOptions options = buildPipelineOptions();
Pipeline p = buildDataflowPipeline(options);
DataflowPipelineJob job = (DataflowPipelineJob) p.run();
assertEquals("newid", job.getJobId());
ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class);
Mockito.verify(mockJobs).create(eq(PROJECT_ID), eq(REGION_ID), jobCaptor.capture());
assertValidJob(jobCaptor.getValue());
}
Aggregations