From 38488d79915eb5b7cf4bcba3a420387c235ccb2c Mon Sep 17 00:00:00 2001 From: SuperSaiyan101 Date: Fri, 30 Oct 2026 23:58:49 +0530 Subject: [PATCH 1/2] [#18690] Port Iceberg IO tests off deprecated DoFnTester --- .../beam/sdk/io/iceberg/BundleLifterTest.java | 117 +++++++++++------- .../iceberg/cdc/ApplyWatermarkColumnTest.java | 24 ++-- .../io/iceberg/cdc/LocalResolveDoFnTest.java | 49 ++++++-- .../io/iceberg/cdc/ResolveChangesTest.java | 67 +++++++--- 4 files changed, 175 insertions(+), 82 deletions(-) diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java index 1eaa0920e6c6..c7a5a3421ab2 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java @@ -21,9 +21,17 @@ import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.empty; +import java.util.HashMap; +import java.util.Map; import org.apache.beam.sdk.io.iceberg.BundleLifter.BundleLiftDoFn; -import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.testing.TestOutputReceiver; +import org.apache.beam.sdk.transforms.DoFn.MultiOutputReceiver; +import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; +import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.TupleTag; +import org.checkerframework.checker.initialization.qual.Initialized; +import org.checkerframework.checker.nullness.qual.NonNull; +import org.checkerframework.checker.nullness.qual.UnknownKeyFor; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -36,64 +44,89 @@ public class BundleLifterTest { private static final TupleTag STRING_SMALL = new TupleTag() {}; private static final TupleTag STRING_LARGE = new TupleTag() {}; + private static class SimpleMultiOutputReceiver implements MultiOutputReceiver { + private final Map, TestOutputReceiver> receivers = new HashMap<>(); + + @SuppressWarnings("unchecked") + @Override + public @UnknownKeyFor @NonNull @Initialized OutputReceiver get( + @UnknownKeyFor @NonNull @Initialized TupleTag tag) { + return (OutputReceiver) receivers.computeIfAbsent(tag, k -> new TestOutputReceiver()); + } + + @Override + public @UnknownKeyFor @NonNull @Initialized + OutputReceiver<@UnknownKeyFor @NonNull @Initialized Row> getRowReceiver( + @UnknownKeyFor @NonNull @Initialized TupleTag tag) { + throw new UnsupportedOperationException("getRowReceiver not supported in test"); + } + + @SuppressWarnings("unchecked") + public java.util.List getOutput(TupleTag tag) { + TestOutputReceiver receiver = (TestOutputReceiver) receivers.get(tag); + return receiver == null ? java.util.Collections.emptyList() : receiver.getOutputs(); + } + } + @Test public void testSmallBundle() throws Exception { - DoFnTester tester = - DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1)); + BundleLiftDoFn fn = new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1); + SimpleMultiOutputReceiver receiver = new SimpleMultiOutputReceiver(); - tester.startBundle(); - tester.processElement(1); - tester.processElement(2); - tester.finishBundle(); + fn.startBundle(); + fn.processElement(1, receiver); + fn.processElement(2, receiver); + fn.finishBundle(); - assertThat(tester.peekOutputElements(INTEGER_SMALL), containsInAnyOrder(1, 2)); - assertThat(tester.peekOutputElements(INTEGER_LARGE), empty()); + assertThat(receiver.getOutput(INTEGER_SMALL), containsInAnyOrder(1, 2)); + assertThat(receiver.getOutput(INTEGER_LARGE), empty()); } @Test public void testLargeBundle() throws Exception { - DoFnTester tester = - DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1)); + BundleLiftDoFn fn = new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1); + SimpleMultiOutputReceiver receiver = new SimpleMultiOutputReceiver(); - tester.startBundle(); - tester.processElement(1); - tester.processElement(2); - tester.processElement(3); - tester.finishBundle(); + fn.startBundle(); + fn.processElement(1, receiver); + fn.processElement(2, receiver); + fn.processElement(3, receiver); + fn.finishBundle(); - assertThat(tester.peekOutputElements(INTEGER_SMALL), empty()); - assertThat(tester.peekOutputElements(INTEGER_LARGE), containsInAnyOrder(1, 2, 3)); + assertThat(receiver.getOutput(INTEGER_SMALL), empty()); + assertThat(receiver.getOutput(INTEGER_LARGE), containsInAnyOrder(1, 2, 3)); } @Test public void testSmallBundleWithSizer() throws Exception { - DoFnTester tester = - DoFnTester.of(new BundleLiftDoFn<>(STRING_SMALL, STRING_LARGE, 10, e -> e.length())); - - tester.startBundle(); - tester.processElement("123"); - tester.processElement("456"); - tester.processElement("789"); - tester.finishBundle(); - - assertThat(tester.peekOutputElements(STRING_SMALL), containsInAnyOrder("123", "456", "789")); - assertThat(tester.peekOutputElements(STRING_LARGE), empty()); + BundleLiftDoFn fn = + new BundleLiftDoFn<>(STRING_SMALL, STRING_LARGE, 10, e -> e.length()); + SimpleMultiOutputReceiver receiver = new SimpleMultiOutputReceiver(); + + fn.startBundle(); + fn.processElement("123", receiver); + fn.processElement("456", receiver); + fn.processElement("789", receiver); + fn.finishBundle(); + + assertThat(receiver.getOutput(STRING_SMALL), containsInAnyOrder("123", "456", "789")); + assertThat(receiver.getOutput(STRING_LARGE), empty()); } @Test public void testLargeBundleWithSizer() throws Exception { - DoFnTester tester = - DoFnTester.of(new BundleLiftDoFn<>(STRING_SMALL, STRING_LARGE, 10, e -> e.length())); - - tester.startBundle(); - tester.processElement("123"); - tester.processElement("456"); - tester.processElement("789"); - tester.processElement("0"); - tester.finishBundle(); - - assertThat(tester.peekOutputElements(STRING_SMALL), empty()); - assertThat( - tester.peekOutputElements(STRING_LARGE), containsInAnyOrder("123", "456", "789", "0")); + BundleLiftDoFn fn = + new BundleLiftDoFn<>(STRING_SMALL, STRING_LARGE, 10, e -> e.length()); + SimpleMultiOutputReceiver receiver = new SimpleMultiOutputReceiver(); + + fn.startBundle(); + fn.processElement("123", receiver); + fn.processElement("456", receiver); + fn.processElement("789", receiver); + fn.processElement("0", receiver); + fn.finishBundle(); + + assertThat(receiver.getOutput(STRING_SMALL), empty()); + assertThat(receiver.getOutput(STRING_LARGE), containsInAnyOrder("123", "456", "789", "0")); } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java index b28f5c2c91f8..1b4fd2cc83da 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java @@ -25,9 +25,9 @@ import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes; import org.apache.beam.sdk.schemas.logicaltypes.Timestamp; import org.apache.beam.sdk.testing.PAssert; +import org.apache.beam.sdk.testing.TestOutputReceiver; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.transforms.DoFnTester; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.Reify; import org.apache.beam.sdk.values.PCollection; @@ -116,30 +116,20 @@ public void nullWatermarkValuePreservesInputTimestamp() { public void unsupportedWatermarkTypeThrows() { Schema schema = Schema.builder().addStringField("wm").build(); Row row = Row.withSchema(schema).addValue("2026-05-24T00:00:00Z").build(); + ApplyWatermarkColumn fn = new ApplyWatermarkColumn("wm", "microseconds"); + TestOutputReceiver receiver = new TestOutputReceiver<>(); - assertThrows( - UnsupportedOperationException.class, - () -> { - try (DoFnTester tester = - DoFnTester.of(new ApplyWatermarkColumn("wm", "microseconds"))) { - tester.processElement(row); - } - }); + assertThrows(UnsupportedOperationException.class, () -> fn.process(row, receiver)); } @Test public void missingWatermarkColumnThrows() { Schema schema = Schema.builder().addStringField("other").build(); Row row = Row.withSchema(schema).addValue("value").build(); + ApplyWatermarkColumn fn = new ApplyWatermarkColumn("wm", "microseconds"); + TestOutputReceiver receiver = new TestOutputReceiver<>(); - assertThrows( - IllegalArgumentException.class, - () -> { - try (DoFnTester tester = - DoFnTester.of(new ApplyWatermarkColumn("wm", "microseconds"))) { - tester.processElement(row); - } - }); + assertThrows(IllegalArgumentException.class, () -> fn.process(row, receiver)); } private void assertWatermarkTimestamp( diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java index 2dfbb854b374..acb3e82f7e36 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java @@ -23,6 +23,7 @@ import static org.hamcrest.Matchers.empty; import static org.junit.Assert.assertEquals; +import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.stream.Collectors; @@ -30,12 +31,15 @@ import org.apache.beam.sdk.io.iceberg.IcebergScanConfig; import org.apache.beam.sdk.io.iceberg.IcebergUtils; import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; -import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.OutputBuilder; import org.apache.beam.sdk.values.Row; -import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; @@ -247,11 +251,42 @@ private List> process( List tasks, Instant timestamp) throws Exception { - try (DoFnTester>, Row> tester = - DoFnTester.of(new LocalResolveDoFn(scanConfig))) { - tester.processTimestampedElement(TimestampedValue.of(KV.of(descriptor, tasks), timestamp)); - return tester.getMutableOutput(tester.getMainOutputTag()); - } + LocalResolveDoFn fn = new LocalResolveDoFn(scanConfig); + fn.setup(); + List> outputs = new ArrayList<>(); + OutputReceiver receiver = + new OutputReceiver() { + @Override + public void output(Row output) { + outputWithTimestamp(output, timestamp); + } + + @Override + public void outputWithTimestamp(Row output, Instant ts) { + outputs.add( + ValueInSingleWindow.of(output, ts, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)); + } + + @Override + public OutputBuilder builder(Row output) { + return WindowedValues.builder() + .setValue(output) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setTimestamp(timestamp) + .setReceiver( + wv -> + outputs.add( + ValueInSingleWindow.of( + wv.getValue(), + wv.getTimestamp(), + wv.getWindows().iterator().next(), + wv.getPaneInfo()))); + } + }; + + fn.process(KV.of(descriptor, tasks), receiver); + return outputs; } private TableIdentifier tableId() { diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java index 4d9e709b8d42..daf1cf659539 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java @@ -22,6 +22,7 @@ import static org.hamcrest.Matchers.empty; import static org.junit.Assert.assertEquals; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -31,13 +32,16 @@ import org.apache.beam.sdk.io.iceberg.IcebergUtils; import org.apache.beam.sdk.io.iceberg.TestDataWarehouse; import org.apache.beam.sdk.schemas.Schema; -import org.apache.beam.sdk.transforms.DoFnTester; +import org.apache.beam.sdk.transforms.DoFn.OutputReceiver; import org.apache.beam.sdk.transforms.join.CoGbkResult; +import org.apache.beam.sdk.transforms.windowing.GlobalWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.OutputBuilder; import org.apache.beam.sdk.values.Row; -import org.apache.beam.sdk.values.TimestampedValue; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.sdk.values.WindowedValues; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet; @@ -175,20 +179,51 @@ private List> process( throws Exception { CoGbkResult result = CoGbkResult.of(ResolveChanges.DELETES, deletes).and(ResolveChanges.INSERTS, inserts); - try (DoFnTester, Row> tester = - DoFnTester.of(new ResolveChanges(scanConfig(icebergSchema)))) { - tester.processTimestampedElement( - TimestampedValue.of( - KV.of( - CdcRowDescriptor.builder() - .setCommitSnapshotId(123) - .setSnapshotSequenceNumber(456) - .setPrimaryKey(pk) - .build(), - result), - timestamp)); - return tester.getMutableOutput(tester.getMainOutputTag()); - } + ResolveChanges fn = new ResolveChanges(scanConfig(icebergSchema)); + fn.setup(); + List> outputs = new ArrayList<>(); + OutputReceiver receiver = + new OutputReceiver() { + @Override + public void output(Row output) { + outputWithTimestamp(output, timestamp); + } + + @Override + public void outputWithTimestamp(Row output, Instant ts) { + outputs.add( + ValueInSingleWindow.of(output, ts, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)); + } + + @Override + public OutputBuilder builder(Row output) { + return WindowedValues.builder() + .setValue(output) + .setWindow(GlobalWindow.INSTANCE) + .setPaneInfo(PaneInfo.NO_FIRING) + .setTimestamp(timestamp) + .setReceiver( + wv -> + outputs.add( + ValueInSingleWindow.of( + wv.getValue(), + wv.getTimestamp(), + wv.getWindows().iterator().next(), + wv.getPaneInfo()))); + } + }; + + fn.processElement( + KV.of( + CdcRowDescriptor.builder() + .setCommitSnapshotId(123) + .setSnapshotSequenceNumber(456) + .setPrimaryKey(pk) + .build(), + result), + timestamp, + receiver); + return outputs; } private IcebergScanConfig scanConfig(org.apache.iceberg.Schema icebergSchema) { From aec65ec646517e761ccbb08648ac34d6c18a1e49 Mon Sep 17 00:00:00 2001 From: SuperSaiyan101 Date: Sat, 31 Oct 2026 00:12:55 +0530 Subject: [PATCH 2/2] Address review feedback: propagate ValueKind and clean up imports --- .../beam/sdk/io/iceberg/BundleLifterTest.java | 6 +++-- .../io/iceberg/cdc/LocalResolveDoFnTest.java | 27 ++++++------------- .../io/iceberg/cdc/ResolveChangesTest.java | 27 ++++++------------- 3 files changed, 20 insertions(+), 40 deletions(-) diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java index c7a5a3421ab2..f11204094c54 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java @@ -21,7 +21,9 @@ import static org.hamcrest.Matchers.containsInAnyOrder; import static org.hamcrest.Matchers.empty; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import org.apache.beam.sdk.io.iceberg.BundleLifter.BundleLiftDoFn; import org.apache.beam.sdk.testing.TestOutputReceiver; @@ -62,9 +64,9 @@ private static class SimpleMultiOutputReceiver implements MultiOutputReceiver { } @SuppressWarnings("unchecked") - public java.util.List getOutput(TupleTag tag) { + public List getOutput(TupleTag tag) { TestOutputReceiver receiver = (TestOutputReceiver) receivers.get(tag); - return receiver == null ? java.util.Collections.emptyList() : receiver.getOutputs(); + return receiver == null ? Collections.emptyList() : receiver.getOutputs(); } } diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java index acb3e82f7e36..072c58babf2d 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java @@ -35,7 +35,6 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.OutputBuilder; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.beam.sdk.values.ValueKind; @@ -255,21 +254,8 @@ private List> process( fn.setup(); List> outputs = new ArrayList<>(); OutputReceiver receiver = - new OutputReceiver() { - @Override - public void output(Row output) { - outputWithTimestamp(output, timestamp); - } - - @Override - public void outputWithTimestamp(Row output, Instant ts) { - outputs.add( - ValueInSingleWindow.of(output, ts, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)); - } - - @Override - public OutputBuilder builder(Row output) { - return WindowedValues.builder() + output -> + WindowedValues.builder() .setValue(output) .setWindow(GlobalWindow.INSTANCE) .setPaneInfo(PaneInfo.NO_FIRING) @@ -281,9 +267,12 @@ public OutputBuilder builder(Row output) { wv.getValue(), wv.getTimestamp(), wv.getWindows().iterator().next(), - wv.getPaneInfo()))); - } - }; + wv.getPaneInfo(), + wv.getRecordId(), + wv.getRecordOffset(), + wv.causedByDrain(), + wv.getOpenTelemetryContext(), + wv.getValueKind()))); fn.process(KV.of(descriptor, tasks), receiver); return outputs; diff --git a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java index daf1cf659539..6b0af3fb5e0e 100644 --- a/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java +++ b/sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java @@ -37,7 +37,6 @@ import org.apache.beam.sdk.transforms.windowing.GlobalWindow; import org.apache.beam.sdk.transforms.windowing.PaneInfo; import org.apache.beam.sdk.values.KV; -import org.apache.beam.sdk.values.OutputBuilder; import org.apache.beam.sdk.values.Row; import org.apache.beam.sdk.values.ValueInSingleWindow; import org.apache.beam.sdk.values.ValueKind; @@ -183,21 +182,8 @@ private List> process( fn.setup(); List> outputs = new ArrayList<>(); OutputReceiver receiver = - new OutputReceiver() { - @Override - public void output(Row output) { - outputWithTimestamp(output, timestamp); - } - - @Override - public void outputWithTimestamp(Row output, Instant ts) { - outputs.add( - ValueInSingleWindow.of(output, ts, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING)); - } - - @Override - public OutputBuilder builder(Row output) { - return WindowedValues.builder() + output -> + WindowedValues.builder() .setValue(output) .setWindow(GlobalWindow.INSTANCE) .setPaneInfo(PaneInfo.NO_FIRING) @@ -209,9 +195,12 @@ public OutputBuilder builder(Row output) { wv.getValue(), wv.getTimestamp(), wv.getWindows().iterator().next(), - wv.getPaneInfo()))); - } - }; + wv.getPaneInfo(), + wv.getRecordId(), + wv.getRecordOffset(), + wv.causedByDrain(), + wv.getOpenTelemetryContext(), + wv.getValueKind()))); fn.processElement( KV.of(