RelocationUtil { public static List<String> buildTasksFromNodesAndJobsFilter(Map<String, Node> nodesById, Set<String> jobIds, ReadOnlyJobOperations jobOperations) { Map<String, Task> tasksById = buildTaskByIdMap(jobOperations); Set<String> taskIdsOnNodes = buildTasksToInstanceMap(nodesById, tasksById).keySet(); return taskIdsOnNodes.stream().filter(taskId -> { if (tasksById.containsKey(taskId)) { Task task = tasksById.get(taskId); return jobIds.contains(task.getJobId()); } return false; }).collect(Collectors.toList()); }  static Map<String, Task> buildTaskByIdMap(ReadOnlyJobOperations jobOperations); static Map<String, Node> buildTasksToInstanceMap(Map<String, Node> nodesById, Map<String, Task> taskByIdMap); static Map<String, Node> buildTasksToInstanceMap(Map<String, Node> nodesById,
                                                            ReadOnlyJobOperations jobOperations); static List<String> buildTasksFromNodesAndJobsFilter(Map<String, Node> nodesById, Set<String> jobIds,
                                                                ReadOnlyJobOperations jobOperations); static List<Task> findTasksOnInstance(Node instance, Collection<Task> tasks); static boolean isAssignedToAgent(Task task); static boolean isOnInstance(Node instance, Task task); static String doFormat(TaskRelocationPlan plan);  }
@Test public void buildTasksFromNodesAndJobsFilter() { String node1 = "node1"; String node2 = "node2"; String node3 = "node3"; Job<BatchJobExt> job1 = JobGenerator.oneBatchJob(); Job<BatchJobExt> job2 = JobGenerator.oneBatchJob(); Job<BatchJobExt> job3 = JobGenerator.oneBatchJob(); BatchJobTask task1 = JobGenerator.batchTasks(job1).getValue().toBuilder() .addToTaskContext(TaskAttributes.TASK_ATTRIBUTES_AGENT_INSTANCE_ID, node1).build(); BatchJobTask task2 = JobGenerator.batchTasks(job2).getValue().toBuilder() .addToTaskContext(TaskAttributes.TASK_ATTRIBUTES_AGENT_INSTANCE_ID, node2).build(); BatchJobTask task3 = JobGenerator.batchTasks(job3).getValue().toBuilder() .addToTaskContext(TaskAttributes.TASK_ATTRIBUTES_AGENT_INSTANCE_ID, node3).build(); ReadOnlyJobOperations jobOperations = mock(ReadOnlyJobOperations.class); when(jobOperations.getJobs()).thenReturn(Arrays.asList(job1, job2, job3)); when(jobOperations.getTasks(job1.getId())).thenReturn(Collections.singletonList(task1)); when(jobOperations.getTasks(job2.getId())).thenReturn(Collections.singletonList(task2)); when(jobOperations.getTasks(job3.getId())).thenReturn(Collections.singletonList(task3)); Map<String, Node> nodes = new HashMap<>(3); nodes.put(node1, buildNode(node1)); nodes.put(node2, buildNode(node2)); nodes.put(node3, buildNode(node3)); Set<String> jobIds = new HashSet<>(2); jobIds.addAll(Arrays.asList(job1.getId(), job3.getId())); List<String> taskIdsOnBadNodes = RelocationUtil.buildTasksFromNodesAndJobsFilter(nodes, jobIds, jobOperations); assertThat(taskIdsOnBadNodes.size()).isEqualTo(2); }