Code review comment: Maybe a stupid question, but how are the `partitionWriters` connected to the `StreamTask` under the test?

Code under review:
public void testRpcTriggerCheckpointWithSourceChain() throws Exception {
        ResultPartition[] partitionWriters = new ResultPartition[2];
        try {
            for (int i = 0; i < partitionWriters.length; ++i) {
                partitionWriters[i] =
                        PartitionTestUtils.createPartition(ResultPartitionType.PIPELINED_BOUNDED);
                partitionWriters[i].setup();
            }

            try (StreamTaskMailboxTestHarness<String> testHarness =
                    new StreamTaskMailboxTestHarnessBuilder<>(
                                    MultipleInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO)
                            .modifyStreamConfig(config -> config.setCheckpointingEnabled(true))
                            .modifyExecutionConfig(ExecutionConfig::enableObjectReuse)
                            .addInput(BasicTypeInfo.INT_TYPE_INFO)
                            .addInput(BasicTypeInfo.STRING_TYPE_INFO)
                            .addSourceInput(
                                    new SourceOperatorFactory<>(
                                            new MultipleInputStreamTaskTest
                                                    .LifeCycleTrackingMockSource(
                                                    Boundedness.BOUNDED, 1),
                                            WatermarkStrategy.noWatermarks()))
                            .addSourceInput(
                                    new SourceOperatorFactory<>(
                                            new MultipleInputStreamTaskTest
                                                    .LifeCycleTrackingMockSource(
                                                    Boundedness.BOUNDED, 1),
                                            WatermarkStrategy.noWatermarks()))
                            .setupOperatorChain(new MapToStringMultipleInputOperatorFactory(4))
                            .finishForSingletonOperatorChain(StringSerializer.INSTANCE)
                            .build()) {

                testHarness
                        .getStreamTask()
                        .getCheckpointCoordinator()
                        .setEnableCheckpointAfterTasksFinished(true);

                
                

                
                
                new Thread(
                                () -> {
                                    for (ResultPartition partitionWriter : partitionWriters) {
                                        waitTillResultPartitionHasEnoughBuffers(partitionWriter, 2);
                                    }

                                    for (ResultPartition partitionWriter : partitionWriters) {
                                        partitionWriter.onSubpartitionAllRecordsProcessed(0);
                                    }
                                })
                        .start();

                
                testHarness.processEvent(EndOfPartitionEvent.INSTANCE, 0, 0);
                testHarness.processEvent(EndOfPartitionEvent.INSTANCE, 1, 0);
                Future<Boolean> checkpointFuture = triggerCheckpoint(testHarness, 4);

                
                
                
                testHarness.finishProcessing();
                assertTrue(checkpointFuture.isDone());
            }
        } finally {
            for (ResultPartitionWriter writer : partitionWriters) {
                if (writer != null) {
                    writer.close();
                }
            }
        }
    }

Line commented on: for (int i = 0; i < partitionWriters.length; ++i) {

Revised code:
public void testRpcTriggerCheckpointWithSourceChain() throws Exception {
        ResultPartition[] partitionWriters = new ResultPartition[2];
        try {
            for (int i = 0; i < partitionWriters.length; ++i) {
                partitionWriters[i] =
                        PartitionTestUtils.createPartition(ResultPartitionType.PIPELINED_BOUNDED);
                partitionWriters[i].setup();
            }

            try (StreamTaskMailboxTestHarness<String> testHarness =
                    new StreamTaskMailboxTestHarnessBuilder<>(
                                    MultipleInputStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO)
                            .modifyStreamConfig(config -> config.setCheckpointingEnabled(true))
                            .modifyExecutionConfig(ExecutionConfig::enableObjectReuse)
                            .addInput(BasicTypeInfo.INT_TYPE_INFO)
                            .addInput(BasicTypeInfo.STRING_TYPE_INFO)
                            .addSourceInput(
                                    new SourceOperatorFactory<>(
                                            new MultipleInputStreamTaskTest
                                                    .LifeCycleTrackingMockSource(
                                                    Boundedness.BOUNDED, 1),
                                            WatermarkStrategy.noWatermarks()))
                            .addSourceInput(
                                    new SourceOperatorFactory<>(
                                            new MultipleInputStreamTaskTest
                                                    .LifeCycleTrackingMockSource(
                                                    Boundedness.BOUNDED, 1),
                                            WatermarkStrategy.noWatermarks()))
                            .addAdditionalOutput(partitionWriters)
                            .setupOperatorChain(new MapToStringMultipleInputOperatorFactory(4))
                            .finishForSingletonOperatorChain(StringSerializer.INSTANCE)
                            .build()) {

                testHarness
                        .getStreamTask()
                        .getCheckpointCoordinator()
                        .setEnableCheckpointAfterTasksFinished(true);
                testHarness
                        .getStreamTask()
                        .getCheckpointBarrierHandler()
                        .get()
                        .setEnableCheckpointAfterTasksFinished(true);

                Future<Boolean> checkpointFuture = triggerCheckpoint(testHarness, 2);
                testHarness.processAll();

                
                testHarness.processEvent(EndOfPartitionEvent.INSTANCE, 0, 0);
                testHarness.processEvent(EndOfPartitionEvent.INSTANCE, 1, 0);
                testHarness.getTaskStateManager().getWaitForReportLatch().await();
                assertEquals(2, testHarness.getTaskStateManager().getReportedCheckpointId());

                
                checkpointFuture = triggerCheckpoint(testHarness, 4);

                
                
                checkState(
                        checkpointFuture instanceof CompletableFuture,
                        "The trigger future should " + " be also CompletableFuture.");
                ((CompletableFuture<?>) checkpointFuture)
                        .thenAccept(
                                (ignored) -> {
                                    for (ResultPartition resultPartition : partitionWriters) {
                                        resultPartition.onSubpartitionAllRecordsProcessed(0);
                                    }
                                });

                
                
                
                testHarness.finishProcessing();
                assertTrue(checkpointFuture.isDone());

                
                for (ResultPartition resultPartition : partitionWriters) {
                    assertEquals(3, resultPartition.getNumberOfQueuedBuffers());
                }
            }
        } finally {
            for (ResultPartitionWriter writer : partitionWriters) {
                if (writer != null) {
                    writer.close();
                }
            }
        }
    }