@Override
    protected CompletableFuture<Void> getCompletionFuture() {
        return sourceThread.getCompletionFuture();
    }
@Test
    void testTriggeringCheckpointAfterSourceThreadFinished() throws Exception {
        ResultPartition[] partitionWriters = new ResultPartition[2];
        try (NettyShuffleEnvironment env =
                new NettyShuffleEnvironmentBuilder()
                        .setNumNetworkBuffers(partitionWriters.length * 2)
                        .build()) {
            for (int i = 0; i < partitionWriters.length; ++i) {
                partitionWriters[i] =
                        PartitionTestUtils.createPartition(
                                env, ResultPartitionType.PIPELINED_BOUNDED, 1);
                partitionWriters[i].setup();
            }

            final CompletableFuture<Long> checkpointCompleted = new CompletableFuture<>();
            try (StreamTaskMailboxTestHarness<String> testHarness =
                    new StreamTaskMailboxTestHarnessBuilder<>(
                                    SourceStreamTask::new, BasicTypeInfo.STRING_TYPE_INFO)
                            .modifyStreamConfig(config -> config.setCheckpointingEnabled(true))
                            .setCheckpointResponder(
                                    new TestCheckpointResponder() {
                                        @Override
                                        public void acknowledgeCheckpoint(
                                                JobID jobID,
                                                ExecutionAttemptID executionAttemptID,
                                                long checkpointId,
                                                CheckpointMetrics checkpointMetrics,
                                                TaskStateSnapshot subtaskState) {
                                            super.acknowledgeCheckpoint(
                                                    jobID,
                                                    executionAttemptID,
                                                    checkpointId,
                                                    checkpointMetrics,
                                                    subtaskState);
                                            checkpointCompleted.complete(checkpointId);
                                        }
                                    })
                            .addAdditionalOutput(partitionWriters)
                            .setupOperatorChain(new StreamSource<>(new MockSource(0, 0, 1)))
                            .finishForSingletonOperatorChain(StringSerializer.INSTANCE)
                            .build()) {

                testHarness.processAll();
                CompletableFuture<Void> taskFinished =
                        testHarness.getStreamTask().getCompletionFuture();
                do {
                    testHarness.processAll();
                } while (!taskFinished.isDone());

                Future<Boolean> checkpointFuture = triggerCheckpoint(testHarness, 2);
                // Notifies the result partition that all records are processed after the
                // last checkpoint is triggered.
                checkState(
                        checkpointFuture instanceof CompletableFuture,
                        "The trigger future should " + " be also CompletableFuture.");
                ((CompletableFuture<?>) checkpointFuture)
                        .thenAccept(
                                (ignored) -> {
                                    for (ResultPartition resultPartition : partitionWriters) {
                                        resultPartition.onSubpartitionAllDataProcessed(0);
                                    }
                                });

                checkpointCompleted.whenComplete(
                        (id, error) ->
                                testHarness.getStreamTask().notifyCheckpointCompleteAsync(2));
                testHarness.finishProcessing();
                assertThat(checkpointFuture.isDone()).isTrue();

                // Each result partition should have emitted 1 barrier, 1 max watermark and 1
                // EndOfUserRecordEvent.
                for (ResultPartition resultPartition : partitionWriters) {
                    assertThat(resultPartition.getNumberOfQueuedBuffers()).isEqualTo(3);
                }
            }
        } finally {
            for (ResultPartitionWriter writer : partitionWriters) {
                if (writer != null) {
                    writer.close();
                }
            }
        }
    }