[FLINK-36892][Test] Use record processor in operator test harness

Zakelly · Dec 13, 2024 · 4df0aa6 · 4df0aa6
1 parent 58de480
commit 4df0aa6
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 14 deletions.
diff --git a/...test/java/org/apache/flink/streaming/util/AbstractBroadcastStreamOperatorTestHarness.java b/...test/java/org/apache/flink/streaming/util/AbstractBroadcastStreamOperatorTestHarness.java
@@ -21,6 +21,7 @@
 import org.apache.flink.streaming.api.operators.StreamOperator;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.io.RecordProcessorUtils;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 
 /** Base class for broadcast stream operator test harnesses. */
@@ -38,17 +39,15 @@ public TwoInputStreamOperator<IN1, IN2, OUT> getTwoInputOperator() {
     }
 
     public void processElement(StreamRecord<IN1> element) throws Exception {
-        getTwoInputOperator().setKeyContextElement1(element);
-        getTwoInputOperator().processElement1(element);
+        RecordProcessorUtils.getRecordProcessor1(getTwoInputOperator()).accept(element);
     }
 
     public void processElement(IN1 value, long timestamp) throws Exception {
         processElement(new StreamRecord<>(value, timestamp));
     }
 
     public void processBroadcastElement(StreamRecord<IN2> element) throws Exception {
-        getTwoInputOperator().setKeyContextElement2(element);
-        getTwoInputOperator().processElement2(element);
+        RecordProcessorUtils.getRecordProcessor2(getTwoInputOperator()).accept(element);
     }
 
     public void processBroadcastElement(IN2 value, long timestamp) throws Exception {

diff --git a/...me/src/test/java/org/apache/flink/streaming/util/MultiInputStreamOperatorTestHarness.java b/...me/src/test/java/org/apache/flink/streaming/util/MultiInputStreamOperatorTestHarness.java
@@ -22,9 +22,11 @@
 import org.apache.flink.streaming.api.operators.MultipleInputStreamOperator;
 import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.io.RecordProcessorUtils;
 import org.apache.flink.streaming.runtime.streamrecord.RecordAttributes;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 /**
  * A test harness for testing a {@link MultipleInputStreamOperator}.
@@ -52,8 +54,8 @@ public MultiInputStreamOperatorTestHarness(
 
     public void processElement(int idx, StreamRecord<?> element) throws Exception {
         Input input = getCastedOperator().getInputs().get(idx);
-        input.setKeyContextElement(element);
-        input.processElement(element);
+        ((ThrowingConsumer<StreamRecord, Exception>) RecordProcessorUtils.getRecordProcessor(input))
+                .accept(element);
     }
 
     public void processWatermark(int idx, Watermark mark) throws Exception {

diff --git a/...time/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java b/...time/src/test/java/org/apache/flink/streaming/util/OneInputStreamOperatorTestHarness.java
@@ -29,10 +29,12 @@
 import org.apache.flink.streaming.api.operators.SimpleOperatorFactory;
 import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.io.RecordProcessorUtils;
 import org.apache.flink.streaming.runtime.streamrecord.RecordAttributes;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
 import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.function.ThrowingConsumer;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -213,13 +215,13 @@ public void processElement(IN value, long timestamp) throws Exception {
 
     public void processElement(StreamRecord<IN> element) throws Exception {
         if (inputs.isEmpty()) {
-            operator.setKeyContextElement1(element);
-            getOneInputOperator().processElement(element);
+            RecordProcessorUtils.getRecordProcessor(getOneInputOperator()).accept(element);
         } else {
             checkState(inputs.size() == 1);
             Input input = inputs.get(0);
-            input.setKeyContextElement(element);
-            input.processElement(element);
+            ((ThrowingConsumer<StreamRecord, Exception>)
+                            RecordProcessorUtils.getRecordProcessor(input))
+                    .accept(element);
         }
     }
 

diff --git a/...time/src/test/java/org/apache/flink/streaming/util/TwoInputStreamOperatorTestHarness.java b/...time/src/test/java/org/apache/flink/streaming/util/TwoInputStreamOperatorTestHarness.java
@@ -21,6 +21,7 @@
 import org.apache.flink.streaming.api.operators.BoundedMultiInput;
 import org.apache.flink.streaming.api.operators.TwoInputStreamOperator;
 import org.apache.flink.streaming.api.watermark.Watermark;
+import org.apache.flink.streaming.runtime.io.RecordProcessorUtils;
 import org.apache.flink.streaming.runtime.streamrecord.RecordAttributes;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.runtime.watermarkstatus.WatermarkStatus;
@@ -54,17 +55,15 @@ public TwoInputStreamOperatorTestHarness(
     }
 
     public void processElement1(StreamRecord<IN1> element) throws Exception {
-        twoInputOperator.setKeyContextElement1(element);
-        twoInputOperator.processElement1(element);
+        RecordProcessorUtils.getRecordProcessor1(twoInputOperator).accept(element);
     }
 
     public void processElement1(IN1 value, long timestamp) throws Exception {
         processElement1(new StreamRecord<>(value, timestamp));
     }
 
     public void processElement2(StreamRecord<IN2> element) throws Exception {
-        twoInputOperator.setKeyContextElement2(element);
-        twoInputOperator.processElement2(element);
+        RecordProcessorUtils.getRecordProcessor2(twoInputOperator).accept(element);
     }
 
     public void processElement2(IN2 value, long timestamp) throws Exception {