@Test
public void testReportCompletedCheckpoint() throws Exception {

	TaskStateStats task1 = new TaskStateStats(new JobVertexID(), 3);
	TaskStateStats task2 = new TaskStateStats(new JobVertexID(), 4);
	HashMap<JobVertexID, TaskStateStats> taskStats = new HashMap<>();
	taskStats.put(task1.getJobVertexId(), task1);
	taskStats.put(task2.getJobVertexId(), task2);
	CheckpointStatsTracker.PendingCheckpointStatsCallback callback = mock(
		CheckpointStatsTracker.PendingCheckpointStatsCallback.class);
	PendingCheckpointStats pending = new PendingCheckpointStats(
		0,
		1,
		CheckpointProperties.forStandardCheckpoint(),
		task1.getNumberOfSubtasks() + task2.getNumberOfSubtasks(),
		taskStats,
		callback);
	// Report subtasks
	for (int i = 0; i < task1.getNumberOfSubtasks(); i++) {
		pending.reportSubtaskStats(task1.getJobVertexId(), createSubtaskStats(i));
	}
	for (int i = 0; i < task2.getNumberOfSubtasks(); i++) {
		pending.reportSubtaskStats(task2.getJobVertexId(), createSubtaskStats(i));
	}
	// Report completed
	String externalPath = "asdjkasdjkasd";
	CompletedCheckpointStats.DiscardCallback discardCallback = pending.reportCompletedCheckpoint(externalPath);
	ArgumentCaptor<CompletedCheckpointStats> args = ArgumentCaptor.forClass(CompletedCheckpointStats.class);
	verify(callback).reportCompletedCheckpoint(args.capture());
	CompletedCheckpointStats completed = args.getValue();
	assertNotNull(completed);
	assertEquals(CheckpointStatsStatus.COMPLETED, completed.getStatus());
	assertFalse(completed.isDiscarded());
	discardCallback.notifyDiscardedCheckpoint();
	assertTrue(completed.isDiscarded());
	assertEquals(externalPath, completed.getExternalPath());
	assertEquals(pending.getCheckpointId(), completed.getCheckpointId());
	assertEquals(pending.getNumberOfAcknowledgedSubtasks(), completed.getNumberOfAcknowledgedSubtasks());
	assertEquals(pending.getLatestAcknowledgedSubtaskStats(), completed.getLatestAcknowledgedSubtaskStats());
	assertEquals(pending.getLatestAckTimestamp(), completed.getLatestAckTimestamp());
	assertEquals(pending.getEndToEndDuration(), completed.getEndToEndDuration());
	assertEquals(pending.getStateSize(), completed.getStateSize());
	assertEquals(pending.getAlignmentBuffered(), completed.getAlignmentBuffered());
	assertEquals(task1, completed.getTaskStateStats(task1.getJobVertexId()));
	assertEquals(task2, completed.getTaskStateStats(task2.getJobVertexId()));
}