@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 producer state
	final Map<ExecutionState, ExecutionState> expected = new HashMap<>(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.SCHEDULED, ExecutionState.RUNNING);
	expected.put(ExecutionState.DEPLOYING, ExecutionState.RUNNING);
	expected.put(ExecutionState.FINISHED, 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, state);
		ExecutionState newTaskState = task.getExecutionState();
		assertEquals(expected.get(state), newTaskState);
	}
	verify(inputGate, times(4)).retriggerPartitionRequest(eq(partitionId.getPartitionId()));
}