Here are the examples of the java api class com.google.api.services.dataflow.model.Job taken from open source projects.
1. DataflowPipelineRunnerTest#testRunReturnDifferentRequestId()
Project: DataflowJavaSDK
File: DataflowPipelineRunnerTest.java
File: DataflowPipelineRunnerTest.java
@Test public void testRunReturnDifferentRequestId() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); Dataflow mockDataflowClient = options.getDataflowClient(); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); when(mockDataflowClient.projects().jobs().create(eq(PROJECT_ID), any(Job.class))).thenReturn(mockRequest); Job resultJob = new Job(); resultJob.setId("newid"); // Return a different request id. resultJob.setClientRequestId("different_request_id"); when(mockRequest.execute()).thenReturn(resultJob); DataflowPipeline p = buildDataflowPipeline(options); try { p.run(); fail("Expected DataflowJobAlreadyExistsException"); } catch (DataflowJobAlreadyExistsException expected) { assertThat(expected.getMessage(), containsString("If you want to submit a second job, try again by setting a " + "different name using --jobName.")); assertEquals(expected.getJob().getJobId(), resultJob.getId()); } }
2. DataflowRunnerTest#testRunReturnDifferentRequestId()
Project: incubator-beam
File: DataflowRunnerTest.java
File: DataflowRunnerTest.java
@Test public void testRunReturnDifferentRequestId() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); Dataflow mockDataflowClient = options.getDataflowClient(); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); when(mockDataflowClient.projects().jobs().create(eq(PROJECT_ID), any(Job.class))).thenReturn(mockRequest); Job resultJob = new Job(); resultJob.setId("newid"); // Return a different request id. resultJob.setClientRequestId("different_request_id"); when(mockRequest.execute()).thenReturn(resultJob); Pipeline p = buildDataflowPipeline(options); try { p.run(); fail("Expected DataflowJobAlreadyExistsException"); } catch (DataflowJobAlreadyExistsException expected) { assertThat(expected.getMessage(), containsString("If you want to submit a second job, try again by setting a " + "different name using --jobName.")); assertEquals(expected.getJob().getJobId(), resultJob.getId()); } }
3. DataflowPipelineJobTest#testGetStateReturnsServiceState()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetStateReturnsServiceState() throws Exception { Dataflow.Projects.Jobs.Get statusRequest = mock(Dataflow.Projects.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + State.RUNNING.name()); when(mockJobs.get(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); DataflowAggregatorTransforms dataflowAggregatorTransforms = mock(DataflowAggregatorTransforms.class); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, dataflowAggregatorTransforms); assertEquals(State.RUNNING, job.getStateWithRetries(DataflowPipelineJob.STATUS_POLLING_ATTEMPTS, fastClock)); }
4. DataflowPipelineJobTest#mockWaitToFinishInState()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
public State mockWaitToFinishInState(State state) throws Exception { Dataflow.Projects.Jobs.Get statusRequest = mock(Dataflow.Projects.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + state.name()); when(mockJobs.get(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); DataflowAggregatorTransforms dataflowAggregatorTransforms = mock(DataflowAggregatorTransforms.class); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, dataflowAggregatorTransforms); return job.waitToFinish(1, TimeUnit.MINUTES, null, fastClock, fastClock); }
5. DataflowPipelineJobTest#testWaitToFinishMessagesFail()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testWaitToFinishMessagesFail() throws Exception { Dataflow.Projects.Jobs.Get statusRequest = mock(Dataflow.Projects.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + State.DONE.name()); when(mockJobs.get(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); MonitoringUtil.JobMessagesHandler jobHandler = mock(MonitoringUtil.JobMessagesHandler.class); Dataflow.Projects.Jobs.Messages mockMessages = mock(Dataflow.Projects.Jobs.Messages.class); Messages.List listRequest = mock(Dataflow.Projects.Jobs.Messages.List.class); when(mockJobs.messages()).thenReturn(mockMessages); when(mockMessages.list(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(listRequest); when(listRequest.execute()).thenThrow(SocketTimeoutException.class); DataflowAggregatorTransforms dataflowAggregatorTransforms = mock(DataflowAggregatorTransforms.class); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, dataflowAggregatorTransforms); State state = job.waitToFinish(5, TimeUnit.MINUTES, jobHandler, fastClock, fastClock); assertEquals(null, state); }
6. DataflowPipelineTranslatorTest#buildMockDataflow()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
private static Dataflow buildMockDataflow(ArgumentMatcher<Job> jobMatcher) throws IOException { Dataflow mockDataflowClient = mock(Dataflow.class); Dataflow.Projects mockProjects = mock(Dataflow.Projects.class); Dataflow.Projects.Jobs mockJobs = mock(Dataflow.Projects.Jobs.class); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); when(mockDataflowClient.projects()).thenReturn(mockProjects); when(mockProjects.jobs()).thenReturn(mockJobs); when(mockJobs.create(eq("someProject"), argThat(jobMatcher))).thenReturn(mockRequest); Job resultJob = new Job(); resultJob.setId("newid"); when(mockRequest.execute()).thenReturn(resultJob); return mockDataflowClient; }
7. DataflowPipelineJobTest#testGetStateReturnsServiceState()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetStateReturnsServiceState() throws Exception { Dataflow.Projects.Jobs.Get statusRequest = mock(Dataflow.Projects.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + State.RUNNING.name()); when(mockJobs.get(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); DataflowAggregatorTransforms dataflowAggregatorTransforms = mock(DataflowAggregatorTransforms.class); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, dataflowAggregatorTransforms); assertEquals(State.RUNNING, job.getStateWithRetries(DataflowPipelineJob.STATUS_POLLING_ATTEMPTS, fastClock)); }
8. DataflowPipelineJobTest#mockWaitToFinishInState()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
public State mockWaitToFinishInState(State state) throws Exception { Dataflow.Projects.Jobs.Get statusRequest = mock(Dataflow.Projects.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + state.name()); when(mockJobs.get(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); DataflowAggregatorTransforms dataflowAggregatorTransforms = mock(DataflowAggregatorTransforms.class); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, dataflowAggregatorTransforms); return job.waitToFinish(1, TimeUnit.MINUTES, null, fastClock, fastClock); }
9. DataflowPipelineJobTest#testWaitToFinishMessagesFail()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testWaitToFinishMessagesFail() throws Exception { Dataflow.Projects.Jobs.Get statusRequest = mock(Dataflow.Projects.Jobs.Get.class); Job statusResponse = new Job(); statusResponse.setCurrentState("JOB_STATE_" + State.DONE.name()); when(mockJobs.get(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(statusRequest); when(statusRequest.execute()).thenReturn(statusResponse); MonitoringUtil.JobMessagesHandler jobHandler = mock(MonitoringUtil.JobMessagesHandler.class); Dataflow.Projects.Jobs.Messages mockMessages = mock(Dataflow.Projects.Jobs.Messages.class); Messages.List listRequest = mock(Dataflow.Projects.Jobs.Messages.List.class); when(mockJobs.messages()).thenReturn(mockMessages); when(mockMessages.list(eq(PROJECT_ID), eq(JOB_ID))).thenReturn(listRequest); when(listRequest.execute()).thenThrow(SocketTimeoutException.class); DataflowAggregatorTransforms dataflowAggregatorTransforms = mock(DataflowAggregatorTransforms.class); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, dataflowAggregatorTransforms); State state = job.waitToFinish(5, TimeUnit.MINUTES, jobHandler, fastClock, fastClock); assertEquals(null, state); }
10. DataflowPipelineTranslatorTest#buildMockDataflow()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
private static Dataflow buildMockDataflow(ArgumentMatcher<Job> jobMatcher) throws IOException { Dataflow mockDataflowClient = mock(Dataflow.class); Dataflow.Projects mockProjects = mock(Dataflow.Projects.class); Dataflow.Projects.Jobs mockJobs = mock(Dataflow.Projects.Jobs.class); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); when(mockDataflowClient.projects()).thenReturn(mockProjects); when(mockProjects.jobs()).thenReturn(mockJobs); when(mockJobs.create(eq("someProject"), argThat(jobMatcher))).thenReturn(mockRequest); Job resultJob = new Job(); resultJob.setId("newid"); when(mockRequest.execute()).thenReturn(resultJob); return mockDataflowClient; }
11. DataflowPipelineTranslatorTest#createPredefinedStep()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
/** * Returns a Step for a DoFn by creating and translating a pipeline. */ private static Step createPredefinedStep() throws Exception { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); Pipeline pipeline = Pipeline.create(options); String stepName = "DoFn1"; pipeline.apply("ReadMyFile", TextIO.Read.from("gs://bucket/in")).apply(stepName, ParDo.of(new NoOpFn())).apply("WriteMyFile", TextIO.Write.to("gs://bucket/out")); Job job = translator.translate(pipeline, (DataflowRunner) pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(13, job.getSteps().size()); Step step = job.getSteps().get(1); assertEquals(stepName, getString(step.getProperties(), PropertyNames.USER_NAME)); return step; }
12. DataflowPipelineTranslatorTest#testDiskSizeGbConfig()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testDiskSizeGbConfig() throws IOException { final Integer diskSizeGb = 1234; DataflowPipelineOptions options = buildPipelineOptions(); options.setDiskSizeGb(diskSizeGb); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(diskSizeGb, job.getEnvironment().getWorkerPools().get(0).getDiskSizeGb()); }
13. DataflowPipelineTranslatorTest#testWorkerMachineTypeConfig()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testWorkerMachineTypeConfig() throws IOException { final String testMachineType = "test-machine-type"; DataflowPipelineOptions options = buildPipelineOptions(); options.setWorkerMachineType(testMachineType); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); WorkerPool workerPool = job.getEnvironment().getWorkerPools().get(0); assertEquals(testMachineType, workerPool.getMachineType()); }
14. DataflowPipelineTranslatorTest#testZoneConfig()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testZoneConfig() throws IOException { final String testZone = "test-zone-1"; DataflowPipelineOptions options = buildPipelineOptions(); options.setZone(testZone); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(testZone, job.getEnvironment().getWorkerPools().get(0).getZone()); }
15. DataflowPipelineTranslatorTest#testSubnetworkConfigMissing()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testSubnetworkConfigMissing() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertNull(job.getEnvironment().getWorkerPools().get(0).getSubnetwork()); }
16. DataflowPipelineTranslatorTest#testSubnetworkConfig()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testSubnetworkConfig() throws IOException { final String testSubnetwork = "regions/REGION/subnetworks/SUBNETWORK"; DataflowPipelineOptions options = buildPipelineOptions(); options.setSubnetwork(testSubnetwork); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(testSubnetwork, job.getEnvironment().getWorkerPools().get(0).getSubnetwork()); }
17. DataflowPipelineTranslatorTest#testNetworkConfigMissing()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testNetworkConfigMissing() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertNull(job.getEnvironment().getWorkerPools().get(0).getNetwork()); }
18. DataflowPipelineTranslatorTest#testNetworkConfig()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testNetworkConfig() throws IOException { final String testNetwork = "test-network"; DataflowPipelineOptions options = buildPipelineOptions(); options.setNetwork(testNetwork); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(testNetwork, job.getEnvironment().getWorkerPools().get(0).getNetwork()); }
19. DataflowPipelineTranslatorTest#testToIterableTranslation()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToIterableTranslation() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<Iterable<T>> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); options.setExperiments(ImmutableList.of("disable_ism_side_input")); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); pipeline.apply(Create.of(1, 2, 3)).apply(View.<Integer>asIterable()); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(2, steps.size()); Step createStep = steps.get(0); assertEquals("CreateCollection", createStep.getKind()); Step collectionToSingletonStep = steps.get(1); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
20. DataflowPipelineTranslatorTest#createPredefinedStep()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
/** * Returns a Step for a DoFn by creating and translating a pipeline. */ private static Step createPredefinedStep() throws Exception { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); String stepName = "DoFn1"; pipeline.apply(TextIO.Read.named("ReadMyFile").from("gs://bucket/in")).apply(ParDo.of(new NoOpFn()).named(stepName)).apply(TextIO.Write.named("WriteMyFile").to("gs://bucket/out")); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(13, job.getSteps().size()); Step step = job.getSteps().get(1); assertEquals(stepName, getString(step.getProperties(), PropertyNames.USER_NAME)); return step; }
21. DataflowPipelineTranslatorTest#testDiskSizeGbConfig()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testDiskSizeGbConfig() throws IOException { final Integer diskSizeGb = 1234; DataflowPipelineOptions options = buildPipelineOptions(); options.setDiskSizeGb(diskSizeGb); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(diskSizeGb, job.getEnvironment().getWorkerPools().get(0).getDiskSizeGb()); }
22. DataflowPipelineTranslatorTest#testWorkerMachineTypeConfig()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testWorkerMachineTypeConfig() throws IOException { final String testMachineType = "test-machine-type"; DataflowPipelineOptions options = buildPipelineOptions(); options.setWorkerMachineType(testMachineType); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); WorkerPool workerPool = job.getEnvironment().getWorkerPools().get(0); assertEquals(testMachineType, workerPool.getMachineType()); }
23. DataflowPipelineTranslatorTest#testZoneConfig()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testZoneConfig() throws IOException { final String testZone = "test-zone-1"; DataflowPipelineOptions options = buildPipelineOptions(); options.setZone(testZone); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(testZone, job.getEnvironment().getWorkerPools().get(0).getZone()); }
24. DataflowPipelineTranslatorTest#testSubnetworkConfigMissing()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testSubnetworkConfigMissing() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertNull(job.getEnvironment().getWorkerPools().get(0).getSubnetwork()); }
25. DataflowPipelineTranslatorTest#testSubnetworkConfig()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testSubnetworkConfig() throws IOException { final String testSubnetwork = "regions/REGION/subnetworks/SUBNETWORK"; DataflowPipelineOptions options = buildPipelineOptions(); options.setSubnetwork(testSubnetwork); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(testSubnetwork, job.getEnvironment().getWorkerPools().get(0).getSubnetwork()); }
26. DataflowPipelineTranslatorTest#testNetworkConfigMissing()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testNetworkConfigMissing() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertNull(job.getEnvironment().getWorkerPools().get(0).getNetwork()); }
27. DataflowPipelineTranslatorTest#testNetworkConfig()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testNetworkConfig() throws IOException { final String testNetwork = "test-network"; DataflowPipelineOptions options = buildPipelineOptions(); options.setNetwork(testNetwork); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals(testNetwork, job.getEnvironment().getWorkerPools().get(0).getNetwork()); }
28. DataflowPipelineJob#cancel()
Project: incubator-beam
File: DataflowPipelineJob.java
File: DataflowPipelineJob.java
/** * Cancels the job. * @throws IOException if there is a problem executing the cancel request. */ public void cancel() throws IOException { Job content = new Job(); content.setProjectId(projectId); content.setId(jobId); content.setRequestedState("JOB_STATE_CANCELLED"); dataflowClient.projects().jobs().update(projectId, jobId, content).execute(); }
29. DataflowPipelineJob#cancel()
Project: DataflowJavaSDK
File: DataflowPipelineJob.java
File: DataflowPipelineJob.java
/** * Cancels the job. * @throws IOException if there is a problem executing the cancel request. */ public void cancel() throws IOException { Job content = new Job(); content.setProjectId(projectId); content.setId(jobId); content.setRequestedState("JOB_STATE_CANCELLED"); dataflowClient.projects().jobs().update(projectId, jobId, content).execute(); }
30. DataflowPipelineRunnerTest#testUpdateAlreadyUpdatedPipeline()
Project: DataflowJavaSDK
File: DataflowPipelineRunnerTest.java
File: DataflowPipelineRunnerTest.java
@Test public void testUpdateAlreadyUpdatedPipeline() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); options.setUpdate(true); options.setJobName("oldJobName"); Dataflow mockDataflowClient = options.getDataflowClient(); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); when(mockDataflowClient.projects().jobs().create(eq(PROJECT_ID), any(Job.class))).thenReturn(mockRequest); final Job resultJob = new Job(); resultJob.setId("newid"); // Return a different request id. resultJob.setClientRequestId("different_request_id"); when(mockRequest.execute()).thenReturn(resultJob); DataflowPipeline p = buildDataflowPipeline(options); thrown.expect(DataflowJobAlreadyUpdatedException.class); thrown.expect(new TypeSafeMatcher<DataflowJobAlreadyUpdatedException>() { @Override public void describeTo(Description description) { description.appendText("Expected job ID: " + resultJob.getId()); } @Override protected boolean matchesSafely(DataflowJobAlreadyUpdatedException item) { return resultJob.getId().equals(item.getJob().getJobId()); } }); thrown.expectMessage("The job named oldjobname with id: oldJobId has already been updated " + "into job id: newid and cannot be updated again."); p.run(); }
31. DataflowRunnerTest#testUpdateAlreadyUpdatedPipeline()
Project: incubator-beam
File: DataflowRunnerTest.java
File: DataflowRunnerTest.java
@Test public void testUpdateAlreadyUpdatedPipeline() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); options.setUpdate(true); options.setJobName("oldJobName"); Dataflow mockDataflowClient = options.getDataflowClient(); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); when(mockDataflowClient.projects().jobs().create(eq(PROJECT_ID), any(Job.class))).thenReturn(mockRequest); final Job resultJob = new Job(); resultJob.setId("newid"); // Return a different request id. resultJob.setClientRequestId("different_request_id"); when(mockRequest.execute()).thenReturn(resultJob); Pipeline p = buildDataflowPipeline(options); thrown.expect(DataflowJobAlreadyUpdatedException.class); thrown.expect(new TypeSafeMatcher<DataflowJobAlreadyUpdatedException>() { @Override public void describeTo(Description description) { description.appendText("Expected job ID: " + resultJob.getId()); } @Override protected boolean matchesSafely(DataflowJobAlreadyUpdatedException item) { return resultJob.getId().equals(item.getJob().getJobId()); } }); thrown.expectMessage("The job named oldjobname with id: oldJobId has already been updated " + "into job id: newid and cannot be updated again."); p.run(); }
32. DataflowRunnerTest#buildMockDataflow()
Project: incubator-beam
File: DataflowRunnerTest.java
File: DataflowRunnerTest.java
private static Dataflow buildMockDataflow(final ArgumentCaptor<Job> jobCaptor) throws IOException { Dataflow mockDataflowClient = mock(Dataflow.class); Dataflow.Projects mockProjects = mock(Dataflow.Projects.class); Dataflow.Projects.Jobs mockJobs = mock(Dataflow.Projects.Jobs.class); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); Dataflow.Projects.Jobs.List mockList = mock(Dataflow.Projects.Jobs.List.class); when(mockDataflowClient.projects()).thenReturn(mockProjects); when(mockProjects.jobs()).thenReturn(mockJobs); when(mockJobs.create(eq(PROJECT_ID), jobCaptor.capture())).thenReturn(mockRequest); when(mockJobs.list(eq(PROJECT_ID))).thenReturn(mockList); when(mockList.setPageToken(anyString())).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; }
33. DataflowPipelineJobTest#testGetAggregatorValuesWhenClientThrowsExceptionThrowsAggregatorRetrievalException()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWhenClientThrowsExceptionThrowsAggregatorRetrievalException() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); IOException cause = new IOException(); when(getMetrics.execute()).thenThrow(cause); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); thrown.expect(AggregatorRetrievalException.class); thrown.expectCause(is(cause)); thrown.expectMessage(aggregator.toString()); thrown.expectMessage("when retrieving Aggregator values for"); job.getAggregatorValues(aggregator); }
34. DataflowPipelineJobTest#testGetAggregatorValuesWithUnrelatedMetricUpdateIgnoresUpdate()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithUnrelatedMetricUpdateIgnoresUpdate() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); MetricUpdate ignoredUpdate = new MetricUpdate(); ignoredUpdate.setScalar(null); MetricStructuredName ignoredName = new MetricStructuredName(); ignoredName.setName("ignoredAggregator.elementCount.out0"); ignoredName.setContext(null); ignoredUpdate.setName(ignoredName); jobMetrics.setMetrics(ImmutableList.of(ignoredUpdate)); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<Long> values = job.getAggregatorValues(aggregator); assertThat(values.getValuesAtSteps().entrySet(), empty()); assertThat(values.getValues(), empty()); }
35. DataflowPipelineJobTest#testGetAggregatorValuesWithMultipleMetricUpdatesReturnsCollection()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithMultipleMetricUpdatesReturnsCollection() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); @SuppressWarnings("unchecked") PTransform<PInput, POutput> otherTransform = mock(PTransform.class); String otherStepName = "s88"; String otherFullName = "Spam/Ham/Eggs"; AppliedPTransform<?, ?, ?> otherAppliedTransform = appliedPTransform(otherFullName, otherTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform, aggregator, otherTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName, otherAppliedTransform, otherStepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); MetricUpdate updateOne = new MetricUpdate(); long stepValue = 1234L; updateOne.setScalar(new BigDecimal(stepValue)); MetricStructuredName structuredNameOne = new MetricStructuredName(); structuredNameOne.setName(aggregatorName); structuredNameOne.setContext(ImmutableMap.of("step", stepName)); updateOne.setName(structuredNameOne); MetricUpdate updateTwo = new MetricUpdate(); long stepValueTwo = 1024L; updateTwo.setScalar(new BigDecimal(stepValueTwo)); MetricStructuredName structuredNameTwo = new MetricStructuredName(); structuredNameTwo.setName(aggregatorName); structuredNameTwo.setContext(ImmutableMap.of("step", otherStepName)); updateTwo.setName(structuredNameTwo); jobMetrics.setMetrics(ImmutableList.of(updateOne, updateTwo)); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<Long> values = job.getAggregatorValues(aggregator); assertThat(values.getValuesAtSteps(), hasEntry(fullName, stepValue)); assertThat(values.getValuesAtSteps(), hasEntry(otherFullName, stepValueTwo)); assertThat(values.getValuesAtSteps().size(), equalTo(2)); assertThat(values.getValues(), containsInAnyOrder(stepValue, stepValueTwo)); assertThat(values.getTotalValue(combineFn), equalTo(Long.valueOf(stepValue + stepValueTwo))); }
36. DataflowPipelineJobTest#testGetAggregatorValuesWithSingleMetricUpdateReturnsSingletonCollection()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithSingleMetricUpdateReturnsSingletonCollection() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); MetricUpdate update = new MetricUpdate(); long stepValue = 1234L; update.setScalar(new BigDecimal(stepValue)); MetricStructuredName structuredName = new MetricStructuredName(); structuredName.setName(aggregatorName); structuredName.setContext(ImmutableMap.of("step", stepName)); update.setName(structuredName); jobMetrics.setMetrics(ImmutableList.of(update)); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<Long> values = job.getAggregatorValues(aggregator); assertThat(values.getValuesAtSteps(), hasEntry(fullName, stepValue)); assertThat(values.getValuesAtSteps().size(), equalTo(1)); assertThat(values.getValues(), contains(stepValue)); assertThat(values.getTotalValue(combineFn), equalTo(Long.valueOf(stepValue))); }
37. DataflowPipelineJobTest#testGetAggregatorValuesWithNullMetricUpdatesReturnsEmptyValue()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithNullMetricUpdatesReturnsEmptyValue() throws IOException, AggregatorRetrievalException { Aggregator<?, ?> aggregator = mock(Aggregator.class); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); jobMetrics.setMetrics(null); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<?> values = job.getAggregatorValues(aggregator); assertThat(values.getValues(), empty()); }
38. DataflowPipelineJobTest#testGetAggregatorValuesWithNoMetricUpdatesReturnsEmptyValue()
Project: incubator-beam
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithNoMetricUpdatesReturnsEmptyValue() throws IOException, AggregatorRetrievalException { Aggregator<?, ?> aggregator = mock(Aggregator.class); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); jobMetrics.setMetrics(ImmutableList.<MetricUpdate>of()); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<?> values = job.getAggregatorValues(aggregator); assertThat(values.getValues(), empty()); }
39. DataflowPipelineRunnerTest#buildMockDataflow()
Project: DataflowJavaSDK
File: DataflowPipelineRunnerTest.java
File: DataflowPipelineRunnerTest.java
private static Dataflow buildMockDataflow(final ArgumentCaptor<Job> jobCaptor) throws IOException { Dataflow mockDataflowClient = mock(Dataflow.class); Dataflow.Projects mockProjects = mock(Dataflow.Projects.class); Dataflow.Projects.Jobs mockJobs = mock(Dataflow.Projects.Jobs.class); Dataflow.Projects.Jobs.Create mockRequest = mock(Dataflow.Projects.Jobs.Create.class); Dataflow.Projects.Jobs.List mockList = mock(Dataflow.Projects.Jobs.List.class); when(mockDataflowClient.projects()).thenReturn(mockProjects); when(mockProjects.jobs()).thenReturn(mockJobs); when(mockJobs.create(eq(PROJECT_ID), jobCaptor.capture())).thenReturn(mockRequest); when(mockJobs.list(eq(PROJECT_ID))).thenReturn(mockList); when(mockList.setPageToken(anyString())).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; }
40. DataflowPipelineJobTest#testGetAggregatorValuesWhenClientThrowsExceptionThrowsAggregatorRetrievalException()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWhenClientThrowsExceptionThrowsAggregatorRetrievalException() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); IOException cause = new IOException(); when(getMetrics.execute()).thenThrow(cause); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); thrown.expect(AggregatorRetrievalException.class); thrown.expectCause(is(cause)); thrown.expectMessage(aggregator.toString()); thrown.expectMessage("when retrieving Aggregator values for"); job.getAggregatorValues(aggregator); }
41. DataflowPipelineJobTest#testGetAggregatorValuesWithUnrelatedMetricUpdateIgnoresUpdate()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithUnrelatedMetricUpdateIgnoresUpdate() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); MetricUpdate ignoredUpdate = new MetricUpdate(); ignoredUpdate.setScalar(null); MetricStructuredName ignoredName = new MetricStructuredName(); ignoredName.setName("ignoredAggregator.elementCount.out0"); ignoredName.setContext(null); ignoredUpdate.setName(ignoredName); jobMetrics.setMetrics(ImmutableList.of(ignoredUpdate)); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<Long> values = job.getAggregatorValues(aggregator); assertThat(values.getValuesAtSteps().entrySet(), empty()); assertThat(values.getValues(), empty()); }
42. DataflowPipelineJobTest#testGetAggregatorValuesWithMultipleMetricUpdatesReturnsCollection()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithMultipleMetricUpdatesReturnsCollection() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); @SuppressWarnings("unchecked") PTransform<PInput, POutput> otherTransform = mock(PTransform.class); String otherStepName = "s88"; String otherFullName = "Spam/Ham/Eggs"; AppliedPTransform<?, ?, ?> otherAppliedTransform = appliedPTransform(otherFullName, otherTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform, aggregator, otherTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName, otherAppliedTransform, otherStepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); MetricUpdate updateOne = new MetricUpdate(); long stepValue = 1234L; updateOne.setScalar(new BigDecimal(stepValue)); MetricStructuredName structuredNameOne = new MetricStructuredName(); structuredNameOne.setName(aggregatorName); structuredNameOne.setContext(ImmutableMap.of("step", stepName)); updateOne.setName(structuredNameOne); MetricUpdate updateTwo = new MetricUpdate(); long stepValueTwo = 1024L; updateTwo.setScalar(new BigDecimal(stepValueTwo)); MetricStructuredName structuredNameTwo = new MetricStructuredName(); structuredNameTwo.setName(aggregatorName); structuredNameTwo.setContext(ImmutableMap.of("step", otherStepName)); updateTwo.setName(structuredNameTwo); jobMetrics.setMetrics(ImmutableList.of(updateOne, updateTwo)); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<Long> values = job.getAggregatorValues(aggregator); assertThat(values.getValuesAtSteps(), hasEntry(fullName, stepValue)); assertThat(values.getValuesAtSteps(), hasEntry(otherFullName, stepValueTwo)); assertThat(values.getValuesAtSteps().size(), equalTo(2)); assertThat(values.getValues(), containsInAnyOrder(stepValue, stepValueTwo)); assertThat(values.getTotalValue(combineFn), equalTo(Long.valueOf(stepValue + stepValueTwo))); }
43. DataflowPipelineJobTest#testGetAggregatorValuesWithSingleMetricUpdateReturnsSingletonCollection()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithSingleMetricUpdateReturnsSingletonCollection() throws IOException, AggregatorRetrievalException { CombineFn<Long, long[], Long> combineFn = new Sum.SumLongFn(); String aggregatorName = "agg"; Aggregator<Long, Long> aggregator = new TestAggregator<>(combineFn, aggregatorName); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); MetricUpdate update = new MetricUpdate(); long stepValue = 1234L; update.setScalar(new BigDecimal(stepValue)); MetricStructuredName structuredName = new MetricStructuredName(); structuredName.setName(aggregatorName); structuredName.setContext(ImmutableMap.of("step", stepName)); update.setName(structuredName); jobMetrics.setMetrics(ImmutableList.of(update)); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<Long> values = job.getAggregatorValues(aggregator); assertThat(values.getValuesAtSteps(), hasEntry(fullName, stepValue)); assertThat(values.getValuesAtSteps().size(), equalTo(1)); assertThat(values.getValues(), contains(stepValue)); assertThat(values.getTotalValue(combineFn), equalTo(Long.valueOf(stepValue))); }
44. DataflowPipelineJobTest#testGetAggregatorValuesWithNullMetricUpdatesReturnsEmptyValue()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithNullMetricUpdatesReturnsEmptyValue() throws IOException, AggregatorRetrievalException { Aggregator<?, ?> aggregator = mock(Aggregator.class); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); jobMetrics.setMetrics(null); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<?> values = job.getAggregatorValues(aggregator); assertThat(values.getValues(), empty()); }
45. DataflowPipelineJobTest#testGetAggregatorValuesWithNoMetricUpdatesReturnsEmptyValue()
Project: DataflowJavaSDK
File: DataflowPipelineJobTest.java
File: DataflowPipelineJobTest.java
@Test public void testGetAggregatorValuesWithNoMetricUpdatesReturnsEmptyValue() throws IOException, AggregatorRetrievalException { Aggregator<?, ?> aggregator = mock(Aggregator.class); @SuppressWarnings("unchecked") PTransform<PInput, POutput> pTransform = mock(PTransform.class); String stepName = "s1"; String fullName = "Foo/Bar/Baz"; AppliedPTransform<?, ?, ?> appliedTransform = appliedPTransform(fullName, pTransform); DataflowAggregatorTransforms aggregatorTransforms = new DataflowAggregatorTransforms(ImmutableSetMultimap.<Aggregator<?, ?>, PTransform<?, ?>>of(aggregator, pTransform).asMap(), ImmutableMap.<AppliedPTransform<?, ?, ?>, String>of(appliedTransform, stepName)); GetMetrics getMetrics = mock(GetMetrics.class); when(mockJobs.getMetrics(PROJECT_ID, JOB_ID)).thenReturn(getMetrics); JobMetrics jobMetrics = new JobMetrics(); when(getMetrics.execute()).thenReturn(jobMetrics); jobMetrics.setMetrics(ImmutableList.<MetricUpdate>of()); Get getState = mock(Get.class); when(mockJobs.get(PROJECT_ID, JOB_ID)).thenReturn(getState); Job modelJob = new Job(); when(getState.execute()).thenReturn(modelJob); modelJob.setCurrentState(State.RUNNING.toString()); DataflowPipelineJob job = new DataflowPipelineJob(PROJECT_ID, JOB_ID, mockWorkflowClient, aggregatorTransforms); AggregatorValues<?> values = job.getAggregatorValues(aggregator); assertThat(values.getValues(), empty()); }
46. DataflowRunnerTest#testRunWithFiles()
Project: incubator-beam
File: DataflowRunnerTest.java
File: DataflowRunnerTest.java
@Test public void testRunWithFiles() throws IOException { // Test that the function DataflowRunner.stageFiles works as // expected. GcsUtil mockGcsUtil = buildMockGcsUtil(true); final String gcsStaging = "gs://somebucket/some/path"; final String gcsTemp = "gs://somebucket/some/temp/path"; final String cloudDataflowDataset = "somedataset"; // Create some temporary files. File temp1 = File.createTempFile("DataflowRunnerTest", "txt"); temp1.deleteOnExit(); File temp2 = File.createTempFile("DataflowRunnerTest2", "txt"); temp2.deleteOnExit(); String overridePackageName = "alias.txt"; ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class); DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); options.setFilesToStage(ImmutableList.of(temp1.getAbsolutePath(), overridePackageName + "=" + temp2.getAbsolutePath())); options.setStagingLocation(gcsStaging); options.setTempLocation(gcsTemp); options.setTempDatasetId(cloudDataflowDataset); options.setProject(PROJECT_ID); options.setJobName("job"); options.setDataflowClient(buildMockDataflow(jobCaptor)); options.setGcsUtil(mockGcsUtil); options.setGcpCredential(new TestCredential()); Pipeline p = buildDataflowPipeline(options); DataflowPipelineJob job = (DataflowPipelineJob) p.run(); assertEquals("newid", job.getJobId()); Job workflowJob = jobCaptor.getValue(); assertValidJob(workflowJob); assertEquals(2, workflowJob.getEnvironment().getWorkerPools().get(0).getPackages().size()); DataflowPackage workflowPackage1 = workflowJob.getEnvironment().getWorkerPools().get(0).getPackages().get(0); assertThat(workflowPackage1.getName(), startsWith(temp1.getName())); DataflowPackage workflowPackage2 = workflowJob.getEnvironment().getWorkerPools().get(0).getPackages().get(1); assertEquals(overridePackageName, workflowPackage2.getName()); assertEquals("storage.googleapis.com/somebucket/some/temp/path", workflowJob.getEnvironment().getTempStoragePrefix()); assertEquals(cloudDataflowDataset, workflowJob.getEnvironment().getDataset()); assertEquals(ReleaseInfo.getReleaseInfo().getName(), workflowJob.getEnvironment().getUserAgent().get("name")); assertEquals(ReleaseInfo.getReleaseInfo().getVersion(), workflowJob.getEnvironment().getUserAgent().get("version")); }
47. DataflowPipelineTranslatorTest#testStepDisplayData()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testStepDisplayData() throws Exception { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); Pipeline pipeline = Pipeline.create(options); DoFn<Integer, Integer> fn1 = new DoFn<Integer, Integer>() { @Override public void processElement(ProcessContext c) throws Exception { c.output(c.element()); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "bar")).add(DisplayData.item("foo2", DataflowPipelineTranslatorTest.class).withLabel("Test Class").withLinkUrl("http://www.google.com")); } }; DoFn<Integer, Integer> fn2 = new DoFn<Integer, Integer>() { @Override public void processElement(ProcessContext c) throws Exception { c.output(c.element()); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo3", 1234)); } }; ParDo.Bound<Integer, Integer> parDo1 = ParDo.of(fn1); ParDo.Bound<Integer, Integer> parDo2 = ParDo.of(fn2); pipeline.apply(Create.of(1, 2, 3)).apply(parDo1).apply(parDo2); Job job = translator.translate(pipeline, (DataflowRunner) pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(3, steps.size()); Map<String, Object> parDo1Properties = steps.get(1).getProperties(); Map<String, Object> parDo2Properties = steps.get(2).getProperties(); assertThat(parDo1Properties, hasKey("display_data")); @SuppressWarnings("unchecked") Collection<Map<String, String>> fn1displayData = (Collection<Map<String, String>>) parDo1Properties.get("display_data"); @SuppressWarnings("unchecked") Collection<Map<String, String>> fn2displayData = (Collection<Map<String, String>>) parDo2Properties.get("display_data"); ImmutableSet<ImmutableMap<String, Object>> expectedFn1DisplayData = ImmutableSet.of(ImmutableMap.<String, Object>builder().put("key", "foo").put("type", "STRING").put("value", "bar").put("namespace", fn1.getClass().getName()).build(), ImmutableMap.<String, Object>builder().put("key", "fn").put("label", "Transform Function").put("type", "JAVA_CLASS").put("value", fn1.getClass().getName()).put("shortValue", fn1.getClass().getSimpleName()).put("namespace", parDo1.getClass().getName()).build(), ImmutableMap.<String, Object>builder().put("key", "foo2").put("type", "JAVA_CLASS").put("value", DataflowPipelineTranslatorTest.class.getName()).put("shortValue", DataflowPipelineTranslatorTest.class.getSimpleName()).put("namespace", fn1.getClass().getName()).put("label", "Test Class").put("linkUrl", "http://www.google.com").build()); ImmutableSet<ImmutableMap<String, Object>> expectedFn2DisplayData = ImmutableSet.of(ImmutableMap.<String, Object>builder().put("key", "fn").put("label", "Transform Function").put("type", "JAVA_CLASS").put("value", fn2.getClass().getName()).put("shortValue", fn2.getClass().getSimpleName()).put("namespace", parDo2.getClass().getName()).build(), ImmutableMap.<String, Object>builder().put("key", "foo3").put("type", "INTEGER").put("value", 1234L).put("namespace", fn2.getClass().getName()).build()); assertEquals(expectedFn1DisplayData, ImmutableSet.copyOf(fn1displayData)); assertEquals(expectedFn2DisplayData, ImmutableSet.copyOf(fn2displayData)); }
48. DataflowPipelineTranslatorTest#testToIterableTranslationWithIsmSideInput()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToIterableTranslationWithIsmSideInput() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<Iterable<T>> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); Pipeline pipeline = Pipeline.create(options); pipeline.apply(Create.of(1, 2, 3)).apply(View.<Integer>asIterable()); Job job = translator.translate(pipeline, (DataflowRunner) pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(3, steps.size()); @SuppressWarnings("unchecked") List<Map<String, Object>> toIsmRecordOutputs = (List<Map<String, Object>>) steps.get(1).getProperties().get(PropertyNames.OUTPUT_INFO); assertTrue(Structs.getBoolean(Iterables.getOnlyElement(toIsmRecordOutputs), "use_indexed_format")); Step collectionToSingletonStep = steps.get(2); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
49. DataflowPipelineTranslatorTest#testToSingletonTranslationWithIsmSideInput()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToSingletonTranslationWithIsmSideInput() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<T> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); Pipeline pipeline = Pipeline.create(options); pipeline.apply(Create.of(1)).apply(View.<Integer>asSingleton()); Job job = translator.translate(pipeline, (DataflowRunner) pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(5, steps.size()); @SuppressWarnings("unchecked") List<Map<String, Object>> toIsmRecordOutputs = (List<Map<String, Object>>) steps.get(3).getProperties().get(PropertyNames.OUTPUT_INFO); assertTrue(Structs.getBoolean(Iterables.getOnlyElement(toIsmRecordOutputs), "use_indexed_format")); Step collectionToSingletonStep = steps.get(4); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
50. DataflowPipelineTranslatorTest#testToIterableTranslation()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToIterableTranslation() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<Iterable<T>> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); options.setExperiments(ImmutableList.of("disable_ism_side_input")); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); Pipeline pipeline = Pipeline.create(options); pipeline.apply(Create.of(1, 2, 3)).apply(View.<Integer>asIterable()); Job job = translator.translate(pipeline, (DataflowRunner) pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(2, steps.size()); Step createStep = steps.get(0); assertEquals("ParallelRead", createStep.getKind()); Step collectionToSingletonStep = steps.get(1); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
51. DataflowPipelineTranslatorTest#testToSingletonTranslation()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToSingletonTranslation() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<T> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); options.setExperiments(ImmutableList.of("disable_ism_side_input")); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); Pipeline pipeline = Pipeline.create(options); pipeline.apply(Create.of(1)).apply(View.<Integer>asSingleton()); Job job = translator.translate(pipeline, (DataflowRunner) pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(2, steps.size()); Step createStep = steps.get(0); assertEquals("ParallelRead", createStep.getKind()); Step collectionToSingletonStep = steps.get(1); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
52. DataflowPipelineTranslatorTest#testPredefinedAddStep()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testPredefinedAddStep() throws Exception { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipelineTranslator.registerTransformTranslator(EmbeddedTransform.class, new EmbeddedTranslator()); // Create a predefined step using another pipeline Step predefinedStep = createPredefinedStep(); // Create a pipeline that the predefined step will be embedded into Pipeline pipeline = Pipeline.create(options); pipeline.apply("ReadMyFile", TextIO.Read.from("gs://bucket/in")).apply(ParDo.of(new NoOpFn())).apply(new EmbeddedTransform(predefinedStep.clone())).apply(ParDo.of(new NoOpFn())); Job job = translator.translate(pipeline, (DataflowRunner) pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(4, steps.size()); // The input to the embedded step should match the output of the step before Map<String, Object> step1Out = getOutputPortReference(steps.get(1)); Map<String, Object> step2In = getDictionary(steps.get(2).getProperties(), PropertyNames.PARALLEL_INPUT); assertEquals(step1Out, step2In); // The output from the embedded step should match the input of the step after Map<String, Object> step2Out = getOutputPortReference(steps.get(2)); Map<String, Object> step3In = getDictionary(steps.get(3).getProperties(), PropertyNames.PARALLEL_INPUT); assertEquals(step2Out, step3In); // The step should not have been modified other than remapping the input Step predefinedStepClone = predefinedStep.clone(); Step embeddedStepClone = steps.get(2).clone(); predefinedStepClone.getProperties().remove(PropertyNames.PARALLEL_INPUT); embeddedStepClone.getProperties().remove(PropertyNames.PARALLEL_INPUT); assertEquals(predefinedStepClone, embeddedStepClone); }
53. DataflowPipelineTranslatorTest#testMaxNumWorkersIsPassedWhenNoAlgorithmIsSet()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testMaxNumWorkersIsPassedWhenNoAlgorithmIsSet() throws IOException { final DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType noScaling = null; DataflowPipelineOptions options = buildPipelineOptions(); options.setMaxNumWorkers(42); options.setAutoscalingAlgorithm(noScaling); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertNull(job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); assertEquals(42, job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getMaxNumWorkers().intValue()); }
54. DataflowPipelineTranslatorTest#testScalingAlgorithmNone()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testScalingAlgorithmNone() throws IOException { final DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType noScaling = DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType.NONE; DataflowPipelineOptions options = buildPipelineOptions(); options.setAutoscalingAlgorithm(noScaling); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals("AUTOSCALING_ALGORITHM_NONE", job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); assertEquals(0, job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getMaxNumWorkers().intValue()); }
55. DataflowPipelineTranslatorTest#testScalingAlgorithmMissing()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testScalingAlgorithmMissing() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); // Autoscaling settings are always set. assertNull(job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); assertEquals(0, job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getMaxNumWorkers().intValue()); }
56. DataflowPipelineTranslatorTest#testSettingOfSdkPipelineOptions()
Project: incubator-beam
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testSettingOfSdkPipelineOptions() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); options.setRunner(DataflowRunner.class); Pipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, (DataflowRunner) p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); // Note that the contents of this materialized map may be changed by the act of reading an // option, which will cause the default to get materialized whereas it would otherwise be // left absent. It is permissible to simply alter this test to reflect current behavior. Map<String, Object> settings = new HashMap<>(); settings.put("appName", "DataflowPipelineTranslatorTest"); settings.put("project", "some-project"); settings.put("pathValidatorClass", "org.apache.beam.runners.dataflow.util.DataflowPathValidator"); settings.put("runner", "org.apache.beam.runners.dataflow.DataflowRunner"); settings.put("jobName", "some-job-name"); settings.put("tempLocation", "gs://somebucket/some/path"); settings.put("stagingLocation", "gs://somebucket/some/path/staging"); settings.put("stableUniqueNames", "WARNING"); settings.put("streaming", false); settings.put("numberOfWorkerHarnessThreads", 0); settings.put("experiments", null); Map<String, Object> sdkPipelineOptions = job.getEnvironment().getSdkPipelineOptions(); assertThat(sdkPipelineOptions, hasKey("options")); assertEquals(settings, sdkPipelineOptions.get("options")); }
57. DataflowPipelineTranslator#translate()
Project: incubator-beam
File: DataflowPipelineTranslator.java
File: DataflowPipelineTranslator.java
/** * Translates a {@link Pipeline} into a {@code JobSpecification}. */ public JobSpecification translate(Pipeline pipeline, DataflowRunner runner, List<DataflowPackage> packages) { Translator translator = new Translator(pipeline, runner); Job result = translator.translate(packages); return new JobSpecification(result, Collections.unmodifiableMap(translator.stepNames)); }
58. DataflowPipelineTranslatorTest#testStepDisplayData()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testStepDisplayData() throws Exception { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); DoFn<Integer, Integer> fn1 = new DoFn<Integer, Integer>() { @Override public void processElement(ProcessContext c) throws Exception { c.output(c.element()); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo", "bar")).add(DisplayData.item("foo2", DataflowPipelineTranslatorTest.class).withLabel("Test Class").withLinkUrl("http://www.google.com")); } }; DoFn<Integer, Integer> fn2 = new DoFn<Integer, Integer>() { @Override public void processElement(ProcessContext c) throws Exception { c.output(c.element()); } @Override public void populateDisplayData(DisplayData.Builder builder) { builder.add(DisplayData.item("foo3", 1234)); } }; ParDo.Bound<Integer, Integer> parDo1 = ParDo.of(fn1); ParDo.Bound<Integer, Integer> parDo2 = ParDo.of(fn2); pipeline.apply(Create.of(1, 2, 3)).apply(parDo1).apply(parDo2); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(3, steps.size()); Map<String, Object> parDo1Properties = steps.get(1).getProperties(); Map<String, Object> parDo2Properties = steps.get(2).getProperties(); assertThat(parDo1Properties, hasKey("display_data")); @SuppressWarnings("unchecked") Collection<Map<String, String>> fn1displayData = (Collection<Map<String, String>>) parDo1Properties.get("display_data"); @SuppressWarnings("unchecked") Collection<Map<String, String>> fn2displayData = (Collection<Map<String, String>>) parDo2Properties.get("display_data"); ImmutableSet<ImmutableMap<String, Object>> expectedFn1DisplayData = ImmutableSet.of(ImmutableMap.<String, Object>builder().put("key", "foo").put("type", "STRING").put("value", "bar").put("namespace", fn1.getClass().getName()).build(), ImmutableMap.<String, Object>builder().put("key", "fn").put("label", "Transform Function").put("type", "JAVA_CLASS").put("value", fn1.getClass().getName()).put("shortValue", fn1.getClass().getSimpleName()).put("namespace", parDo1.getClass().getName()).build(), ImmutableMap.<String, Object>builder().put("key", "foo2").put("type", "JAVA_CLASS").put("value", DataflowPipelineTranslatorTest.class.getName()).put("shortValue", DataflowPipelineTranslatorTest.class.getSimpleName()).put("namespace", fn1.getClass().getName()).put("label", "Test Class").put("linkUrl", "http://www.google.com").build()); ImmutableSet<ImmutableMap<String, Object>> expectedFn2DisplayData = ImmutableSet.of(ImmutableMap.<String, Object>builder().put("key", "fn").put("label", "Transform Function").put("type", "JAVA_CLASS").put("value", fn2.getClass().getName()).put("shortValue", fn2.getClass().getSimpleName()).put("namespace", parDo2.getClass().getName()).build(), ImmutableMap.<String, Object>builder().put("key", "foo3").put("type", "INTEGER").put("value", 1234L).put("namespace", fn2.getClass().getName()).build()); assertEquals(expectedFn1DisplayData, ImmutableSet.copyOf(fn1displayData)); assertEquals(expectedFn2DisplayData, ImmutableSet.copyOf(fn2displayData)); }
59. DataflowPipelineTranslatorTest#testToIterableTranslationWithIsmSideInput()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToIterableTranslationWithIsmSideInput() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<Iterable<T>> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); pipeline.apply(Create.of(1, 2, 3)).apply(View.<Integer>asIterable()); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(3, steps.size()); @SuppressWarnings("unchecked") List<Map<String, Object>> toIsmRecordOutputs = (List<Map<String, Object>>) steps.get(1).getProperties().get(PropertyNames.OUTPUT_INFO); assertTrue(Structs.getBoolean(Iterables.getOnlyElement(toIsmRecordOutputs), "use_indexed_format")); Step collectionToSingletonStep = steps.get(2); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
60. DataflowPipelineTranslatorTest#testToSingletonTranslationWithIsmSideInput()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToSingletonTranslationWithIsmSideInput() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<T> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); pipeline.apply(Create.of(1)).apply(View.<Integer>asSingleton()); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(5, steps.size()); @SuppressWarnings("unchecked") List<Map<String, Object>> toIsmRecordOutputs = (List<Map<String, Object>>) steps.get(3).getProperties().get(PropertyNames.OUTPUT_INFO); assertTrue(Structs.getBoolean(Iterables.getOnlyElement(toIsmRecordOutputs), "use_indexed_format")); Step collectionToSingletonStep = steps.get(4); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
61. DataflowPipelineTranslatorTest#testPubsubIOWriteTranslation()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testPubsubIOWriteTranslation() throws Exception { // A "change detector" test that makes sure the translation of PubsubIO.Read // does not change in bad ways during refactor String project = "fakeproject"; String topicName = "faketopic"; PubsubTopic topic = PubsubTopic.fromPath(String.format("projects/%s/topics/%s", project, topicName)); DataflowPipelineOptions options = buildPipelineOptions(); options.setStreaming(true); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); pipeline.apply(Create.of("hello", "goodbye")).apply(PubsubIO.Write.topic(topic.asPath())); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); // There are 6 steps from the Create and 1 from the PubsubIO.Write List<Step> steps = job.getSteps(); assertEquals(7, steps.size()); Step prepStep = steps.get(5); Step writeStep = steps.get(6); // Ensure the core translation does not change; if we get this far, then the right translator // was invoked, so we don't try to combinatorially test every combination of configuration // options. That would be for the unit tests of that translator. assertEquals("ParallelWrite", writeStep.getKind()); assertEquals("pubsub", Structs.getString(writeStep.getProperties(), PropertyNames.FORMAT)); assertEquals(topic.asV1Beta1Path(), Structs.getString(writeStep.getProperties(), PropertyNames.PUBSUB_TOPIC)); assertEquals("PubsubIO.Write/StreamingPubsubIOWrite", Structs.getString(writeStep.getProperties(), PropertyNames.USER_NAME)); }
62. DataflowPipelineTranslatorTest#testPubsubIOReadTranslation()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testPubsubIOReadTranslation() throws Exception { // A "change detector" test that makes sure the translation of PubsubIO.Read // does not change in bad ways during refactor String project = "fakeproject"; String topicName = "faketopic"; PubsubTopic topic = PubsubTopic.fromPath(String.format("projects/%s/topics/%s", project, topicName)); DataflowPipelineOptions options = buildPipelineOptions(); options.setStreaming(true); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); pipeline.apply(PubsubIO.Read.topic(topic.asPath())); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(1, steps.size()); Step pubsubStep = steps.get(0); // Ensure the core translation does not change; if we get this far, then the right translator // was invoked, so we don't try to combinatorially test every combination of configuration // options. That would be for the unit tests of that translator. assertEquals("ParallelRead", pubsubStep.getKind()); assertEquals("pubsub", Structs.getString(pubsubStep.getProperties(), PropertyNames.FORMAT)); assertEquals(topic.asV1Beta1Path(), Structs.getString(pubsubStep.getProperties(), PropertyNames.PUBSUB_TOPIC)); assertEquals("PubsubIO.Read", Structs.getString(pubsubStep.getProperties(), PropertyNames.USER_NAME)); }
63. DataflowPipelineTranslatorTest#testToSingletonTranslation()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testToSingletonTranslation() throws Exception { // A "change detector" test that makes sure the translation // of getting a PCollectionView<T> does not change // in bad ways during refactor DataflowPipelineOptions options = buildPipelineOptions(); options.setExperiments(ImmutableList.of("disable_ism_side_input")); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipeline pipeline = DataflowPipeline.create(options); pipeline.apply(Create.of(1)).apply(View.<Integer>asSingleton()); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(2, steps.size()); Step createStep = steps.get(0); assertEquals("CreateCollection", createStep.getKind()); Step collectionToSingletonStep = steps.get(1); assertEquals("CollectionToSingleton", collectionToSingletonStep.getKind()); }
64. DataflowPipelineTranslatorTest#testPredefinedAddStep()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testPredefinedAddStep() throws Exception { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipelineTranslator translator = DataflowPipelineTranslator.fromOptions(options); DataflowPipelineTranslator.registerTransformTranslator(EmbeddedTransform.class, new EmbeddedTranslator()); // Create a predefined step using another pipeline Step predefinedStep = createPredefinedStep(); // Create a pipeline that the predefined step will be embedded into DataflowPipeline pipeline = DataflowPipeline.create(options); pipeline.apply(TextIO.Read.named("ReadMyFile").from("gs://bucket/in")).apply(ParDo.of(new NoOpFn())).apply(new EmbeddedTransform(predefinedStep.clone())).apply(ParDo.of(new NoOpFn())); Job job = translator.translate(pipeline, pipeline.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); List<Step> steps = job.getSteps(); assertEquals(4, steps.size()); // The input to the embedded step should match the output of the step before Map<String, Object> step1Out = getOutputPortReference(steps.get(1)); Map<String, Object> step2In = getDictionary(steps.get(2).getProperties(), PropertyNames.PARALLEL_INPUT); assertEquals(step1Out, step2In); // The output from the embedded step should match the input of the step after Map<String, Object> step2Out = getOutputPortReference(steps.get(2)); Map<String, Object> step3In = getDictionary(steps.get(3).getProperties(), PropertyNames.PARALLEL_INPUT); assertEquals(step2Out, step3In); // The step should not have been modified other than remapping the input Step predefinedStepClone = predefinedStep.clone(); Step embeddedStepClone = steps.get(2).clone(); predefinedStepClone.getProperties().remove(PropertyNames.PARALLEL_INPUT); embeddedStepClone.getProperties().remove(PropertyNames.PARALLEL_INPUT); assertEquals(predefinedStepClone, embeddedStepClone); }
65. DataflowPipelineTranslatorTest#testMaxNumWorkersIsPassedWhenNoAlgorithmIsSet()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testMaxNumWorkersIsPassedWhenNoAlgorithmIsSet() throws IOException { final DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType noScaling = null; DataflowPipelineOptions options = buildPipelineOptions(); options.setMaxNumWorkers(42); options.setAutoscalingAlgorithm(noScaling); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertNull(job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); assertEquals(42, job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getMaxNumWorkers().intValue()); }
66. DataflowPipelineTranslatorTest#testScalingAlgorithmNone()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testScalingAlgorithmNone() throws IOException { final DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType noScaling = DataflowPipelineWorkerPoolOptions.AutoscalingAlgorithmType.NONE; DataflowPipelineOptions options = buildPipelineOptions(); options.setAutoscalingAlgorithm(noScaling); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); assertEquals("AUTOSCALING_ALGORITHM_NONE", job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); assertEquals(0, job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getMaxNumWorkers().intValue()); }
67. DataflowPipelineTranslatorTest#testScalingAlgorithmMissing()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testScalingAlgorithmMissing() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); assertEquals(1, job.getEnvironment().getWorkerPools().size()); // Autoscaling settings are always set. assertNull(job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getAlgorithm()); assertEquals(0, job.getEnvironment().getWorkerPools().get(0).getAutoscalingSettings().getMaxNumWorkers().intValue()); }
68. DataflowPipelineTranslatorTest#testSettingOfSdkPipelineOptions()
Project: DataflowJavaSDK
File: DataflowPipelineTranslatorTest.java
File: DataflowPipelineTranslatorTest.java
@Test public void testSettingOfSdkPipelineOptions() throws IOException { DataflowPipelineOptions options = buildPipelineOptions(); options.setRunner(DataflowPipelineRunner.class); DataflowPipeline p = buildPipeline(options); p.traverseTopologically(new RecordingPipelineVisitor()); Job job = DataflowPipelineTranslator.fromOptions(options).translate(p, p.getRunner(), Collections.<DataflowPackage>emptyList()).getJob(); // Note that the contents of this materialized map may be changed by the act of reading an // option, which will cause the default to get materialized whereas it would otherwise be // left absent. It is permissible to simply alter this test to reflect current behavior. Map<String, Object> settings = new HashMap<>(); settings.put("appName", "DataflowPipelineTranslatorTest"); settings.put("project", "some-project"); settings.put("pathValidatorClass", "com.google.cloud.dataflow.sdk.util.DataflowPathValidator"); settings.put("runner", "com.google.cloud.dataflow.sdk.runners.DataflowPipelineRunner"); settings.put("jobName", "some-job-name"); settings.put("tempLocation", "gs://somebucket/some/path"); settings.put("stagingLocation", "gs://somebucket/some/path/staging"); settings.put("stableUniqueNames", "WARNING"); settings.put("streaming", false); settings.put("numberOfWorkerHarnessThreads", 0); settings.put("experiments", null); Map<String, Object> sdkPipelineOptions = job.getEnvironment().getSdkPipelineOptions(); assertThat(sdkPipelineOptions, hasKey("options")); assertEquals(settings, sdkPipelineOptions.get("options")); }
69. DataflowPipelineRunnerTest#testRunWithFiles()
Project: DataflowJavaSDK
File: DataflowPipelineRunnerTest.java
File: DataflowPipelineRunnerTest.java
@Test public void testRunWithFiles() throws IOException { // Test that the function DataflowPipelineRunner.stageFiles works as // expected. GcsUtil mockGcsUtil = buildMockGcsUtil(true); final String gcsStaging = "gs://somebucket/some/path"; final String gcsTemp = "gs://somebucket/some/temp/path"; final String cloudDataflowDataset = "somedataset"; // Create some temporary files. File temp1 = File.createTempFile("DataflowPipelineRunnerTest", "txt"); temp1.deleteOnExit(); File temp2 = File.createTempFile("DataflowPipelineRunnerTest2", "txt"); temp2.deleteOnExit(); String overridePackageName = "alias.txt"; ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class); DataflowPipelineOptions options = PipelineOptionsFactory.as(DataflowPipelineOptions.class); options.setFilesToStage(ImmutableList.of(temp1.getAbsolutePath(), overridePackageName + "=" + temp2.getAbsolutePath())); options.setStagingLocation(gcsStaging); options.setTempLocation(gcsTemp); options.setTempDatasetId(cloudDataflowDataset); options.setProject(PROJECT_ID); options.setJobName("job"); options.setDataflowClient(buildMockDataflow(jobCaptor)); options.setGcsUtil(mockGcsUtil); options.setGcpCredential(new TestCredential()); DataflowPipeline p = buildDataflowPipeline(options); DataflowPipelineJob job = p.run(); assertEquals("newid", job.getJobId()); Job workflowJob = jobCaptor.getValue(); assertValidJob(workflowJob); assertEquals(2, workflowJob.getEnvironment().getWorkerPools().get(0).getPackages().size()); DataflowPackage workflowPackage1 = workflowJob.getEnvironment().getWorkerPools().get(0).getPackages().get(0); assertThat(workflowPackage1.getName(), startsWith(temp1.getName())); DataflowPackage workflowPackage2 = workflowJob.getEnvironment().getWorkerPools().get(0).getPackages().get(1); assertEquals(overridePackageName, workflowPackage2.getName()); assertEquals("storage.googleapis.com/somebucket/some/temp/path", workflowJob.getEnvironment().getTempStoragePrefix()); assertEquals(cloudDataflowDataset, workflowJob.getEnvironment().getDataset()); assertEquals(DataflowReleaseInfo.getReleaseInfo().getName(), workflowJob.getEnvironment().getUserAgent().get("name")); assertEquals(DataflowReleaseInfo.getReleaseInfo().getVersion(), workflowJob.getEnvironment().getUserAgent().get("version")); }
70. DataflowPipelineTranslator#translate()
Project: DataflowJavaSDK
File: DataflowPipelineTranslator.java
File: DataflowPipelineTranslator.java
/** * Translates a {@link Pipeline} into a {@code JobSpecification}. */ public JobSpecification translate(Pipeline pipeline, DataflowPipelineRunner runner, List<DataflowPackage> packages) { Translator translator = new Translator(pipeline, runner); Job result = translator.translate(packages); return new JobSpecification(result, Collections.unmodifiableMap(translator.stepNames)); }