@Setup public void setup() { invoker = DoFnInvokers.invokerFor(splittableFn); invoker.invokeSetup(); }
@Teardown public void tearDown() { invoker.invokeTeardown(); invoker = null; } }
@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"); } }); }
@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(); }
/** @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; } }
RestrictionWithDefaultTracker.class, CoderForDefaultTracker.class)); assertThat( invoker.<RestrictionWithDefaultTracker>invokeGetRestrictionCoder(coderRegistry), instanceOf(CoderForDefaultTracker.class)); invoker.invokeSplitRestriction( "blah", "foo", assertEquals(stop(), invoker.invokeProcessElement(mockArgumentProvider)); assertThat( invoker.invokeNewTracker(new RestrictionWithDefaultTracker()), instanceOf(DefaultTracker.class));
private void invokeProcessElement(WindowedValue<InputT> elem) { // This can contain user code. Wrap it in case it throws an exception. try { invoker.invokeProcessElement(new DoFnProcessContext(elem)); } catch (Exception ex) { throw wrapUserCodeException(ex); } }
@StartBundle public void startBundle(StartBundleContext c) { invoker.invokeStartBundle( new DoFn<InputT, OutputT>.StartBundleContext() { @Override public PipelineOptions getPipelineOptions() { return c.getPipelineOptions(); } }); }
@ProcessElement public void process(ProcessContext c, BoundedWindow w) { InputT element = c.element().getKey(); RestrictionT restriction = c.element().getValue(); while (true) { RestrictionTracker<RestrictionT, PositionT> tracker = invoker.invokeNewTracker(restriction); ProcessContinuation continuation = invoker.invokeProcessElement(new NestedProcessContext<>(fn, c, element, w, tracker)); if (continuation.shouldResume()) { restriction = tracker.checkpoint(); Uninterruptibles.sleepUninterruptibly( continuation.resumeDelay().getMillis(), TimeUnit.MILLISECONDS); } else { break; } } }
invoker.invokeNewTracker(elementAndRestriction.getValue()); SplittableProcessElementInvoker<InputT, OutputT, RestrictionT, PositionT>.Result result = processElementInvoker.invokeProcessElement(
when(fn.processElement(mockProcessContext, tracker)).thenReturn(resume()); assertEquals(coder, invoker.invokeGetRestrictionCoder(CoderRegistry.createDefault())); assertEquals(restriction, invoker.invokeGetInitialRestriction("blah")); final List<SomeRestriction> outputs = new ArrayList<>(); invoker.invokeSplitRestriction( "blah", restriction, assertEquals(tracker, invoker.invokeNewTracker(restriction)); assertEquals( resume(), invoker.invokeProcessElement( new FakeArgumentProvider<String, String>() { @Override
private DoFn.ProcessContinuation invokeProcessElement(DoFn<String, String> fn) { return DoFnInvokers.invokerFor(fn).invokeProcessElement(mockArgumentProvider); }
@StartBundle public void startBundle(StartBundleContext c) throws Exception { invoker.invokeStartBundle(wrapContextAsStartBundle(c)); }
@Override public void open(Configuration parameters) throws Exception { doFnInvoker = DoFnInvokers.invokerFor(doFn); doFnInvoker.invokeSetup(); }
@Teardown public void tearDown() { invoker.invokeTeardown(); invoker = null; } }
createProcessContext( ValueInSingleWindow.of(element, timestamp, window, PaneInfo.NO_FIRING)); fnInvoker.invokeProcessElement( new DoFnInvoker.ArgumentProvider<InputT, OutputT>() { @Override
@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); } }
@FinishBundle public void finishBundle(FinishBundleContext c) throws Exception { invoker.invokeFinishBundle(wrapContextAsFinishBundle(c)); }
@Setup public void setup() throws Exception { invoker = DoFnInvokers.invokerFor(fn); invoker.invokeSetup(); }
@Teardown public void tearDown() throws Exception { invoker.invokeTeardown(); }