ClusterState { public void removeTask(TaskInfo taskInfo) throws InvalidParameterException { List<TaskInfo> taskList = getTaskList(); LOGGER.debug("Removing TaskInfo from cluster for task: " + taskInfo.getTaskId().getValue()); if (!taskList.remove(taskInfo)) { throw new InvalidParameterException("TaskInfo does not exist in list: " + taskInfo.getTaskId().getValue()); } getStatus(taskInfo).destroy(); setTaskInfoList(taskList); } ClusterState(@NotNull SerializableState zooKeeperStateDriver, @NotNull FrameworkState frameworkState); List<TaskInfo> getTaskList(); Integer getElasticNodeId(); Map<String, Task> getGuiTaskList(); ESTaskStatus getStatus(TaskID taskID); void addTask(ESTaskStatus esTask); void addTask(TaskInfo taskInfo); void removeTask(TaskInfo taskInfo); Boolean exists(TaskID taskId); TaskInfo getTask(TaskID taskID); TaskInfo getTask(Protos.ExecutorID executorID); void update(Protos.TaskStatus status); boolean taskInError(Protos.TaskStatus status); void destroy(); static final Logger LOGGER; static final String STATE_LIST; }
@Test public void shouldDeleteTask() throws IOException { ArrayList<Protos.TaskInfo> mock = Mockito.spy(new ArrayList<>()); Protos.TaskInfo defaultTaskInfo = ProtoTestUtil.getDefaultTaskInfo(); mock.add(defaultTaskInfo); when(state.get(anyString())).thenReturn(mock).thenReturn(ProtoTestUtil.getDefaultTaskStatus(Protos.TaskState.TASK_FINISHED)); clusterState.removeTask(defaultTaskInfo); verify(state, times(1)).set(anyString(), any()); verify(mock, times(1)).remove(eq(defaultTaskInfo)); }