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 <[email protected]>
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});
