Examples with UserCodeException org.apache.beam.sdk.util.UserCodeException used on opensource projects

Example 6 with UserCodeException

use of org.apache.beam.sdk.util.UserCodeException in project beam by apache.

the class SimpleParDoFnTest method testErrorPropagation.

@Test
public void testErrorPropagation() throws Exception {
    TestErrorDoFn fn = new TestErrorDoFn();
    DoFnInfo<?, ?> fnInfo = DoFnInfo.forFn(fn, WindowingStrategy.globalDefault(), null, /* side input views */
    null, /* input coder */
    MAIN_OUTPUT, DoFnSchemaInformation.create(), Collections.emptyMap());
    TestReceiver receiver = new TestReceiver();
    ParDoFn userParDoFn = new SimpleParDoFn<>(options, DoFnInstanceManagers.singleInstance(fnInfo), new EmptySideInputReader(), MAIN_OUTPUT, ImmutableMap.of(MAIN_OUTPUT, 0), BatchModeExecutionContext.forTesting(options, "testStage").getStepContext(operationContext), operationContext, DoFnSchemaInformation.create(), Collections.emptyMap(), SimpleDoFnRunnerFactory.INSTANCE);
    try {
        userParDoFn.startBundle(receiver);
        userParDoFn.processElement(null);
        fail("should have failed");
    } catch (Exception exn) {
        // Because we're calling this from inside the SDK and not from a
        // user's program (e.g. through Pipeline.run), the error should
        // be thrown as a UserCodeException. The cause of the
        // UserCodeError shouldn't contain any of the stack from within
        // the SDK, since we don't want to overwhelm users with stack
        // frames outside of their control.
        assertThat(exn, instanceOf(UserCodeException.class));
        // Stack trace of the cause should contain three frames:
        // TestErrorDoFn.nestedFunctionBeta
        // TestErrorDoFn.nestedFunctionAlpha
        // TestErrorDoFn.startBundle
        assertThat(stackTraceFrameStrings(exn.getCause()), contains(containsString("TestErrorDoFn.nestedFunctionBeta"), containsString("TestErrorDoFn.nestedFunctionAlpha"), containsString("TestErrorDoFn.startBundle")));
        assertThat(exn.toString(), containsString("test error in initialize"));
    }
    try {
        userParDoFn.processElement(WindowedValue.valueInGlobalWindow(3));
        fail("should have failed");
    } catch (Exception exn) {
        // Exception should be a UserCodeException since we're calling
        // from inside the SDK.
        assertThat(exn, instanceOf(UserCodeException.class));
        // Stack trace of the cause should contain two frames:
        // TestErrorDoFn.nestedFunctionBeta
        // TestErrorDoFn.processElement
        assertThat(stackTraceFrameStrings(exn.getCause()), contains(containsString("TestErrorDoFn.nestedFunctionBeta"), containsString("TestErrorDoFn.processElement")));
        assertThat(exn.toString(), containsString("test error in process"));
    }
    try {
        userParDoFn.finishBundle();
        fail("should have failed");
    } catch (Exception exn) {
        // Exception should be a UserCodeException since we're calling
        // from inside the SDK.
        assertThat(exn, instanceOf(UserCodeException.class));
        // Stack trace should only contain a single frame:
        // TestErrorDoFn.finishBundle
        assertThat(stackTraceFrameStrings(exn.getCause()), contains(containsString("TestErrorDoFn.finishBundle")));
        assertThat(exn.toString(), containsString("test error in finalize"));
    }
}