@Test
public void testReportSubtaskStats() throws Exception {

	long checkpointId = Integer.MAX_VALUE + 1222L;
	long triggerTimestamp = Integer.MAX_VALUE - 1239L;
	CheckpointProperties props = CheckpointProperties.forStandardCheckpoint();
	TaskStateStats task1 = new TaskStateStats(new JobVertexID(), 3);
	TaskStateStats task2 = new TaskStateStats(new JobVertexID(), 4);
	int totalSubtaskCount = task1.getNumberOfSubtasks() + task2.getNumberOfSubtasks();
	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(
			checkpointId,
			triggerTimestamp,
			props,
			totalSubtaskCount,
			taskStats,
			callback);
	// Check initial state
	assertEquals(checkpointId, pending.getCheckpointId());
	assertEquals(triggerTimestamp, pending.getTriggerTimestamp());
	assertEquals(props, pending.getProperties());
	assertEquals(CheckpointStatsStatus.IN_PROGRESS, pending.getStatus());
	assertEquals(0, pending.getNumberOfAcknowledgedSubtasks());
	assertEquals(0, pending.getStateSize());
	assertEquals(totalSubtaskCount, pending.getNumberOfSubtasks());
	assertNull(pending.getLatestAcknowledgedSubtaskStats());
	assertEquals(-1, pending.getLatestAckTimestamp());
	assertEquals(-1, pending.getEndToEndDuration());
	assertEquals(task1, pending.getTaskStateStats(task1.getJobVertexId()));
	assertEquals(task2, pending.getTaskStateStats(task2.getJobVertexId()));
	assertNull(pending.getTaskStateStats(new JobVertexID()));
	// Report subtasks and check getters
	assertFalse(pending.reportSubtaskStats(new JobVertexID(), createSubtaskStats(0)));
	long stateSize = 0;
	long alignmentBuffered = 0;
	// Report 1st task
	for (int i = 0; i < task1.getNumberOfSubtasks(); i++) {
		SubtaskStateStats subtask = createSubtaskStats(i);
		stateSize += subtask.getStateSize();
		alignmentBuffered += subtask.getAlignmentBuffered();
		pending.reportSubtaskStats(task1.getJobVertexId(), subtask);
		assertEquals(subtask, pending.getLatestAcknowledgedSubtaskStats());
		assertEquals(subtask.getAckTimestamp(), pending.getLatestAckTimestamp());
		assertEquals(subtask.getAckTimestamp() - triggerTimestamp, pending.getEndToEndDuration());
		assertEquals(stateSize, pending.getStateSize());
		assertEquals(alignmentBuffered, pending.getAlignmentBuffered());
	}
	// Don't allow overwrite
	assertFalse(pending.reportSubtaskStats(task1.getJobVertexId(), task1.getSubtaskStats()[0]));
	// Report 2nd task
	for (int i = 0; i < task2.getNumberOfSubtasks(); i++) {
		SubtaskStateStats subtask = createSubtaskStats(i);
		stateSize += subtask.getStateSize();
		alignmentBuffered += subtask.getAlignmentBuffered();
		pending.reportSubtaskStats(task2.getJobVertexId(), subtask);
		assertEquals(subtask, pending.getLatestAcknowledgedSubtaskStats());
		assertEquals(subtask.getAckTimestamp(), pending.getLatestAckTimestamp());
		assertEquals(subtask.getAckTimestamp() - triggerTimestamp, pending.getEndToEndDuration());
		assertEquals(stateSize, pending.getStateSize());
		assertEquals(alignmentBuffered, pending.getAlignmentBuffered());
	}
	assertEquals(task1.getNumberOfSubtasks(), task1.getNumberOfAcknowledgedSubtasks());
	assertEquals(task2.getNumberOfSubtasks(), task2.getNumberOfAcknowledgedSubtasks());
}