public void onCheckpoint(long checkpointId) throws Exception {
        // Include the uncheckpointed assignments to the snapshot.
        assignmentsByCheckpointId.put(checkpointId, uncheckpointedAssignments);
        uncheckpointedAssignments = new HashMap<>();
    }
@Test
    void testOnCheckpoint() throws Exception {
        final long checkpointId = 123L;
        SplitAssignmentTracker<MockSourceSplit> tracker = new SplitAssignmentTracker<>();
        tracker.recordSplitAssignment(getSplitsAssignment(3, 0));

        // Serialize
        tracker.onCheckpoint(checkpointId);

        // Verify the uncheckpointed assignments.
        assertThat(tracker.uncheckpointedAssignments()).isEmpty();

        // verify assignments put into the checkpoints.
        Map<Long, Map<Integer, LinkedHashSet<MockSourceSplit>>> assignmentsByCheckpoints =
                tracker.assignmentsByCheckpointId();
        assertThat(assignmentsByCheckpoints.size()).isOne();

        Map<Integer, LinkedHashSet<MockSourceSplit>> assignmentForCheckpoint =
                assignmentsByCheckpoints.get(checkpointId);
        assertThat(assignmentForCheckpoint).isNotNull();

        verifyAssignment(Arrays.asList("0"), assignmentForCheckpoint.get(0));
        verifyAssignment(Arrays.asList("1", "2"), assignmentForCheckpoint.get(1));
        verifyAssignment(Arrays.asList("3", "4", "5"), assignmentForCheckpoint.get(2));
    }