This is an automated email from the ASF dual-hosted git repository.
stankiewicz pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 7b60339d86b Fix output bugs in FnApiDoFnRunner @OnTimer and
@OnWindowExpiration (#40024)
7b60339d86b is described below
commit 7b60339d86b8de685ce66833c0f868ad9362e2aa
Author: Kenneth Knowles <[email protected]>
AuthorDate: Thu Sep 17 10:59:38 2026 -0400
Fix output bugs in FnApiDoFnRunner @OnTimer and @OnWindowExpiration (#40024)
The three argument providers each carry their own near-identical copy of
the output-receiver plumbing, and the @OnTimer and @OnWindowExpiration
copies had drifted from the working @ProcessElement one:
- OnTimerContext.outputWindowedValue(tag, ...) had an empty body, so
tagged windowed output from @OnTimer was silently dropped.
- OnTimerContext's tagged receivers passed no tag to
outputWindowedValue, so output to a side tag went to the main output.
- The @OnTimer row receivers built from WindowedValues.builder(
currentElement) or called withValue() on a fresh builder. Neither
works during timer processing: currentElement is only set while
processing an element, and withValue() reads timestamp/window/pane
off the builder it is called on. Both now seed from currentTimer,
matching the @OnWindowExpiration equivalents.
- OnWindowExpirationContext read currentElement.getValueKind() while
currentElement was null. Window expiration emits new records, so this
is ValueKind.INSERT.
- OnWindowExpirationContext.timeDomain() returned currentTimeDomain,
which processOnWindowExpiration never sets. TimeDomain is not an
allowed @OnWindowExpiration parameter, so drop the override and let
BaseArgumentProvider reject it.
- outputWindowedValue(tag, ...) skipped the unknown-tag check its
siblings perform, so a bad tag produced an NPE.
- processTimer's finally did not clear causedByDrain, unlike
processOnWindowExpiration. Not observable today since every entry
point sets it before invoking user code, but the asymmetry is a trap.
The two inner Context classes are renamed to TimerContext and
WindowExpirationContext so the fields holding them can drop their raw
types; the enclosing providers are generic in K, and errorprone's
SameNameButDifferent rejects the bare Context name.
Co-authored-by: Claude Opus 4.8 <[email protected]>
---
.../apache/beam/fn/harness/FnApiDoFnRunner.java | 61 +++--
.../beam/fn/harness/FnApiDoFnRunnerTest.java | 292 +++++++++++++++++++++
2 files changed, 333 insertions(+), 20 deletions(-)
diff --git
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
index a5914a25f79..d391807d798 100644
---
a/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
+++
b/sdks/java/harness/src/main/java/org/apache/beam/fn/harness/FnApiDoFnRunner.java
@@ -1235,6 +1235,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
currentTimer = null;
currentTimeDomain = null;
currentWindow = null;
+ causedByDrain = null;
}
}
@@ -2295,10 +2296,10 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
* DoFn.OnWindowExpiration @OnWindowExpiration}.
*/
private class OnWindowExpirationContext<K> extends
BaseArgumentProvider<InputT, OutputT> {
- private class Context extends DoFn<InputT,
OutputT>.OnWindowExpirationContext
+ private class WindowExpirationContext extends DoFn<InputT,
OutputT>.OnWindowExpirationContext
implements OutputReceiver<OutputT> {
- private Context() {
+ private WindowExpirationContext() {
doFn.super();
}
@@ -2369,7 +2370,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
null,
currentTimer.causedByDrain(),
null,
- currentElement.getValueKind()));
+ ValueKind.INSERT));
}
@Override
@@ -2395,6 +2396,9 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
checkOnWindowExpirationTimestamp(timestamp);
FnDataReceiver<WindowedValue<T>> consumer =
(FnDataReceiver) localNameToConsumer.get(tag.getId());
+ if (consumer == null) {
+ throw new IllegalArgumentException(String.format("Unknown output tag
%s", tag));
+ }
outputTo(consumer, WindowedValues.of(output, timestamp, windows,
paneInfo));
}
@@ -2405,7 +2409,12 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
@Override
public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T>
windowedValue) {
- outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()),
windowedValue);
+ FnDataReceiver<WindowedValue<T>> consumer =
+ (FnDataReceiver) localNameToConsumer.get(tag.getId());
+ if (consumer == null) {
+ throw new IllegalArgumentException(String.format("Unknown output tag
%s", tag));
+ }
+ outputTo(consumer, windowedValue);
}
@SuppressWarnings(
@@ -2435,8 +2444,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
}
}
- private final OnWindowExpirationContext.Context context =
- new OnWindowExpirationContext.Context();
+ private final WindowExpirationContext context = new
WindowExpirationContext();
@Override
public DoFn<InputT, OutputT>.OnWindowExpirationContext
onWindowExpirationContext(
@@ -2459,11 +2467,6 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
return currentTimer.getHoldTimestamp();
}
- @Override
- public TimeDomain timeDomain(DoFn<InputT, OutputT> doFn) {
- return currentTimeDomain;
- }
-
@Override
public K key() {
return (K) currentTimer.getUserKey();
@@ -2622,9 +2625,9 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
/** Provides arguments for a {@link DoFnInvoker} for {@link DoFn.OnTimer
@OnTimer}. */
private class OnTimerContext<K> extends BaseArgumentProvider<InputT,
OutputT> {
- private class Context extends DoFn<InputT, OutputT>.OnTimerContext
+ private class TimerContext extends DoFn<InputT, OutputT>.OnTimerContext
implements OutputReceiver<OutputT> {
- private Context() {
+ private TimerContext() {
doFn.super();
}
@@ -2719,7 +2722,12 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
@Override
public <T> void outputWindowedValue(TupleTag<T> tag, WindowedValue<T>
windowedValue) {
- outputTo((FnDataReceiver) localNameToConsumer.get(tag.getId()),
windowedValue);
+ FnDataReceiver<WindowedValue<T>> consumer =
+ (FnDataReceiver) localNameToConsumer.get(tag.getId());
+ if (consumer == null) {
+ throw new IllegalArgumentException(String.format("Unknown output tag
%s", tag));
+ }
+ outputTo(consumer, windowedValue);
}
@Override
@@ -2728,7 +2736,15 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
T output,
Instant timestamp,
Collection<? extends BoundedWindow> windows,
- PaneInfo paneInfo) {}
+ PaneInfo paneInfo) {
+ checkTimerTimestamp(timestamp);
+ FnDataReceiver<WindowedValue<T>> consumer =
+ (FnDataReceiver) localNameToConsumer.get(tag.getId());
+ if (consumer == null) {
+ throw new IllegalArgumentException(String.format("Unknown output tag
%s", tag));
+ }
+ outputTo(consumer, WindowedValues.of(output, timestamp, windows,
paneInfo));
+ }
@Override
public TimeDomain timeDomain() {
@@ -2772,7 +2788,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
}
}
- private final OnTimerContext.Context context = new
OnTimerContext.Context();
+ private final TimerContext context = new TimerContext();
@Override
public BoundedWindow window() {
@@ -2822,8 +2838,12 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
@Override
public OutputBuilder<Row> builder(Row value) {
- return WindowedValues.builder(currentElement)
- .withValue(value)
+ return WindowedValues.<Row>builder()
+ .setValue(value)
+ .setTimestamp(currentTimer.getHoldTimestamp())
+ .setWindow(currentWindow)
+ .setPaneInfo(currentTimer.getPaneInfo())
+ .setCausedByDrain(currentTimer.causedByDrain())
.setReceiver(
windowedValue ->
context.outputWindowedValue(
@@ -2860,7 +2880,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
.setWindow(currentWindow)
.setCausedByDrain(currentTimer.causedByDrain())
.setPaneInfo(currentTimer.getPaneInfo())
- .setReceiver(windowedValue ->
context.outputWindowedValue(windowedValue));
+ .setReceiver(windowedValue ->
context.outputWindowedValue(tag, windowedValue));
}
};
}
@@ -2887,7 +2907,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
@Override
public OutputBuilder<Row> builder(Row value) {
return WindowedValues.<Row>builder()
- .withValue(value)
+ .setValue(value)
.setTimestamp(currentTimer.getHoldTimestamp())
.setWindow(currentWindow)
.setPaneInfo(currentTimer.getPaneInfo())
@@ -2895,6 +2915,7 @@ public class FnApiDoFnRunner<InputT, RestrictionT,
PositionT, WatermarkEstimator
.setReceiver(
windowedValue ->
context.outputWindowedValue(
+ tag,
windowedValue.withValue(
fromRowFunction.apply(windowedValue.getValue()))));
}
diff --git
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
index d24ab39c047..4166b2a48dc 100644
---
a/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
+++
b/sdks/java/harness/src/test/java/org/apache/beam/fn/harness/FnApiDoFnRunnerTest.java
@@ -87,6 +87,7 @@ import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.metrics.MetricsEnvironment;
import org.apache.beam.sdk.options.ExperimentalOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.state.BagState;
import org.apache.beam.sdk.state.CombiningState;
import org.apache.beam.sdk.state.StateSpec;
@@ -128,6 +129,7 @@ import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionTuple;
import org.apache.beam.sdk.values.PCollectionView;
+import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.TupleTag;
import org.apache.beam.sdk.values.TupleTagList;
import org.apache.beam.sdk.values.WindowedValue;
@@ -1140,6 +1142,296 @@ public class FnApiDoFnRunnerTest implements
Serializable {
assertThat(result, containsInAnyOrder(expected.toArray()));
}
+ private static class TestTimerTaggedOutputDoFn extends DoFn<KV<String,
String>, String> {
+ @TimerId("event")
+ private final TimerSpec eventTimerSpec =
TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ private final TupleTag<String> additionalOutput;
+
+ private TestTimerTaggedOutputDoFn(TupleTag<String> additionalOutput) {
+ this.additionalOutput = additionalOutput;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext context, @TimerId("event")
Timer eventTimeTimer) {
+
eventTimeTimer.withOutputTimestamp(context.timestamp()).set(context.timestamp());
+ }
+
+ @OnTimer("event")
+ public void eventTimer(
+ OnTimerContext context, @Key String key, MultiOutputReceiver
receiver) {
+ context.output("main:" + key);
+ context.output(additionalOutput, "output:" + key);
+ context.outputWindowedValue(
+ additionalOutput,
+ "outputWindowedValue:" + key,
+ context.timestamp(),
+ Collections.singletonList(GlobalWindow.INSTANCE),
+ PaneInfo.NO_FIRING);
+ receiver.get(additionalOutput).output("receiver:" + key);
+ }
+ }
+
+ @Test
+ public void testTimerTaggedOutputs() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<KV<String, String>> valuePCollection =
+ p.apply(Create.of(KV.of("unused", "unused")));
+ TupleTag<String> mainOutput = new TupleTag<String>("main") {};
+ TupleTag<String> additionalOutput = new TupleTag<String>("additional")
{};
+ PCollectionTuple outputPCollection =
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID,
+ ParDo.of(new TestTimerTaggedOutputDoFn(additionalOutput))
+ .withOutputTags(mainOutput,
TupleTagList.of(additionalOutput)));
+
+ SdkComponents sdkComponents = SdkComponents.create();
+ sdkComponents.registerEnvironment(Environment.getDefaultInstance());
+ RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p,
sdkComponents);
+ String outputPCollectionId =
+ sdkComponents.registerPCollection(outputPCollection.get(mainOutput));
+ String additionalPCollectionId =
+
sdkComponents.registerPCollection(outputPCollection.get(additionalOutput));
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(TEST_TRANSFORM_ID);
+
+ List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
+ List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID,
pTransform)
+ .beamFnStateClient(new
FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of()))
+ .processBundleInstructionId("57L")
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+
.putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+
.putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ .outboundAggregators(
+ ImmutableMap.of(
+ ApiServiceDescriptor.getDefaultInstance(),
+ new TestBeamFnDataOutboundAggregator(() -> "57L")))
+
.timerApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance())
+ .build();
+ context.addPCollectionConsumer(
+ outputPCollectionId,
+ (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
mainOutputValues::add);
+ context.addPCollectionConsumer(
+ additionalPCollectionId,
+ (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
additionalOutputValues::add);
+
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+ Iterables.getOnlyElement(context.getStartBundleFunctions()).run();
+
+ context
+ .getIncomingTimerEndpoint("ts-event")
+ .getReceiver()
+ .accept(timerInGlobalWindow("A", new Instant(1400L), new
Instant(2400L)));
+
+ assertThat(mainOutputValues, contains(isValueInGlobalWindow("main:A",
new Instant(1400L))));
+ assertThat(
+ additionalOutputValues,
+ contains(
+ isValueInGlobalWindow("output:A", new Instant(1400L)),
+ isValueInGlobalWindow("outputWindowedValue:A", new
Instant(1400L)),
+ isValueInGlobalWindow("receiver:A", new Instant(1400L))));
+
+ Iterables.getOnlyElement(context.getFinishBundleFunctions()).run();
+ Iterables.getOnlyElement(context.getTearDownFunctions()).run();
+ }
+
+ private static final Schema ROW_SCHEMA =
+ Schema.of(Schema.Field.of("field", Schema.FieldType.STRING));
+
+ private static Row row(String value) {
+ return Row.withSchema(ROW_SCHEMA).addValue(value).build();
+ }
+
+ private static class TestTimerRowOutputDoFn extends DoFn<KV<String,
String>, Row> {
+ @TimerId("event")
+ private final TimerSpec eventTimerSpec =
TimerSpecs.timer(TimeDomain.EVENT_TIME);
+
+ private final TupleTag<Row> mainOutput;
+ private final TupleTag<Row> additionalOutput;
+
+ private TestTimerRowOutputDoFn(TupleTag<Row> mainOutput, TupleTag<Row>
additionalOutput) {
+ this.mainOutput = mainOutput;
+ this.additionalOutput = additionalOutput;
+ }
+
+ @ProcessElement
+ public void processElement(ProcessContext context, @TimerId("event")
Timer eventTimeTimer) {
+
eventTimeTimer.withOutputTimestamp(context.timestamp()).set(context.timestamp());
+ }
+
+ @OnTimer("event")
+ public void eventTimer(@Key String key, MultiOutputReceiver receiver) {
+ receiver.getRowReceiver(mainOutput).output(row("mainRow:" + key));
+ receiver.getRowReceiver(additionalOutput).output(row("taggedRow:" +
key));
+ }
+ }
+
+ @Test
+ public void testTimerRowOutputReceivers() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<KV<String, String>> valuePCollection =
+ p.apply(Create.of(KV.of("unused", "unused")));
+ TupleTag<Row> mainOutput = new TupleTag<Row>("main") {};
+ TupleTag<Row> additionalOutput = new TupleTag<Row>("additional") {};
+ PCollectionTuple outputPCollection =
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID,
+ ParDo.of(new TestTimerRowOutputDoFn(mainOutput,
additionalOutput))
+ .withOutputTags(mainOutput,
TupleTagList.of(additionalOutput)));
+ outputPCollection.get(mainOutput).setRowSchema(ROW_SCHEMA);
+ outputPCollection.get(additionalOutput).setRowSchema(ROW_SCHEMA);
+
+ SdkComponents sdkComponents = SdkComponents.create();
+ sdkComponents.registerEnvironment(Environment.getDefaultInstance());
+ RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p,
sdkComponents);
+ String outputPCollectionId =
+ sdkComponents.registerPCollection(outputPCollection.get(mainOutput));
+ String additionalPCollectionId =
+
sdkComponents.registerPCollection(outputPCollection.get(additionalOutput));
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(TEST_TRANSFORM_ID);
+
+ List<WindowedValue<Row>> mainOutputValues = new ArrayList<>();
+ List<WindowedValue<Row>> additionalOutputValues = new ArrayList<>();
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID,
pTransform)
+ .beamFnStateClient(new
FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of()))
+ .processBundleInstructionId("57L")
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+
.putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+
.putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ .outboundAggregators(
+ ImmutableMap.of(
+ ApiServiceDescriptor.getDefaultInstance(),
+ new TestBeamFnDataOutboundAggregator(() -> "57L")))
+
.timerApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance())
+ .build();
+ context.addPCollectionConsumer(
+ outputPCollectionId,
+ (FnDataReceiver) (FnDataReceiver<WindowedValue<Row>>)
mainOutputValues::add);
+ context.addPCollectionConsumer(
+ additionalPCollectionId,
+ (FnDataReceiver) (FnDataReceiver<WindowedValue<Row>>)
additionalOutputValues::add);
+
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+ Iterables.getOnlyElement(context.getStartBundleFunctions()).run();
+
+ context
+ .getIncomingTimerEndpoint("ts-event")
+ .getReceiver()
+ .accept(timerInGlobalWindow("A", new Instant(1400L), new
Instant(2400L)));
+
+ assertThat(
+ mainOutputValues, contains(isValueInGlobalWindow(row("mainRow:A"),
new Instant(1400L))));
+ assertThat(
+ additionalOutputValues,
+ contains(isValueInGlobalWindow(row("taggedRow:A"), new
Instant(1400L))));
+
+ Iterables.getOnlyElement(context.getFinishBundleFunctions()).run();
+ Iterables.getOnlyElement(context.getTearDownFunctions()).run();
+ }
+
+ private static class TestWindowExpirationDoFn extends DoFn<KV<String,
String>, String> {
+ @StateId("bag")
+ private final StateSpec<BagState<String>> bagStateSpec =
StateSpecs.bag(StringUtf8Coder.of());
+
+ private final TupleTag<String> additionalOutput;
+
+ private TestWindowExpirationDoFn(TupleTag<String> additionalOutput) {
+ this.additionalOutput = additionalOutput;
+ }
+
+ @ProcessElement
+ public void processElement(
+ ProcessContext context, @StateId("bag") BagState<String> bagState) {
+ bagState.add(context.element().getValue());
+ }
+
+ @OnWindowExpiration
+ public void onWindowExpiration(OnWindowExpirationContext context, @Key
String key) {
+ context.output("main:" + key);
+ context.output(additionalOutput, "output:" + key);
+ }
+ }
+
+ @Test
+ public void testOnWindowExpirationTaggedOutputs() throws Exception {
+ Pipeline p = Pipeline.create();
+ PCollection<KV<String, String>> valuePCollection =
+ p.apply(Create.of(KV.of("unused", "unused")));
+ TupleTag<String> mainOutput = new TupleTag<String>("main") {};
+ TupleTag<String> additionalOutput = new TupleTag<String>("additional")
{};
+ PCollectionTuple outputPCollection =
+ valuePCollection.apply(
+ TEST_TRANSFORM_ID,
+ ParDo.of(new TestWindowExpirationDoFn(additionalOutput))
+ .withOutputTags(mainOutput,
TupleTagList.of(additionalOutput)));
+
+ SdkComponents sdkComponents = SdkComponents.create();
+ sdkComponents.registerEnvironment(Environment.getDefaultInstance());
+ RunnerApi.Pipeline pProto = PipelineTranslation.toProto(p,
sdkComponents);
+ String outputPCollectionId =
+ sdkComponents.registerPCollection(outputPCollection.get(mainOutput));
+ String additionalPCollectionId =
+
sdkComponents.registerPCollection(outputPCollection.get(additionalOutput));
+ RunnerApi.PTransform pTransform =
+ pProto.getComponents().getTransformsOrThrow(TEST_TRANSFORM_ID);
+ String onWindowExpirationFamilyId =
+ RunnerApi.ParDoPayload.parseFrom(pTransform.getSpec().getPayload())
+ .getOnWindowExpirationTimerFamilySpec();
+
+ List<WindowedValue<String>> mainOutputValues = new ArrayList<>();
+ List<WindowedValue<String>> additionalOutputValues = new ArrayList<>();
+ PTransformRunnerFactoryTestContext context =
+ PTransformRunnerFactoryTestContext.builder(TEST_TRANSFORM_ID,
pTransform)
+ .beamFnStateClient(new
FakeBeamFnStateClient(StringUtf8Coder.of(), ImmutableMap.of()))
+ .processBundleInstructionId("57L")
+ .components(
+ RunnerApi.Components.newBuilder()
+ .putAllCoders(pProto.getComponents().getCodersMap())
+ .putAllEnvironments(Collections.emptyMap())
+
.putAllWindowingStrategies(pProto.getComponents().getWindowingStrategiesMap())
+
.putAllPcollections(pProto.getComponentsOrBuilder().getPcollectionsMap())
+ .build())
+ .outboundAggregators(
+ ImmutableMap.of(
+ ApiServiceDescriptor.getDefaultInstance(),
+ new TestBeamFnDataOutboundAggregator(() -> "57L")))
+
.timerApiServiceDescriptor(ApiServiceDescriptor.getDefaultInstance())
+ .build();
+ context.addPCollectionConsumer(
+ outputPCollectionId,
+ (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
mainOutputValues::add);
+ context.addPCollectionConsumer(
+ additionalPCollectionId,
+ (FnDataReceiver) (FnDataReceiver<WindowedValue<String>>)
additionalOutputValues::add);
+
+ new FnApiDoFnRunner.Factory<>().addRunnerForPTransform(context);
+ Iterables.getOnlyElement(context.getStartBundleFunctions()).run();
+
+ context
+ .getIncomingTimerEndpoint(onWindowExpirationFamilyId)
+ .getReceiver()
+ .accept(timerInGlobalWindow("A", new Instant(1400L), new
Instant(2400L)));
+
+ assertThat(mainOutputValues, contains(isValueInGlobalWindow("main:A",
new Instant(1400L))));
+ assertThat(
+ additionalOutputValues, contains(isValueInGlobalWindow("output:A",
new Instant(1400L))));
+
+ Iterables.getOnlyElement(context.getFinishBundleFunctions()).run();
+ Iterables.getOnlyElement(context.getTearDownFunctions()).run();
+ }
+
private <K> org.apache.beam.sdk.util.construction.Timer<K>
timerInGlobalWindow(
K userKey, Instant holdTimestamp, Instant fireTimestamp) {
return dynamicTimerInGlobalWindow(userKey, "", holdTimestamp,
fireTimestamp);
