@Test
public void testOnPartitionStateUpdate() throws Exception {

	IntermediateDataSetID resultId = new IntermediateDataSetID();
	ResultPartitionID partitionId = new ResultPartitionID();
	SingleInputGate inputGate = mock(SingleInputGate.class);
	when(inputGate.getConsumedResultId()).thenReturn(resultId);
	final Task task = createTask(InvokableBlockingInInvoke.class);
	// Set the mock input gate
	setInputGate(task, inputGate);
	// Expected task state for each partition state
	final Map<ExecutionState, ExecutionState> expected = Maps
			.newHashMapWithExpectedSize(ExecutionState.values().length);
	// Fail the task for unexpected states
	for (ExecutionState state : ExecutionState.values()) {
		expected.put(state, ExecutionState.FAILED);
	}
	expected.put(ExecutionState.RUNNING, ExecutionState.RUNNING);
	expected.put(ExecutionState.CANCELED, ExecutionState.CANCELING);
	expected.put(ExecutionState.CANCELING, ExecutionState.CANCELING);
	expected.put(ExecutionState.FAILED, ExecutionState.CANCELING);
	for (ExecutionState state : ExecutionState.values()) {
		setState(task, ExecutionState.RUNNING);
		task.onPartitionStateUpdate(resultId, partitionId.getPartitionId(), state);
		ExecutionState newTaskState = task.getExecutionState();
		assertEquals(expected.get(state), newTaskState);
	}
	verify(inputGate, times(1)).retriggerPartitionRequest(eq(partitionId.getPartitionId()));
}