@FinishBundle public void finishBundle(FinishBundleContext c) { invoker.invokeFinishBundle( new DoFn<InputT, OutputT>.FinishBundleContext() { @Override public PipelineOptions getPipelineOptions() { return c.getPipelineOptions(); } @Override public void output(@Nullable OutputT output, Instant timestamp, BoundedWindow window) { throw new UnsupportedOperationException( "Output from FinishBundle for SDF is not supported"); } @Override public <T> void output( TupleTag<T> tag, T output, Instant timestamp, BoundedWindow window) { throw new UnsupportedOperationException( "Output from FinishBundle for SDF is not supported"); } }); }
@FinishBundle public void finishBundle(FinishBundleContext c) throws Exception { invoker.invokeFinishBundle(wrapContextAsFinishBundle(c)); }
@Override public void finishBundle() { // This can contain user code. Wrap it in case it throws an exception. try { invoker.invokeFinishBundle(new DoFnFinishBundleContext()); } catch (Throwable t) { // Exception in user code. throw wrapUserCodeException(t); } }
/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */ @Deprecated public void finishBundle() throws Exception { checkState( state == State.BUNDLE_STARTED, "Must be inside bundle to call finishBundle, but was: %s", state); try { fnInvoker.invokeFinishBundle(new TestFinishBundleContext()); } catch (UserCodeException e) { unwrapUserCodeException(e); } if (cloningBehavior == CloningBehavior.CLONE_PER_BUNDLE) { fnInvoker.invokeTeardown(); fn = null; fnInvoker = null; state = State.UNINITIALIZED; } else { state = State.BUNDLE_FINISHED; } }
@Test public void testFinishBundleException() throws Exception { DoFnInvoker<Integer, Integer> invoker = DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @FinishBundle public void finishBundle(@SuppressWarnings("unused") FinishBundleContext c) { throw new IllegalArgumentException("bogus"); } @ProcessElement public void processElement(@SuppressWarnings("unused") ProcessContext c) {} }); thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); invoker.invokeFinishBundle(null); }
@Test public void testDoFnWithStartBundleSetupTeardown() throws Exception { class MockFn extends DoFn<String, String> { @ProcessElement public void processElement(ProcessContext c) {} @StartBundle public void startBundle(StartBundleContext c) {} @FinishBundle public void finishBundle(FinishBundleContext c) {} @Setup public void before() {} @Teardown public void after() {} } MockFn fn = mock(MockFn.class); DoFnInvoker<String, String> invoker = DoFnInvokers.invokerFor(fn); invoker.invokeSetup(); invoker.invokeStartBundle(mockStartBundleContext); invoker.invokeFinishBundle(mockFinishBundleContext); invoker.invokeTeardown(); verify(fn).before(); verify(fn).startBundle(mockStartBundleContext); verify(fn).finishBundle(mockFinishBundleContext); verify(fn).after(); }