### Question:
StreamPartitionAssignor implements PartitionAssignor, Configurable { @Override public void onAssignment(Assignment assignment) { List<TopicPartition> partitions = new ArrayList<>(assignment.partitions()); Collections.sort(partitions, PARTITION_COMPARATOR); AssignmentInfo info = AssignmentInfo.decode(assignment.userData()); this.standbyTasks = info.standbyTasks; this.activeTasks = new HashMap<>(); if (partitions.size() != info.activeTasks.size()) { throw new TaskAssignmentException( String.format("stream-thread [%s] Number of assigned partitions %d is not equal to the number of active taskIds %d" + ", assignmentInfo=%s", streamThread.getName(), partitions.size(), info.activeTasks.size(), info.toString()) ); } for (int i = 0; i < partitions.size(); i++) { TopicPartition partition = partitions.get(i); TaskId id = info.activeTasks.get(i); Set<TopicPartition> assignedPartitions = activeTasks.get(id); if (assignedPartitions == null) { assignedPartitions = new HashSet<>(); activeTasks.put(id, assignedPartitions); } assignedPartitions.add(partition); } this.partitionsByHostState = info.partitionsByHost; final Collection<Set<TopicPartition>> values = partitionsByHostState.values(); final Map<TopicPartition, PartitionInfo> topicToPartitionInfo = new HashMap<>(); for (Set<TopicPartition> value : values) { for (TopicPartition topicPartition : value) { topicToPartitionInfo.put(topicPartition, new PartitionInfo(topicPartition.topic(), topicPartition.partition(), null, new Node[0], new Node[0])); } } metadataWithInternalTopics = Cluster.empty().withPartitions(topicToPartitionInfo); checkForNewTopicAssignments(assignment); }  @Override void configure(Map<String, ?> configs); @Override String name(); @Override Subscription subscription(Set<String> topics); @Override Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions); @Override void onAssignment(Assignment assignment); void close(); final static int NOT_AVAILABLE; }### Answer:
@Test public void testOnAssignment() throws Exception { TopicPartition t2p3 = new TopicPartition("topic2", 3); TopologyBuilder builder = new TopologyBuilder(); builder.addSource("source1", "topic1"); builder.addSource("source2", "topic2"); builder.addProcessor("processor", new MockProcessorSupplier(), "source1", "source2"); UUID uuid = UUID.randomUUID(); String client1 = "client1"; StreamThread thread = new StreamThread(builder, config, mockClientSupplier, "test", client1, uuid, new Metrics(), Time.SYSTEM, new StreamsMetadataState(builder, StreamsMetadataState.UNKNOWN_HOST), 0); partitionAssignor.configure(config.getConsumerConfigs(thread, "test", client1)); List<TaskId> activeTaskList = Utils.mkList(task0, task3); Map<TaskId, Set<TopicPartition>> activeTasks = new HashMap<>(); Map<TaskId, Set<TopicPartition>> standbyTasks = new HashMap<>(); activeTasks.put(task0, Utils.mkSet(t1p0)); activeTasks.put(task3, Utils.mkSet(t2p3)); standbyTasks.put(task1, Utils.mkSet(t1p0)); standbyTasks.put(task2, Utils.mkSet(t2p0)); AssignmentInfo info = new AssignmentInfo(activeTaskList, standbyTasks, new HashMap<HostInfo, Set<TopicPartition>>()); PartitionAssignor.Assignment assignment = new PartitionAssignor.Assignment(Utils.mkList(t1p0, t2p3), info.encode()); partitionAssignor.onAssignment(assignment); assertEquals(activeTasks, partitionAssignor.activeTasks()); assertEquals(standbyTasks, partitionAssignor.standbyTasks()); }