diff --git a/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java b/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java index 73b3f9ddd57e..e86fcf740947 100644 --- a/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java +++ b/runners/core-java/src/test/java/org/apache/beam/runners/core/SplittableParDoProcessFnTest.java @@ -46,12 +46,14 @@ import org.apache.beam.sdk.coders.SerializableCoder; import org.apache.beam.sdk.coders.VoidCoder; import org.apache.beam.sdk.io.range.OffsetRange; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.state.TimeDomain; import org.apache.beam.sdk.testing.ResetDateTimeProvider; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; import org.apache.beam.sdk.transforms.DoFn; -import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.DoFnSchemaInformation; import org.apache.beam.sdk.transforms.View; import org.apache.beam.sdk.transforms.splittabledofn.HasDefaultTracker; import org.apache.beam.sdk.transforms.splittabledofn.ManualWatermarkEstimator; @@ -138,7 +140,8 @@ public RestrictionTracker.IsBounded isBounded() { private static class ProcessFnTester< InputT, OutputT, RestrictionT, PositionT, WatermarkEstimatorStateT> implements AutoCloseable { - private final DoFnTester>, OutputT> tester; + private final DoFnRunner>, OutputT> runner; + private final List> outputs = new ArrayList<>(); private Instant currentProcessingTime; private InMemoryTimerInternals timerInternals; @@ -191,27 +194,66 @@ private static class ProcessFnTester< watermarkEstimatorStateCoder, windowingStrategy, sideInputMapping); - this.tester = DoFnTester.of(processFn); this.timerInternals = new InMemoryTimerInternals(); this.stateInternals = new TestInMemoryStateInternals<>("dummy"); processFn.setStateInternalsFactory(key -> stateInternals); processFn.setTimerInternalsFactory(key -> timerInternals); processFn.setSideInputReader(sideInputReader); + PipelineOptions options = PipelineOptionsFactory.create(); + TupleTag mainOutputTag = new TupleTag<>(); + WindowedValueMultiReceiver outputReceiver = + new WindowedValueMultiReceiver() { + @Override + @SuppressWarnings("unchecked") + public void output(TupleTag tag, WindowedValue output) { + for (BoundedWindow window : output.getWindows()) { + outputs.add( + (ValueInSingleWindow) + ValueInSingleWindow.of( + output.getValue(), + output.getTimestamp(), + window, + output.getPaneInfo())); + } + } + }; processFn.setProcessElementInvoker( new OutputAndTimeBoundedSplittableProcessElementInvoker<>( fn, - tester.getPipelineOptions(), - new DoFnTesterWindowedValueReceiver(tester), - tester.getMainOutputTag(), + options, + outputReceiver, + mainOutputTag, sideInputReader, Executors.newSingleThreadScheduledExecutor(Executors.defaultThreadFactory()), maxOutputsPerBundle, maxBundleDuration, () -> bundleFinalizer)); - // Do not clone since ProcessFn references non-serializable DoFnTester itself - // through the state/timer/output callbacks. - this.tester.setCloningBehavior(DoFnTester.CloningBehavior.DO_NOT_CLONE); - this.tester.startBundle(); + this.runner = + DoFnRunners.simpleRunner( + options, + processFn, + sideInputReader, + outputReceiver, + mainOutputTag, + Collections.emptyList(), + new StepContext() { + @Override + public StateInternals stateInternals() { + return stateInternals; + } + + @Override + public TimerInternals timerInternals() { + return timerInternals; + } + }, + null, + Collections.emptyMap(), + windowingStrategy, + DoFnSchemaInformation.create(), + sideInputMapping); + processFn.setup(options); + runner.startBundle(); timerInternals.advanceProcessingTime(currentProcessingTime); this.currentProcessingTime = currentProcessingTime; @@ -219,7 +261,8 @@ private static class ProcessFnTester< @Override public void close() throws Exception { - tester.close(); + runner.finishBundle(); + processFn.tearDown(); } /** Performs a seed {@link DoFn.ProcessElement} call feeding the element and restriction. */ @@ -233,9 +276,11 @@ void startElement(InputT element, RestrictionT restriction) throws Exception { } void startElement(WindowedValue> windowedValue) throws Exception { - tester.processElement( - KeyedWorkItems.elementsWorkItem( - "key".getBytes(StandardCharsets.UTF_8), Collections.singletonList(windowedValue))); + runner.processElement( + WindowedValues.valueInGlobalWindow( + KeyedWorkItems.elementsWorkItem( + "key".getBytes(StandardCharsets.UTF_8), + Collections.singletonList(windowedValue)))); } boolean advanceDrain() throws Exception { @@ -249,8 +294,9 @@ boolean advanceDrain() throws Exception { new Instant(Long.MAX_VALUE), TimeDomain.EVENT_TIME, CausedByDrain.CAUSED_BY_DRAIN)); - tester.processElement( - KeyedWorkItems.timersWorkItem("key".getBytes(StandardCharsets.UTF_8), timers)); + runner.processElement( + WindowedValues.valueInGlobalWindow( + KeyedWorkItems.timersWorkItem("key".getBytes(StandardCharsets.UTF_8), timers))); return true; } @@ -270,17 +316,29 @@ boolean advanceProcessingTimeBy(Duration duration) throws Exception { if (timers.isEmpty()) { return false; } - tester.processElement( - KeyedWorkItems.timersWorkItem("key".getBytes(StandardCharsets.UTF_8), timers)); + runner.processElement( + WindowedValues.valueInGlobalWindow( + KeyedWorkItems.timersWorkItem("key".getBytes(StandardCharsets.UTF_8), timers))); return true; } List> peekOutputElementsInWindow(BoundedWindow window) { - return tester.peekOutputElementsInWindow(window); + List> res = new ArrayList<>(); + for (ValueInSingleWindow value : outputs) { + if (value.getWindow().equals(window)) { + res.add(TimestampedValue.of(value.getValue(), value.getTimestamp())); + } + } + return res; } List takeOutputElements() { - return tester.takeOutputElements(); + List res = new ArrayList<>(); + for (ValueInSingleWindow value : outputs) { + res.add(value.getValue()); + } + outputs.clear(); + return res; } public Instant getWatermarkHold() { @@ -288,25 +346,6 @@ public Instant getWatermarkHold() { } } - private static class DoFnTesterWindowedValueReceiver implements WindowedValueMultiReceiver { - private final DoFnTester tester; - - private DoFnTesterWindowedValueReceiver(DoFnTester tester) { - this.tester = tester; - } - - @Override - public void output(TupleTag tag, WindowedValue output) { - for (BoundedWindow window : output.getWindows()) { - tester - .getMutableOutput(tag) - .add( - ValueInSingleWindow.of( - output.getValue(), output.getTimestamp(), window, output.getPaneInfo())); - } - } - } - /** A simple splittable {@link DoFn} that's actually monolithic. */ private static class ToStringFn extends DoFn { @ProcessElement