com.google.api.services.dataflow.model.Job

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
@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
@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
@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
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
@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
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
@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
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
@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
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
/**
   * 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
@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
@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
@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
@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
@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
@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
@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
@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
/**
   * 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
@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
@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
@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
@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
@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
@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
@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
/**
   * 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
/**
   * 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
@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
@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
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
@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
@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
@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
@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
@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
@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
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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
/**
   * 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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
@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
/**
   * 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));
}