Examples with TestPipelineOptions org.apache.beam.sdk.testing.TestPipelineOptions used on opensource projects

Example 1 with TestPipelineOptions

use of org.apache.beam.sdk.testing.TestPipelineOptions in project beam by apache.

the class TestDataflowRunnerTest method testStreamingOnSuccessMatcherWhenPipelineSucceeds.

/**
   * Tests that when a streaming pipeline terminates and doesn't fail due to {@link PAssert} that
   * the {@link TestPipelineOptions#setOnSuccessMatcher(SerializableMatcher) on success matcher} is
   * invoked.
   */
@Test
public void testStreamingOnSuccessMatcherWhenPipelineSucceeds() throws Exception {
    options.setStreaming(true);
    Pipeline p = TestPipeline.create(options);
    PCollection<Integer> pc = p.apply(Create.of(1, 2, 3));
    PAssert.that(pc).containsInAnyOrder(1, 2, 3);
    final DataflowPipelineJob mockJob = Mockito.mock(DataflowPipelineJob.class);
    when(mockJob.getState()).thenReturn(State.DONE);
    when(mockJob.getProjectId()).thenReturn("test-project");
    when(mockJob.getJobId()).thenReturn("test-job");
    DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class);
    when(mockRunner.run(any(Pipeline.class))).thenReturn(mockJob);
    TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient);
    options.as(TestPipelineOptions.class).setOnSuccessMatcher(new TestSuccessMatcher(mockJob, 1));
    when(mockJob.waitUntilFinish(any(Duration.class), any(JobMessagesHandler.class))).thenReturn(State.DONE);
    when(mockClient.getJobMetrics(anyString())).thenReturn(generateMockMetricResponse(true, /* success */
    true));
    runner.run(p, mockRunner);
}

Example 2 with TestPipelineOptions

use of org.apache.beam.sdk.testing.TestPipelineOptions in project beam by apache.

the class TestDataflowRunnerTest method testStreamingOnSuccessMatcherWhenPipelineFails.

/**
   * Tests that when a streaming pipeline terminates in FAIL that the {@link
   * TestPipelineOptions#setOnSuccessMatcher(SerializableMatcher) on success matcher} is not
   * invoked.
   */
@Test
public void testStreamingOnSuccessMatcherWhenPipelineFails() throws Exception {
    options.setStreaming(true);
    Pipeline p = TestPipeline.create(options);
    PCollection<Integer> pc = p.apply(Create.of(1, 2, 3));
    PAssert.that(pc).containsInAnyOrder(1, 2, 3);
    final DataflowPipelineJob mockJob = Mockito.mock(DataflowPipelineJob.class);
    when(mockJob.getState()).thenReturn(State.FAILED);
    when(mockJob.getProjectId()).thenReturn("test-project");
    when(mockJob.getJobId()).thenReturn("test-job");
    DataflowRunner mockRunner = Mockito.mock(DataflowRunner.class);
    when(mockRunner.run(any(Pipeline.class))).thenReturn(mockJob);
    TestDataflowRunner runner = TestDataflowRunner.fromOptionsAndClient(options, mockClient);
    options.as(TestPipelineOptions.class).setOnSuccessMatcher(new TestFailureMatcher());
    when(mockJob.waitUntilFinish(any(Duration.class), any(JobMessagesHandler.class))).thenReturn(State.FAILED);
    expectedException.expect(RuntimeException.class);
    runner.run(p, mockRunner);
// If the onSuccessMatcher were invoked, it would have crashed here with AssertionError
}

Example 3 with TestPipelineOptions

use of org.apache.beam.sdk.testing.TestPipelineOptions in project beam by apache.

the class TestDataflowRunner method run.

DataflowPipelineJob run(Pipeline pipeline, DataflowRunner runner) {
    updatePAssertCount(pipeline);
    TestPipelineOptions testPipelineOptions = options.as(TestPipelineOptions.class);
    final DataflowPipelineJob job;
    job = runner.run(pipeline);
    LOG.info("Running Dataflow job {} with {} expected assertions.", job.getJobId(), expectedNumberOfAssertions);
    assertThat(job, testPipelineOptions.getOnCreateMatcher());
    Boolean jobSuccess;
    Optional<Boolean> allAssertionsPassed;
    ErrorMonitorMessagesHandler messageHandler = new ErrorMonitorMessagesHandler(job, new MonitoringUtil.LoggingHandler());
    if (options.isStreaming()) {
        jobSuccess = waitForStreamingJobTermination(job, messageHandler);
        // No metrics in streaming
        allAssertionsPassed = Optional.absent();
    } else {
        jobSuccess = waitForBatchJobTermination(job, messageHandler);
        allAssertionsPassed = checkForPAssertSuccess(job);
    }
    // the actionable message from the logs it is acceptable.
    if (!allAssertionsPassed.isPresent()) {
        LOG.warn("Dataflow job {} did not output a success or failure metric.", job.getJobId());
    } else if (!allAssertionsPassed.get()) {
        throw new AssertionError(errorMessage(job, messageHandler));
    }
    // as simply job failures.
    if (!jobSuccess) {
        throw new RuntimeException(errorMessage(job, messageHandler));
    }
    // If there is no reason to immediately fail, run the success matcher.
    assertThat(job, testPipelineOptions.getOnSuccessMatcher());
    return job;
}