You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by ro...@apache.org on 2026/07/13 18:35:40 UTC

(flink) branch master updated: [FLINK-40036][tests] Support custom MockEnvironment in multi-input harness and state-parameterized tests (#28589)

This is an automated email from the ASF dual-hosted git repository.

rkhachatryan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git


The following commit(s) were added to refs/heads/master by this push:
     new c404d7b08e4 [FLINK-40036][tests] Support custom MockEnvironment in multi-input harness and state-parameterized tests (#28589)
c404d7b08e4 is described below

commit c404d7b08e4fcc97c892281f597148e398955552
Author: Kaivalya Apte <ap...@gmail.com>
AuthorDate: Mon Jul 13 20:35:32 2026 +0200

    [FLINK-40036][tests] Support custom MockEnvironment in multi-input harness and state-parameterized tests (#28589)
---
 .../operators/testutils/MockEnvironment.java       |  5 +--
 .../testutils/MockEnvironmentBuilder.java          |  9 ++++-
 .../KeyedMultiInputStreamOperatorTestHarness.java  | 12 +++++++
 .../util/MultiInputStreamOperatorTestHarness.java  |  7 ++++
 .../StreamingMultiJoinOperatorTestBase.java        |  6 +++-
 .../util/StateParameterizedHarnessTestBase.java    | 41 ++++++++++++++++++++++
 6 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
index c50eface312..101ff4386cb 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironment.java
@@ -179,13 +179,14 @@ public class MockEnvironment implements Environment, AutoCloseable {
             TaskManagerRuntimeInfo taskManagerRuntimeInfo,
             MemoryManager memManager,
             ExternalResourceInfoProvider externalResourceInfoProvider,
-            ChannelStateWriteRequestExecutorFactory channelStateExecutorFactory) {
+            ChannelStateWriteRequestExecutorFactory channelStateExecutorFactory,
+            Configuration jobConfiguration) {
 
         this.jobInfo = new JobInfoImpl(jobID, jobName);
         this.jobVertexID = jobVertexID;
         this.jobType = jobType;
         this.taskInfo = new TaskInfoImpl(taskName, maxParallelism, subtaskIndex, parallelism, 0);
-        this.jobConfiguration = new Configuration();
+        this.jobConfiguration = jobConfiguration;
         this.taskConfiguration = taskConfiguration;
         this.inputs = new LinkedList<>();
         this.outputs = new LinkedList<ResultPartitionWriter>();
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java
index 2a8c7c273a8..c69a1e275e2 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/operators/testutils/MockEnvironmentBuilder.java
@@ -67,6 +67,7 @@ public class MockEnvironmentBuilder {
             ExternalResourceInfoProvider.NO_EXTERNAL_RESOURCES;
     private ChannelStateWriteRequestExecutorFactory channelStateExecutorFactory =
             new ChannelStateWriteRequestExecutorFactory(jobID);
+    private Configuration jobConfiguration = new Configuration();
 
     private MemoryManager buildMemoryManager(long memorySize) {
         return MemoryManagerBuilder.newBuilder().setMemorySize(memorySize).build();
@@ -181,6 +182,11 @@ public class MockEnvironmentBuilder {
         return this;
     }
 
+    public MockEnvironmentBuilder setJobConfiguration(Configuration jobConfiguration) {
+        this.jobConfiguration = jobConfiguration;
+        return this;
+    }
+
     public MockEnvironment build() {
         if (ioManager == null) {
             ioManager = new IOManagerAsync();
@@ -206,6 +212,7 @@ public class MockEnvironmentBuilder {
                 taskManagerRuntimeInfo,
                 memoryManager,
                 externalResourceInfoProvider,
-                channelStateExecutorFactory);
+                channelStateExecutorFactory,
+                jobConfiguration);
     }
 }
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/util/KeyedMultiInputStreamOperatorTestHarness.java b/flink-runtime/src/test/java/org/apache/flink/streaming/util/KeyedMultiInputStreamOperatorTestHarness.java
index 9d3f5fa431b..15ee270fe85 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/util/KeyedMultiInputStreamOperatorTestHarness.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/util/KeyedMultiInputStreamOperatorTestHarness.java
@@ -22,6 +22,7 @@ import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.ClosureCleaner;
 import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
 import org.apache.flink.streaming.api.operators.MultipleInputStreamOperator;
 import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
 
@@ -43,6 +44,17 @@ public class KeyedMultiInputStreamOperatorTestHarness<KEY, OUT>
         config.serializeAllConfigs();
     }
 
+    public KeyedMultiInputStreamOperatorTestHarness(
+            StreamOperatorFactory<OUT> operator,
+            TypeInformation<KEY> keyType,
+            MockEnvironment environment)
+            throws Exception {
+        super(operator, environment);
+        config.setStateKeySerializer(
+                keyType.createSerializer(executionConfig.getSerializerConfig()));
+        config.serializeAllConfigs();
+    }
+
     public KeyedMultiInputStreamOperatorTestHarness(
             StreamOperatorFactory<OUT> operatorFactory,
             int maxParallelism,
diff --git a/flink-runtime/src/test/java/org/apache/flink/streaming/util/MultiInputStreamOperatorTestHarness.java b/flink-runtime/src/test/java/org/apache/flink/streaming/util/MultiInputStreamOperatorTestHarness.java
index 058fc57ee8b..95f3c0bf57a 100644
--- a/flink-runtime/src/test/java/org/apache/flink/streaming/util/MultiInputStreamOperatorTestHarness.java
+++ b/flink-runtime/src/test/java/org/apache/flink/streaming/util/MultiInputStreamOperatorTestHarness.java
@@ -18,6 +18,7 @@
 
 package org.apache.flink.streaming.util;
 
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
 import org.apache.flink.streaming.api.operators.Input;
 import org.apache.flink.streaming.api.operators.MultipleInputStreamOperator;
 import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
@@ -50,6 +51,12 @@ public class MultiInputStreamOperatorTestHarness<OUT>
         super(operatorFactory, maxParallelism, numSubtasks, subtaskIndex);
     }
 
+    public MultiInputStreamOperatorTestHarness(
+            StreamOperatorFactory<OUT> operatorFactory, MockEnvironment environment)
+            throws Exception {
+        super(operatorFactory, environment);
+    }
+
     public void processElement(int idx, StreamRecord<?> element) throws Exception {
         Input input = getCastedOperator().getInputs().get(idx);
         input.setKeyContextElement(element);
diff --git a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/multijoin/StreamingMultiJoinOperatorTestBase.java b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/multijoin/StreamingMultiJoinOperatorTestBase.java
index 27861fafb11..4b48ec36941 100644
--- a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/multijoin/StreamingMultiJoinOperatorTestBase.java
+++ b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/operators/join/stream/multijoin/StreamingMultiJoinOperatorTestBase.java
@@ -20,6 +20,7 @@ package org.apache.flink.table.runtime.operators.join.stream.multijoin;
 
 import org.apache.flink.api.common.functions.AbstractRichFunction;
 import org.apache.flink.api.java.functions.KeySelector;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.streaming.util.KeyedMultiInputStreamOperatorTestHarness;
 import org.apache.flink.table.data.RowData;
@@ -157,7 +158,9 @@ public abstract class StreamingMultiJoinOperatorTestBase extends StateParameteri
     @AfterEach
     protected void afterEach() throws Exception {
         if (testHarness != null) {
+            MockEnvironment environment = testHarness.getEnvironment();
             testHarness.close();
+            environment.close();
         }
     }
 
@@ -404,7 +407,8 @@ public abstract class StreamingMultiJoinOperatorTestBase extends StateParameteri
                         this.joinAttributeMap);
 
         KeyedMultiInputStreamOperatorTestHarness<RowData, RowData> harness =
-                new KeyedMultiInputStreamOperatorTestHarness<>(factory, partitionKeyTypeInfo);
+                new KeyedMultiInputStreamOperatorTestHarness<>(
+                        factory, partitionKeyTypeInfo, createMockEnvironment());
 
         setupKeySelectorsForTestHarness(harness);
 
diff --git a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/util/StateParameterizedHarnessTestBase.java b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/util/StateParameterizedHarnessTestBase.java
index 4919bb5a09a..269089b36e9 100644
--- a/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/util/StateParameterizedHarnessTestBase.java
+++ b/flink-table/flink-table-runtime/src/test/java/org/apache/flink/table/runtime/util/StateParameterizedHarnessTestBase.java
@@ -18,7 +18,12 @@
 
 package org.apache.flink.table.runtime.util;
 
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.configuration.CheckpointingOptions;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.configuration.MemorySize;
+import org.apache.flink.runtime.operators.testutils.MockEnvironment;
+import org.apache.flink.runtime.operators.testutils.MockEnvironmentBuilder;
 import org.apache.flink.runtime.state.CheckpointStorage;
 import org.apache.flink.runtime.state.StateBackend;
 import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
@@ -37,6 +42,13 @@ import java.util.Collection;
 
 public abstract class StateParameterizedHarnessTestBase {
 
+    /**
+     * Managed memory for the harness environment. Sized for state backends that reserve managed
+     * memory per keyed-state DB (the harness default of 3 MB is too small for them); harmless for
+     * the others.
+     */
+    private static final MemorySize MANAGED_MEMORY = MemorySize.ofMebiBytes(128);
+
     @Rule public TemporaryFolder tempFolder = new TemporaryFolder();
 
     public enum StateBackendMode {
@@ -95,6 +107,35 @@ public abstract class StateParameterizedHarnessTestBase {
         }
     }
 
+    /**
+     * Builds the {@link MockEnvironment} the test harness is constructed with. The harness would
+     * otherwise build its own with only 3 MB of managed memory and no checkpoint storage. We
+     * configure it the same way for every backend: enough managed memory, and a filesystem
+     * checkpoint storage (so a backend that resolves a task-owned state directory from it works).
+     */
+    protected MockEnvironment createMockEnvironment() {
+        String checkpointPath = "file://" + tempFolder.getRoot().getAbsolutePath();
+
+        Configuration jobConfig = new Configuration();
+        jobConfig.set(CheckpointingOptions.CHECKPOINTS_DIRECTORY, checkpointPath);
+
+        MockEnvironment environment =
+                new MockEnvironmentBuilder()
+                        .setManagedMemorySize(MANAGED_MEMORY.getBytes())
+                        .setJobConfiguration(jobConfig)
+                        .build();
+
+        try {
+            environment.setCheckpointStorageAccess(
+                    new FileSystemCheckpointStorage(checkpointPath)
+                            .createCheckpointStorage(new JobID()));
+        } catch (IOException e) {
+            throw new RuntimeException(
+                    "Cannot create checkpoint storage for the test environment", e);
+        }
+        return environment;
+    }
+
     @Parameters(name = "StateBackend={0}")
     public static Collection<Object[]> parameters() {
         return Arrays.asList(new Object[] {HEAP_BACKEND}, new Object[] {ROCKSDB_BACKEND});