Repository: apache/flink
File: flink-runtime/src/test/java/org/apache/flink/streaming/api/operators/SourceOperatorTest.java (Java)

Code before:
                                2L,
                                CheckpointOptions.alignedNoTimeout(
                                        CheckpointType.CHECKPOINT,
                                        CheckpointStorageLocationReference.getDefault()),
                                new MemCheckpointStreamFactory(10240));
                        operator.emitNext(new DataOutputToOutput<>(operator.output));
                    }

                    assertThat(operator.isAvailable()).isTrue();
                    assertThat(out.stream().map(element -> element.asRecord().getValue()))
                            .containsExactly(0);
                });
    }

    @Test
    public void testFinalSavepointRestoredSourceOperator() throws Exception {
        testRestoredSourceOperator(
                (operator, out) -> {
                    DataOutputToOutput<Integer> output = new DataOutputToOutput<>(operator.output);
                    operator.stop(StopMode.NO_DRAIN); // emulate stop with savepoint
                    DataInputStatus status;
                    do {
                        status = operator.emitNext(output);
                    } while (status != END_OF_INPUT);
                });
    }

    private void testRestoredSourceOperator(
            BiConsumerWithException<
                            SourceOperator<Integer, MockSourceSplit>,
                            List<StreamElement>,
                            Exception>
                    test)
            throws Exception {
        final List<StreamElement> out = new ArrayList<>();
        try (SourceOperatorTestContext context =
                SourceOperatorTestContext.builder()
                        .setWatermarkStrategy(
                                WatermarkStrategy.<Integer>forMonotonousTimestamps()
                                        .withWatermarkAlignment("ag-1", Duration.ofMillis(50))
                                        .withTimestampAssigner(
                                                (element, recordTimestamp) -> element))
                        .setOutput(new CollectorOutput<>(out))
                        .setPauseSourcesUntilFirstCheckpoint(pauseSourcesUntilCheckpoint)
                        .setPreInit(
                                // recover with some state, so the source will pause until a
                                // checkpoint
                                // to speedup recovery (if pauseSourcesUntilCheckpoint)
                                (stateManager, operatorID) -> {
                                    long checkpointID = 1L;
                                    stateManager.setReportedCheckpointId(checkpointID);

Review comment (bug):
Does this assert something? I presume it checks against some exception failure, but if so:

- adding some assertions if possible would be still great
- maybe explaining what's the purpose of the test?

Code after:
                                2L,
                                CheckpointOptions.alignedNoTimeout(
                                        CheckpointType.CHECKPOINT,
                                        CheckpointStorageLocationReference.getDefault()),
                                new MemCheckpointStreamFactory(10240));
                        operator.emitNext(new DataOutputToOutput<>(operator.output));
                    }

                    assertThat(operator.isAvailable()).isTrue();
                    assertThat(out.stream().map(element -> element.asRecord().getValue()))
                            .containsExactly(0);
                });
    }

    @TestTemplate
    public void testFinalSavepointRestoredSourceOperator() throws Exception {
        testRestoredSourceOperator(
                (operator, out) -> {
                    DataOutputToOutput<Integer> output = new DataOutputToOutput<>(operator.output);
                    operator.stop(StopMode.NO_DRAIN); // emulate stop with savepoint
                    DataInputStatus status;
                    do {
                        status = operator.emitNext(output); // should not fail
                    } while (status != END_OF_INPUT);
                });
    }

    private void testRestoredSourceOperator(
            BiConsumerWithException<
                            SourceOperator<Integer, MockSourceSplit>,
                            List<StreamElement>,
                            Exception>
                    test)
            throws Exception {
        final List<StreamElement> out = new ArrayList<>();
        try (SourceOperatorTestContext context =
                SourceOperatorTestContext.builder()
                        .setWatermarkStrategy(
                                WatermarkStrategy.<Integer>forMonotonousTimestamps()
                                        .withWatermarkAlignment("ag-1", Duration.ofMillis(50))
                                        .withTimestampAssigner(
                                                (element, recordTimestamp) -> element))
                        .setOutput(new CollectorOutput<>(out))
                        .setPauseSourcesUntilFirstCheckpoint(pauseSourcesUntilCheckpoint)
                        .setPreInit(
                                // recover with some state, so the source will pause until a
                                // checkpoint
                                // to speedup recovery (if pauseSourcesUntilCheckpoint)
                                (stateManager, operatorID) -> {
                                    long checkpointID = 1L;
                                    stateManager.setReportedCheckpointId(checkpointID);