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..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,9 +21,19 @@ 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.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 +46,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 List getOutput(TupleTag tag) { + TestOutputReceiver receiver = (TestOutputReceiver) receivers.get(tag); + return receiver == null ? 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..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 @@ -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,14 @@ 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.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 +250,32 @@ 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 = + output -> + 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(), + wv.getRecordId(), + wv.getRecordOffset(), + wv.causedByDrain(), + wv.getOpenTelemetryContext(), + wv.getValueKind()))); + + 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..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 @@ -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,15 @@ 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.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 +178,41 @@ 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 = + output -> + 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(), + wv.getRecordId(), + wv.getRecordOffset(), + wv.causedByDrain(), + wv.getOpenTelemetryContext(), + wv.getValueKind()))); + + 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) {