Repository: apache/flink
File: flink-tests/src/test/java/org/apache/flink/test/checkpointing/UnalignedCheckpointRescaleITCase.java (Java)

Code before:
    public UnalignedCheckpointRescaleITCase(
            String desc,
            Topology topology,
            int oldParallelism,
            int newParallelism,
            long sourceSleepMs) {
        this.topology = topology;
        this.oldParallelism = oldParallelism;
        this.newParallelism = newParallelism;
        this.sourceSleepMs = sourceSleepMs;
    }

    @Test
    public void shouldRescaleUnalignedCheckpoint() throws Exception {
        final UnalignedSettings prescaleSettings =
                new UnalignedSettings(topology)
                        .setParallelism(oldParallelism)
                        .setExpectedFailures(1)
                        .setSourceSleepMs(sourceSleepMs);
        prescaleSettings.setGenerateCheckpoint(true);
        final File checkpointDir = super.execute(prescaleSettings);

        // resume
        final UnalignedSettings postscaleSettings =
                new UnalignedSettings(topology)
                        .setParallelism(newParallelism)
                        .setExpectedFailures(1);
        postscaleSettings.setRestoreCheckpoint(checkpointDir);
        super.execute(postscaleSettings);
    }

    protected void checkCounters(JobExecutionResult result) {
        collector.checkThat(
                "NUM_OUTPUTS = NUM_INPUTS",
                result.<Long>getAccumulatorResult(NUM_OUTPUTS),
                equalTo(result.getAccumulatorResult(NUM_INPUTS)));
        if (!topology.equals(Topology.CUSTOM_PARTITIONER)) {
            collector.checkThat(
                    "NUM_DUPLICATES",
                    result.<Long>getAccumulatorResult(NUM_DUPLICATES),
                    equalTo(0L));
        }
    }

    /**
     * A sink that checks if the members arrive in the expected order without any missing values.
     */
    protected static class VerifyingSink extends VerifyingSinkBase<VerifyingSink.State> {
        private boolean firstDuplicate = true;

        protected VerifyingSink(long minCheckpoints, long checkpointingInterval) {

Review comment (bug):
The second job expects failover once, and finishes after source generates all records. So removing `@ThrowableAnnotation(ThrowableType.NonRecoverableError)` for `TestException`.

Also, I introduced `ExpectedFinalJobStatus` in `UnalignedSettings` to check the final `JobStatus`.

Code after:
        return Arrays.stream(parameters)
                .map(params -> new Object[][] {ArrayUtils.insert(params.length, params)})
                .flatMap(Arrays::stream)
                .toArray(Object[][]::new);
    }

    public UnalignedCheckpointRescaleITCase(
            String desc,
            Topology topology,
            int oldParallelism,
            int newParallelism,
            long sourceSleepMs) {
        this.topology = topology;
        this.oldParallelism = oldParallelism;
        this.newParallelism = newParallelism;
        this.sourceSleepMs = sourceSleepMs;
    }

    /**
     * Tests unaligned checkpoint rescaling behavior.
     *
     * <p>Prescale phase: Job fails when completedCheckpoints >= minCheckpoints/2 via FailingMapper.
     * Generates checkpoint for rescale test.
     *
     * <p>Postscale phase: Job restores from checkpoint with different parallelism, failovers once,
     * and finishes after source generates all records.
     */
    @Test
    public void shouldRescaleUnalignedCheckpoint() throws Exception {
        final UnalignedSettings prescaleSettings =
                new UnalignedSettings(topology)
                        .setParallelism(oldParallelism)
                        .setExpectedFailures(1)
                        .setSourceSleepMs(sourceSleepMs)
                        .setExpectedFinalJobStatus(JobStatus.FAILED);
        prescaleSettings.setGenerateCheckpoint(true);
        final String checkpointDir = super.execute(prescaleSettings);
        assertThat(checkpointDir)
                .as("First job must generate a checkpoint for rescale test to be valid.")
                .isNotNull();
        // resume
        final UnalignedSettings postscaleSettings =
                new UnalignedSettings(topology)
                        .setParallelism(newParallelism)
                        .setExpectedFailures(1)
                        .setExpectedFinalJobStatus(JobStatus.FINISHED);
        postscaleSettings.setRestoreCheckpoint(checkpointDir);
        super.execute(postscaleSettings);
    }

    protected void checkCounters(JobExecutionResult result) {