JobModelUtil { public static Map<TaskName, Set<SystemStreamPartition>> getTaskToSystemStreamPartitions(JobModel jobModel) { Preconditions.checkArgument(jobModel != null, "JobModel cannot be null"); Map<String, ContainerModel> containers = jobModel.getContainers(); HashMap<TaskName, Set<SystemStreamPartition>> taskToSSPs = new HashMap<>(); for (ContainerModel containerModel : containers.values()) { for (TaskName taskName : containerModel.getTasks().keySet()) { TaskModel taskModel = containerModel.getTasks().get(taskName); if (taskModel.getTaskMode() != TaskMode.Active) { continue; } taskToSSPs.putIfAbsent(taskName, new HashSet<>()); taskToSSPs.get(taskName).addAll(taskModel.getSystemStreamPartitions()); } } return taskToSSPs; } private  JobModelUtil(); static Map<TaskName, Set<SystemStreamPartition>> getTaskToSystemStreamPartitions(JobModel jobModel); static void writeJobModel(JobModel jobModel, String jobModelVersion, MetadataStore metadataStore); static JobModel readJobModel(String jobModelVersion, MetadataStore metadataStore); static Set<SystemStream> getSystemStreams(JobModel jobModel);  }
@Test(expected = IllegalArgumentException.class) public void testTaskToSystemStreamPartitionsWithNullJobModel() { JobModelUtil.getTaskToSystemStreamPartitions(null); }
@Test public void testTaskToSystemStreamPartitions() { MapConfig config = new MapConfig(ImmutableMap.of("a", "b")); Set<SystemStreamPartition> ssps1 = ImmutableSet.of( new SystemStreamPartition("system1", "stream1_1", new Partition(0)), new SystemStreamPartition("system1", "stream1_2", new Partition(0))); Set<SystemStreamPartition> ssps2 = ImmutableSet.of( new SystemStreamPartition("system1", "stream2_1", new Partition(1)), new SystemStreamPartition("system1", "stream2_2", new Partition(1))); Set<SystemStreamPartition> ssps3 = ImmutableSet.of( new SystemStreamPartition("system1", "stream3_1", new Partition(2)), new SystemStreamPartition("system1", "stream3_2", new Partition(2))); Set<SystemStreamPartition> ssps4 = ImmutableSet.of( new SystemStreamPartition("system1", "stream4_1", new Partition(3)), new SystemStreamPartition("system1", "stream4_2", new Partition(3))); TaskName task1 = new TaskName("task1"); TaskName task2 = new TaskName("task2"); Map<TaskName, TaskModel> tasksForContainer1 = ImmutableMap.of( task1, new TaskModel(task1, ssps1, new Partition(0)), task2, new TaskModel(task2, ssps2, new Partition(1))); TaskName task3 = new TaskName("task3"); Map<TaskName, TaskModel> tasksForContainer2 = ImmutableMap.of( task3, new TaskModel(task3, ssps3, new Partition(2)), task1, new TaskModel(task1, ssps4, new Partition(3))); ContainerModel containerModel1 = new ContainerModel("0", tasksForContainer1); ContainerModel containerModel2 = new ContainerModel("1", tasksForContainer2); Map<String, ContainerModel> containers = ImmutableMap.of("0", containerModel1, "1", containerModel2); JobModel jobModel = new JobModel(config, containers); assertEquals(ssps1.size() + ssps4.size(), JobModelUtil.getTaskToSystemStreamPartitions(jobModel).get(task1).size()); assertTrue(JobModelUtil.getTaskToSystemStreamPartitions(jobModel).get(task1).containsAll(ssps1)); assertTrue(JobModelUtil.getTaskToSystemStreamPartitions(jobModel).get(task1).containsAll(ssps4)); assertEquals(ssps2.size(), JobModelUtil.getTaskToSystemStreamPartitions(jobModel).get(task2).size()); assertTrue(JobModelUtil.getTaskToSystemStreamPartitions(jobModel).get(task2).containsAll(ssps2)); assertEquals(ssps3.size(), JobModelUtil.getTaskToSystemStreamPartitions(jobModel).get(task3).size()); assertTrue(JobModelUtil.getTaskToSystemStreamPartitions(jobModel).get(task3).containsAll(ssps3)); }