Skip to content

[#18690] Port Iceberg IO tests off deprecated DoFnTester #40909

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Supersaiyan101 wants to merge 2 commits into
base: master
Choose a base branch
from
Open

[#18690] Port Iceberg IO tests off deprecated DoFnTester #40909

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
38488d7
[#18690] Port Iceberg IO tests off deprecated DoFnTester
Supersaiyan101 Oct 30, 2026
aec65ec
Address review feedback: propagate ValueKind and clean up imports
Supersaiyan101 Oct 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 77 additions & 42 deletions sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/BundleLifterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -36,64 +46,89 @@ public class BundleLifterTest {
private static final TupleTag<String> STRING_SMALL = new TupleTag<String>() {};
private static final TupleTag<String> STRING_LARGE = new TupleTag<String>() {};

private static class SimpleMultiOutputReceiver implements MultiOutputReceiver {
private final Map<TupleTag<?>, TestOutputReceiver<?>> receivers = new HashMap<>();

@SuppressWarnings("unchecked")
@Override
public @UnknownKeyFor @NonNull @Initialized <T> OutputReceiver<T> get(
@UnknownKeyFor @NonNull @Initialized TupleTag<T> tag) {
return (OutputReceiver<T>) receivers.computeIfAbsent(tag, k -> new TestOutputReceiver<T>());
}

@Override
public @UnknownKeyFor @NonNull @Initialized <T>
OutputReceiver<@UnknownKeyFor @NonNull @Initialized Row> getRowReceiver(
@UnknownKeyFor @NonNull @Initialized TupleTag<T> tag) {
throw new UnsupportedOperationException("getRowReceiver not supported in test");
}

@SuppressWarnings("unchecked")
public <T> List<T> getOutput(TupleTag<T> tag) {
TestOutputReceiver<T> receiver = (TestOutputReceiver<T>) receivers.get(tag);
return receiver == null ? Collections.emptyList() : receiver.getOutputs();
}
}

@Test
public void testSmallBundle() throws Exception {
DoFnTester<Integer, Void> tester =
DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1));
BundleLiftDoFn<Integer> 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<Integer, Void> tester =
DoFnTester.of(new BundleLiftDoFn<>(INTEGER_SMALL, INTEGER_LARGE, 3, x -> 1));
BundleLiftDoFn<Integer> 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<String, Void> 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<String> 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<String, Void> 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<String> 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"));
}
}
24 changes: 7 additions & 17 deletions ...io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ApplyWatermarkColumnTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<Row> receiver = new TestOutputReceiver<>();

assertThrows(
UnsupportedOperationException.class,
() -> {
try (DoFnTester<Row, Row> 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<Row> receiver = new TestOutputReceiver<>();

assertThrows(
IllegalArgumentException.class,
() -> {
try (DoFnTester<Row, Row> tester =
DoFnTester.of(new ApplyWatermarkColumn("wm", "microseconds"))) {
tester.processElement(row);
}
});
assertThrows(IllegalArgumentException.class, () -> fn.process(row, receiver));
}

private void assertWatermarkTimestamp(
Expand Down
38 changes: 31 additions & 7 deletions ...ava/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/LocalResolveDoFnTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,19 +23,22 @@
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;
import org.apache.beam.sdk.io.iceberg.IcebergCatalogConfig;
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;
Expand Down Expand Up @@ -247,11 +250,32 @@ private List<ValueInSingleWindow<Row>> process(
List<SerializableChangelogTask> tasks,
Instant timestamp)
throws Exception {
try (DoFnTester<KV<ChangelogDescriptor, List<SerializableChangelogTask>>, 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<ValueInSingleWindow<Row>> outputs = new ArrayList<>();
OutputReceiver<Row> receiver =
output ->
WindowedValues.<Row>builder()
.setValue(output)
.setWindow(GlobalWindow.INSTANCE)
.setPaneInfo(PaneInfo.NO_FIRING)
.setTimestamp(timestamp)
.setReceiver(
wv ->
outputs.add(
ValueInSingleWindow.of(

@github-actions github-actions Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Critical: LocalResolveDoFn emits CDC records using out.builder(row).setValueKind(kind).output(), and the tests assert on value.getValueKind() (via kindAndProjectedRow). However, the 4-argument ValueInSingleWindow.of(value, timestamp, window, paneInfo) overload hardcodes ValueKind.INSERT, discarding wv.getValueKind(). As a result, all emitted records will have ValueKind.INSERT instead of UPDATE_BEFORE, UPDATE_AFTER, or DELETE, causing the tests in this class to fail.

Use the 9-argument ValueInSingleWindow.of(...) overload and pass wv.getValueKind() (and wv.causedByDrain()). You can also remove the redundant output and outputWithTimestamp overrides since OutputReceiver already provides default methods that delegate to builder(...).

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() {
Expand Down
56 changes: 40 additions & 16 deletions .../java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/cdc/ResolveChangesTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -175,20 +178,41 @@ private List<ValueInSingleWindow<Row>> process(
throws Exception {
CoGbkResult result =
CoGbkResult.of(ResolveChanges.DELETES, deletes).and(ResolveChanges.INSERTS, inserts);
try (DoFnTester<KV<CdcRowDescriptor, CoGbkResult>, 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<ValueInSingleWindow<Row>> outputs = new ArrayList<>();
OutputReceiver<Row> receiver =
output ->
WindowedValues.<Row>builder()
.setValue(output)
.setWindow(GlobalWindow.INSTANCE)
.setPaneInfo(PaneInfo.NO_FIRING)
.setTimestamp(timestamp)
.setReceiver(
wv ->
outputs.add(
ValueInSingleWindow.of(

@github-actions github-actions Bot Oct 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Critical: Similar to LocalResolveDoFnTest, ResolveChanges sets the CDC kind via .setValueKind(kind) on the OutputBuilder, and kindAndData inspects value.getValueKind(). Calling the 4-argument ValueInSingleWindow.of(...) overload drops wv.getValueKind() and defaults to ValueKind.INSERT, which breaks updatePairAndExtraRowsPreserveKindsAndTimestamp.

Pass wv.getValueKind() using the 9-argument ValueInSingleWindow.of(wv.getValue(), wv.getTimestamp(), wv.getWindows().iterator().next(), wv.getPaneInfo(), wv.getRecordId(), wv.getRecordOffset(), wv.causedByDrain(), wv.getOpenTelemetryContext(), wv.getValueKind()) overload (and optionally drop the redundant output/outputWithTimestamp overrides).

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) {
Expand Down
Loading