public void onCheckpointComplete(long checkpointId) {
        assignmentsByCheckpointId.entrySet().removeIf(entry -> entry.getKey() <= checkpointId);
    }
@Test
    void testOnCheckpointComplete() throws Exception {
        final long checkpointId1 = 100L;
        final long checkpointId2 = 101L;
        SplitAssignmentTracker<MockSourceSplit> tracker = new SplitAssignmentTracker<>();

        // Assign some splits to subtask 0 and 1.
        tracker.recordSplitAssignment(getSplitsAssignment(2, 0));

        // Take the first snapshot.
        tracker.onCheckpoint(checkpointId1);
        verifyAssignment(
                Arrays.asList("0"), tracker.assignmentsByCheckpointId(checkpointId1).get(0));
        verifyAssignment(
                Arrays.asList("1", "2"), tracker.assignmentsByCheckpointId(checkpointId1).get(1));

        // Assign additional splits to subtask 0 and 1.
        tracker.recordSplitAssignment(getSplitsAssignment(2, 3));

        // Take the second snapshot.
        tracker.onCheckpoint(checkpointId2);
        verifyAssignment(
                Arrays.asList("0"), tracker.assignmentsByCheckpointId(checkpointId1).get(0));
        verifyAssignment(
                Arrays.asList("1", "2"), tracker.assignmentsByCheckpointId(checkpointId1).get(1));
        verifyAssignment(
                Arrays.asList("3"), tracker.assignmentsByCheckpointId(checkpointId2).get(0));
        verifyAssignment(
                Arrays.asList("4", "5"), tracker.assignmentsByCheckpointId(checkpointId2).get(1));

        // Complete the first checkpoint.
        tracker.onCheckpointComplete(checkpointId1);
        assertThat(tracker.assignmentsByCheckpointId(checkpointId1)).isNull();
        verifyAssignment(
                Arrays.asList("3"), tracker.assignmentsByCheckpointId(checkpointId2).get(0));
        verifyAssignment(
                Arrays.asList("4", "5"), tracker.assignmentsByCheckpointId(checkpointId2).get(1));
    }