Source:
public void testCreateCheckpointedMultipleInputGate() throws Exception { try (CloseableRegistry registry = new CloseableRegistry()) { MockEnvironment environment = new MockEnvironmentBuilder().build(); MockStreamTask streamTask = new MockStreamTaskBuilder(environment).build(); StreamConfig streamConfig = new StreamConfig(environment.getJobConfiguration()); streamConfig.setCheckpointMode(CheckpointingMode.EXACTLY_ONCE); streamConfig.setUnalignedCheckpointsEnabled(true); Collection<IndexedInputGate>[] inputGates = new Collection[] { Collections.singletonList(new MockIndexedInputGate(1, 4)), Collections.singletonList(new MockIndexedInputGate(0, 2)), }; new MockChannelStateWriter() { @Override public void addInputData( long checkpointId, InputChannelInfo info, int startSeqNum, Buffer... data) { super.addInputData(checkpointId, info, startSeqNum, data); } }; CheckpointedInputGate[] checkpointedMultipleInputGate = InputProcessorUtil.createCheckpointedMultipleInputGate( streamTask, streamConfig, new MockChannelStateWriter(), environment.getMetricGroup().getIOMetricGroup(), streamTask.getName(), inputGates); for (CheckpointedInputGate checkpointedInputGate : checkpointedMultipleInputGate) { registry.registerCloseable(checkpointedInputGate); } CheckpointBarrierHandler barrierHandler = checkpointedMultipleInputGate[0].getCheckpointBarrierHandler(); List<IndexedInputGate> allInputGates = Arrays.stream(inputGates).flatMap(gates -> gates.stream()).collect(Collectors.toList()); for (IndexedInputGate inputGate : allInputGates) { for (int channelId = 0; channelId < inputGate.getNumberOfInputChannels(); channelId++) { assertTrue(barrierHandler.getBufferReceivedListener().isPresent()); BufferReceivedListener bufferReceivedListener = barrierHandler.getBufferReceivedListener().get(); bufferReceivedListener.notifyBarrierReceived( new CheckpointBarrier(1, 42, CheckpointOptions.forCheckpointWithDefaultLocation(true, true)), new InputChannelInfo(inputGate.getGateIndex(), channelId)); } } assertFalse(barrierHandler.isCheckpointPending()); } }

Target:
public void testCreateCheckpointedMultipleInputGate() throws Exception { try (CloseableRegistry registry = new CloseableRegistry()) { MockEnvironment environment = new MockEnvironmentBuilder().build(); MockStreamTask streamTask = new MockStreamTaskBuilder(environment).build(); StreamConfig streamConfig = new StreamConfig(environment.getJobConfiguration()); streamConfig.setCheckpointMode(CheckpointingMode.EXACTLY_ONCE); streamConfig.setUnalignedCheckpointsEnabled(true); Collection<IndexedInputGate>[] inputGates = new Collection[] { Collections.singletonList(new MockIndexedInputGate(1, 4)), Collections.singletonList(new MockIndexedInputGate(0, 2)), }; CheckpointedInputGate[] checkpointedMultipleInputGate = InputProcessorUtil.createCheckpointedMultipleInputGate( streamTask, streamConfig, new MockChannelStateWriter(), environment.getMetricGroup().getIOMetricGroup(), streamTask.getName(), inputGates); for (CheckpointedInputGate checkpointedInputGate : checkpointedMultipleInputGate) { registry.registerCloseable(checkpointedInputGate); } CheckpointBarrierHandler barrierHandler = checkpointedMultipleInputGate[0].getCheckpointBarrierHandler(); assertTrue(barrierHandler.getBufferReceivedListener().isPresent()); BufferReceivedListener bufferReceivedListener = barrierHandler.getBufferReceivedListener().get(); List<IndexedInputGate> allInputGates = Arrays.stream(inputGates).flatMap(gates -> gates.stream()).collect(Collectors.toList()); for (IndexedInputGate inputGate : allInputGates) { for (int channelId = 0; channelId < inputGate.getNumberOfInputChannels(); channelId++) { bufferReceivedListener.notifyBarrierReceived( new CheckpointBarrier(1, 42, CheckpointOptions.forCheckpointWithDefaultLocation(true, true)), new InputChannelInfo(inputGate.getGateIndex(), channelId)); } } assertTrue(barrierHandler.getAllBarriersReceivedFuture(1).isDone()); } }