From 5937c55f3280920cb5976d7724c452cf5369fa1c Mon Sep 17 00:00:00 2001 From: Rishavbhattarai Date: Tue, 27 Oct 2026 23:07:34 -0700 Subject: [PATCH] Port SplittableParDoProcessFnTest off DoFnTester Run ProcessFn through DoFnRunners.simpleRunner, the same runner the DirectRunner uses, instead of the deprecated DoFnTester. The ProcessFnTester helper keeps its API, so the test bodies don't change. Fixes #18691 --- .../core/SplittableParDoProcessFnTest.java | 117 ++++++++++++------ 1 file changed, 78 insertions(+), 39 deletions(-) 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