Java Code Examples for org.apache.beam.runners.dataflow.options.DataflowPipelineOptions#getDataflowClient()

The following examples show how to use org.apache.beam.runners.dataflow.options.DataflowPipelineOptions#getDataflowClient() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: DataflowRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up 
@Test
public void testRunReturnDifferentRequestId() throws IOException {
  DataflowPipelineOptions options = buildPipelineOptions();
  Dataflow mockDataflowClient = options.getDataflowClient();
  Dataflow.Projects.Locations.Jobs.Create mockRequest =
      mock(Dataflow.Projects.Locations.Jobs.Create.class);
  when(mockDataflowClient
          .projects()
          .locations()
          .jobs()
          .create(eq(PROJECT_ID), eq(REGION_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());
  }
}
Example 2
Source File: DataflowRunnerTest.java    From beam with Apache License 2.0 5 votes vote down vote up 
@Test
public void testUpdateAlreadyUpdatedPipeline() throws IOException {
  DataflowPipelineOptions options = buildPipelineOptions();
  options.setUpdate(true);
  options.setJobName("oldJobName");
  Dataflow mockDataflowClient = options.getDataflowClient();
  Dataflow.Projects.Locations.Jobs.Create mockRequest =
      mock(Dataflow.Projects.Locations.Jobs.Create.class);
  when(mockDataflowClient
          .projects()
          .locations()
          .jobs()
          .create(eq(PROJECT_ID), eq(REGION_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();
}
Example 3
Source File: DataflowClient.java    From beam with Apache License 2.0 4 votes vote down vote up 
public static DataflowClient create(DataflowPipelineOptions options) {
  return new DataflowClient(options.getDataflowClient(), options);
}