@StartBundle public void startBundle(StartBundleContext c) { invoker.invokeStartBundle( new DoFn<InputT, OutputT>.StartBundleContext() { @Override public PipelineOptions getPipelineOptions() { return c.getPipelineOptions(); } }); }
@StartBundle public void startBundle(StartBundleContext c) throws Exception { invoker.invokeStartBundle(wrapContextAsStartBundle(c)); }
@Override public void startBundle() { // This can contain user code. Wrap it in case it throws an exception. try { invoker.invokeStartBundle(new DoFnStartBundleContext()); } catch (Throwable t) { // Exception in user code. throw wrapUserCodeException(t); } }
/** @deprecated Use {@link TestPipeline} with the {@code DirectRunner}. */ @Deprecated public void startBundle() throws Exception { checkState( state == State.UNINITIALIZED || state == State.BUNDLE_FINISHED, "Wrong state during startBundle: %s", state); if (state == State.UNINITIALIZED) { initializeState(); } try { fnInvoker.invokeStartBundle(new TestStartBundleContext()); } catch (UserCodeException e) { unwrapUserCodeException(e); } state = State.BUNDLE_STARTED; }
@Test public void testStartBundleException() throws Exception { DoFnInvoker<Integer, Integer> invoker = DoFnInvokers.invokerFor( new DoFn<Integer, Integer>() { @StartBundle public void startBundle(@SuppressWarnings("unused") StartBundleContext c) { throw new IllegalArgumentException("bogus"); } @ProcessElement public void processElement(@SuppressWarnings("unused") ProcessContext c) {} }); thrown.expect(UserCodeException.class); thrown.expectMessage("bogus"); invoker.invokeStartBundle(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(); }