Code review comment: So far, we've only seen the issue in `#testCheckpointRescalingInKeyedState`. We don't need the two checkpoints here, actually, because we're not relying on elements in the test. We could keep the tests functionality if we make the `waitForOneMoreCheckpoint` configurable as suggested one of my previous comments.

Code under review:
public void testCheckpointRescalingNonPartitionedStateCausesException() throws Exception {
        final int parallelism = totalSlots / 2;
        final int parallelism2 = totalSlots;
        final int maxParallelism = 13;

        ClusterClient<?> client = cluster.getClusterClient();

        try {
            JobGraph jobGraph =
                    createJobGraphWithOperatorState(
                            parallelism, maxParallelism, OperatorCheckpointMethod.NON_PARTITIONED);
            
            StateSourceBase.canFinishLatch = new CountDownLatch(1);

            final JobID jobID = jobGraph.getJobID();

            client.submitJob(jobGraph).get();

            
            waitForAllTaskRunning(cluster.getMiniCluster(), jobGraph.getJobID(), false);
            
            StateSourceBase.workStartedLatch.await();

            waitForNewCheckpoint(jobID, cluster.getMiniCluster());

            JobResourceRequirements.Builder builder = JobResourceRequirements.newBuilder();
            for (JobVertex vertex : jobGraph.getVertices()) {
                builder.setParallelismForJobVertex(vertex.getID(), parallelism2, parallelism2);
            }

            restClusterClient.updateJobResourceRequirements(jobID, builder.build()).join();

            waitForRunningTasks(restClusterClient, jobID, parallelism2);
            waitForAvailableSlots(restClusterClient, totalSlots - parallelism2);

            StateSourceBase.canFinishLatch.countDown();

            client.requestJobResult(jobID).get();
        } catch (JobExecutionException exception) {
            if (!(exception.getCause() instanceof IllegalStateException)) {
                throw exception;
            }
        }
    }

Line commented on: builder.setParallelismForJobVertex(vertex.getID(), parallelism2, parallelism2);

Revised code:
public void testCheckpointRescalingNonPartitionedStateCausesException() throws Exception {
        final int parallelism = totalSlots / 2;
        final int parallelism2 = totalSlots;
        final int maxParallelism = 13;

        ClusterClient<?> client = cluster.getClusterClient();

        try {
            JobGraph jobGraph =
                    createJobGraphWithOperatorState(
                            parallelism, maxParallelism, OperatorCheckpointMethod.NON_PARTITIONED);
            
            StateSourceBase.canFinishLatch = new CountDownLatch(1);

            final JobID jobID = jobGraph.getJobID();

            client.submitJob(jobGraph).get();

            
            waitForAllTaskRunning(cluster.getMiniCluster(), jobGraph.getJobID(), false);
            
            StateSourceBase.workStartedLatch.await();

            waitForNewCheckpoint(jobID, cluster.getMiniCluster());

            JobResourceRequirements.Builder builder = JobResourceRequirements.newBuilder();
            for (JobVertex vertex : jobGraph.getVertices()) {
                builder.setParallelismForJobVertex(vertex.getID(), parallelism2, parallelism2);
            }

            restClusterClient.updateJobResourceRequirements(jobID, builder.build()).join();

            waitForRunningTasks(restClusterClient, jobID, parallelism2);
            waitForAvailableSlots(restClusterClient, totalSlots - parallelism2);

            StateSourceBase.canFinishLatch.countDown();

            client.requestJobResult(jobID).get();
        } catch (JobExecutionException exception) {
            if (!(exception.getCause() instanceof IllegalStateException)) {
                throw exception;
            }
        }
    }