Examples with TestTaskStateManagerBuilder org.apache.flink.runtime.state.TestTaskStateManagerBuilder used on opensource projects

Example 1 with TestTaskStateManagerBuilder

use of org.apache.flink.runtime.state.TestTaskStateManagerBuilder in project flink by apache.

the class StreamTaskMailboxTestHarnessBuilder method buildUnrestored.

public StreamTaskMailboxTestHarness<OUT> buildUnrestored() throws Exception {
    TestTaskStateManagerBuilder taskStateManagerBuilder = TestTaskStateManager.builder().setLocalRecoveryConfig(localRecoveryConfig).setCheckpointResponder(checkpointResponder);
    if (taskStateSnapshots != null) {
        taskStateManagerBuilder.setReportedCheckpointId(taskStateSnapshots.keySet().iterator().next()).setJobManagerTaskStateSnapshotsByCheckpointId(taskStateSnapshots);
    }
    TestTaskStateManager taskStateManager = taskStateManagerBuilder.build();
    StreamMockEnvironment streamMockEnvironment = new StreamMockEnvironment(new JobID(), new ExecutionAttemptID(), jobConfig, streamConfig.getConfiguration(), executionConfig, memorySize, new MockInputSplitProvider(), bufferSize, taskStateManager, collectNetworkEvents);
    streamMockEnvironment.setCheckpointResponder(taskStateManager.getCheckpointResponder());
    streamMockEnvironment.setTaskManagerInfo(taskManagerRuntimeInfo);
    initializeInputs(streamMockEnvironment);
    checkState(inputGates != null, "InputGates hasn't been initialised");
    StreamElementSerializer<OUT> outputStreamRecordSerializer = new StreamElementSerializer<>(outputSerializer);
    Queue<Object> outputList = new ArrayDeque<>();
    streamMockEnvironment.addOutput(outputList, outputStreamRecordSerializer);
    streamMockEnvironment.setTaskMetricGroup(taskMetricGroup);
    for (ResultPartitionWriter writer : additionalOutputs) {
        streamMockEnvironment.addOutput(writer);
    }
    StreamTask<OUT, ?> task = taskFactory.apply(streamMockEnvironment);
    return new StreamTaskMailboxTestHarness<>(task, outputList, inputGates, streamMockEnvironment);
}