content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.streams.processor.internals;
  
  import org.apache.kafka.clients.admin.Admin;
  import org.apache.kafka.clients.consumer.Consumer;
  import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
  import org.apache.kafka.common.Cluster;
  import org.apache.kafka.common.Node;
  import org.apache.kafka.common.PartitionInfo;
  import org.apache.kafka.common.TopicPartition;
  import org.apache.kafka.common.utils.MockTime;
  import org.apache.kafka.streams.StreamsConfig;
  import org.apache.kafka.streams.processor.TaskId;
  import org.apache.kafka.streams.processor.assignment.ProcessId;
  import org.apache.kafka.streams.processor.internals.assignment.AssignmentInfo;
  import org.apache.kafka.streams.processor.internals.assignment.ReferenceContainer;
  import org.apache.kafka.streams.processor.internals.assignment.SubscriptionInfo;
  import org.apache.kafka.test.MockApiProcessorSupplier;
  import org.apache.kafka.test.MockClientSupplier;
  import org.apache.kafka.test.MockInternalTopicManager;
  import org.apache.kafka.test.MockKeyValueStoreBuilder;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.extension.ExtendWith;
  import org.mockito.junit.jupiter.MockitoExtension;
  import org.mockito.junit.jupiter.MockitoSettings;
  import org.mockito.quality.Strictness;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.stream.Collectors;
  
  import static java.util.Arrays.asList;
  import static java.util.Collections.emptySet;
  import static java.util.Collections.singletonList;
  import static org.apache.kafka.common.utils.Utils.mkEntry;
  import static org.apache.kafka.common.utils.Utils.mkMap;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_TASKS;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_1;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_2;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_3;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_4;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_5;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_6;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_7;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_8;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PID_9;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.createMockAdminClientForAssignor;
  import static org.apache.kafka.streams.processor.internals.assignment.StreamsAssignmentProtocolVersions.LATEST_SUPPORTED_VERSION;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  
  @ExtendWith(MockitoExtension.class)
  @MockitoSettings(strictness = Strictness.STRICT_STUBS)
  public class RackAwarenessStreamsPartitionAssignorTest {
  
      private final List<PartitionInfo> infos = asList(
          new PartitionInfo("topic0", 0, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic0", 1, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic0", 2, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic1", 0, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic1", 1, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic1", 2, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic2", 0, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic2", 1, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic2", 2, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic3", 0, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic3", 1, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic3", 2, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic4", 0, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic4", 1, Node.noNode(), new Node[0], new Node[0]),
          new PartitionInfo("topic4", 2, Node.noNode(), new Node[0], new Node[0])
      );
  
      final String consumer1 = "consumer1";
      final String consumer2 = "consumer2";
      final String consumer3 = "consumer3";
      final String consumer4 = "consumer4";
      final String consumer5 = "consumer5";
      final String consumer6 = "consumer6";
      final String consumer7 = "consumer7";
      final String consumer8 = "consumer8";
      final String consumer9 = "consumer9";
  
  
      private final Cluster metadata = new Cluster(
              "cluster",
              singletonList(Node.noNode()),
              infos,
              emptySet(),
              emptySet());
  
      private static final List<String> ALL_TAG_KEYS = new ArrayList<>();
      static {
          for (int i = 0; i < StreamsConfig.MAX_RACK_AWARE_ASSIGNMENT_TAG_LIST_SIZE; i++) {
              ALL_TAG_KEYS.add("key-" + i);
          }
      }
  
      private final StreamsPartitionAssignor partitionAssignor = new StreamsPartitionAssignor();
      private final MockClientSupplier mockClientSupplier = new MockClientSupplier();
      private static final String USER_END_POINT = "localhost:8080";
      private static final String APPLICATION_ID = "stream-partition-assignor-test";
  
      private TaskManager taskManager;
      private Admin adminClient;
      private StreamsConfig streamsConfig = new StreamsConfig(configProps());
      private final InternalTopologyBuilder builder = new InternalTopologyBuilder();
      private TopologyMetadata topologyMetadata = new TopologyMetadata(builder, streamsConfig);
      private final StreamsMetadataState streamsMetadataState = mock(StreamsMetadataState.class);
      private final Map<String, ConsumerPartitionAssignor.Subscription> subscriptions = new HashMap<>();
      private final MockTime time = new MockTime();
  
      @SuppressWarnings("unchecked")
      private Map<String, Object> configProps() {
          final Map<String, Object> configurationMap = new HashMap<>();
          configurationMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
          configurationMap.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, USER_END_POINT);
          final ReferenceContainer referenceContainer = new ReferenceContainer();
          referenceContainer.mainConsumer = (Consumer<byte[], byte[]>) mock(Consumer.class);
          referenceContainer.adminClient = adminClient;
          referenceContainer.taskManager = taskManager;
          referenceContainer.streamsMetadataState = streamsMetadataState;
          referenceContainer.time = time;
          configurationMap.put(StreamsConfig.InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
          configurationMap.put(StreamsConfig.RACK_AWARE_ASSIGNMENT_TAGS_CONFIG, String.join(",", ALL_TAG_KEYS));
          ALL_TAG_KEYS.forEach(key -> configurationMap.put(StreamsConfig.clientTagPrefix(key), "dummy"));
          return configurationMap;
      }
  
      // Make sure to complete setting up any mocks (such as TaskManager or AdminClient) before configuring the assignor
      private void configurePartitionAssignorWith(final Map<String, Object> props) {
          final Map<String, Object> configMap = configProps();
          configMap.putAll(props);
  
          streamsConfig = new StreamsConfig(configMap);
          topologyMetadata = new TopologyMetadata(builder, streamsConfig);
          partitionAssignor.configure(configMap);
  
          overwriteInternalTopicManagerWithMock();
      }
  
      // Useful for tests that don't care about the task offset sums
      private void createMockTaskManager() {
          taskManager = mock(TaskManager.class);
          when(taskManager.topologyMetadata()).thenReturn(topologyMetadata);
          topologyMetadata.buildAndRewriteTopology();
      }
  
      private void overwriteInternalTopicManagerWithMock() {
          final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
                  time,
                  streamsConfig,
                  mockClientSupplier.restoreConsumer,
                  false
          );
          partitionAssignor.setInternalTopicManager(mockInternalTopicManager);
      }
  
      @Test
      public void shouldDistributeWithMaximumNumberOfClientTags() {
          setupTopology(3, 2);
  
          createMockTaskManager();
          adminClient = createMockAdminClientForAssignor(getTopicPartitionOffsetsMap(
                  Arrays.asList(APPLICATION_ID + "-store2-changelog", APPLICATION_ID + "-store3-changelog", APPLICATION_ID + "-store4-changelog"),
                  Arrays.asList(3, 3, 3)),
              true);
          configurePartitionAssignorWith(Collections.singletonMap(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1));
  
          final Map<String, String> clientTags1 = new HashMap<>();
          final Map<String, String> clientTags2 = new HashMap<>();
  
          for (int i = 0; i < ALL_TAG_KEYS.size(); i++) {
              final String key = ALL_TAG_KEYS.get(i);
              clientTags1.put(key, "value-1-" + i);
              clientTags2.put(key, "value-2-" + i);
          }
  
          final Map<String, Map<String, String>> hostTags = new HashMap<>();
          subscriptions.put(consumer1, getSubscription(PID_1, EMPTY_TASKS, clientTags1));
          hostTags.put(consumer1, clientTags1);
          subscriptions.put(consumer2, getSubscription(PID_2, EMPTY_TASKS, clientTags1));
          hostTags.put(consumer2, clientTags1);
          subscriptions.put(consumer3, getSubscription(PID_3, EMPTY_TASKS, clientTags2));
          hostTags.put(consumer3, clientTags2);
  
          Map<String, ConsumerPartitionAssignor.Assignment> assignments = partitionAssignor
              .assign(metadata, new ConsumerPartitionAssignor.GroupSubscription(subscriptions))
              .groupAssignment();
  
          verifyIdealTaskDistributionReached(getClientTagDistributions(assignments, hostTags), ALL_TAG_KEYS);
  
          // kill the first consumer and rebalance, should still achieve ideal distribution
          subscriptions.clear();
          subscriptions.put(consumer2, getSubscription(PID_2, AssignmentInfo.decode(assignments.get(consumer2).userData()).activeTasks(), clientTags1));
          subscriptions.put(consumer3, getSubscription(PID_3, AssignmentInfo.decode(assignments.get(consumer3).userData()).activeTasks(), clientTags2));
  
          assignments = partitionAssignor.assign(metadata, new ConsumerPartitionAssignor.GroupSubscription(subscriptions))
              .groupAssignment();
  
          verifyIdealTaskDistributionReached(getClientTagDistributions(assignments, hostTags), ALL_TAG_KEYS);
      }
  
      @Test
      public void shouldDistributeOnDistinguishingTagSubset() {
          setupTopology(3, 0);
  
          createMockTaskManager();
          adminClient = createMockAdminClientForAssignor(getTopicPartitionOffsetsMap(Arrays.asList(APPLICATION_ID + "-store0-changelog", APPLICATION_ID + "-store1-changelog", APPLICATION_ID + "-store2-changelog"),
                  Arrays.asList(3, 3, 3)),
              true);
          configurePartitionAssignorWith(Collections.singletonMap(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1));
  
          // use the same tag value for key1, and different value for key2
          // then we verify that for key2 we still achieve ideal distribution
          final Map<String, String> clientTags1 = new HashMap<>();
          final Map<String, String> clientTags2 = new HashMap<>();
          clientTags1.put(ALL_TAG_KEYS.get(0), "value-1-all");
          clientTags2.put(ALL_TAG_KEYS.get(0), "value-2-all");
          clientTags1.put(ALL_TAG_KEYS.get(1), "value-1-1");
          clientTags2.put(ALL_TAG_KEYS.get(1), "value-2-2");
  
          final String consumer1 = "consumer1";
          final String consumer2 = "consumer2";
          final String consumer3 = "consumer3";
          final String consumer4 = "consumer4";
          final String consumer5 = "consumer5";
          final String consumer6 = "consumer6";
  
          final Map<String, Map<String, String>> hostTags = new HashMap<>();
          subscriptions.put(consumer1, getSubscription(PID_1, EMPTY_TASKS, clientTags1));
          hostTags.put(consumer1, clientTags1);
          subscriptions.put(consumer2, getSubscription(PID_2, EMPTY_TASKS, clientTags1));
          hostTags.put(consumer2, clientTags1);
          subscriptions.put(consumer3, getSubscription(PID_3, EMPTY_TASKS, clientTags1));
          hostTags.put(consumer3, clientTags1);
          subscriptions.put(consumer4, getSubscription(PID_4, EMPTY_TASKS, clientTags2));
          hostTags.put(consumer4, clientTags2);
          subscriptions.put(consumer5, getSubscription(PID_5, EMPTY_TASKS, clientTags2));
          hostTags.put(consumer5, clientTags2);
          subscriptions.put(consumer6, getSubscription(PID_6, EMPTY_TASKS, clientTags2));
          hostTags.put(consumer6, clientTags2);
  
          final Map<String, ConsumerPartitionAssignor.Assignment> assignments = partitionAssignor
              .assign(metadata, new ConsumerPartitionAssignor.GroupSubscription(subscriptions))
              .groupAssignment();
  
          verifyIdealTaskDistributionReached(getClientTagDistributions(assignments, hostTags), Collections.singletonList(ALL_TAG_KEYS.get(1)));
      }
  
      @Test
      public void shouldDistributeWithMultipleStandbys() {
          setupTopology(3, 0);
  
          createMockTaskManager();
          adminClient = createMockAdminClientForAssignor(getTopicPartitionOffsetsMap(
                  Arrays.asList(APPLICATION_ID + "-store0-changelog", APPLICATION_ID + "-store1-changelog", APPLICATION_ID + "-store2-changelog"),
                  Arrays.asList(3, 3, 3)),
              true);
          configurePartitionAssignorWith(Collections.singletonMap(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 2));
  
          final Map<String, String> clientTags1 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-1"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-1"));
          final Map<String, String> clientTags2 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-1"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-2"));
          final Map<String, String> clientTags3 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-1"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-3"));
          final Map<String, String> clientTags4 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-2"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-1"));
          final Map<String, String> clientTags5 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-2"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-2"));
          final Map<String, String> clientTags6 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-2"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-3"));
          final Map<String, String> clientTags7 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-3"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-1"));
          final Map<String, String> clientTags8 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-3"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-2"));
          final Map<String, String> clientTags9 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-3"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-3"));
  
          final Map<String, Map<String, String>> hostTags = new HashMap<>();
          subscriptions.put(consumer1, getSubscription(PID_1, EMPTY_TASKS, clientTags1));
          hostTags.put(consumer1, clientTags1);
          subscriptions.put(consumer2, getSubscription(PID_2, EMPTY_TASKS, clientTags2));
          hostTags.put(consumer2, clientTags2);
          subscriptions.put(consumer3, getSubscription(PID_3, EMPTY_TASKS, clientTags3));
          hostTags.put(consumer3, clientTags3);
          subscriptions.put(consumer4, getSubscription(PID_4, EMPTY_TASKS, clientTags4));
          hostTags.put(consumer4, clientTags4);
          subscriptions.put(consumer5, getSubscription(PID_5, EMPTY_TASKS, clientTags5));
          hostTags.put(consumer5, clientTags5);
          subscriptions.put(consumer6, getSubscription(PID_6, EMPTY_TASKS, clientTags6));
          hostTags.put(consumer6, clientTags6);
          subscriptions.put(consumer7, getSubscription(PID_7, EMPTY_TASKS, clientTags7));
          hostTags.put(consumer7, clientTags7);
          subscriptions.put(consumer8, getSubscription(PID_8, EMPTY_TASKS, clientTags8));
          hostTags.put(consumer8, clientTags8);
          subscriptions.put(consumer9, getSubscription(PID_9, EMPTY_TASKS, clientTags9));
          hostTags.put(consumer9, clientTags9);
  
          final Map<String, ConsumerPartitionAssignor.Assignment> assignments = partitionAssignor
              .assign(metadata, new ConsumerPartitionAssignor.GroupSubscription(subscriptions))
              .groupAssignment();
  
          verifyIdealTaskDistributionReached(getClientTagDistributions(assignments, hostTags), Arrays.asList(ALL_TAG_KEYS.get(0), ALL_TAG_KEYS.get(1)));
      }
  
      @Test
      public void shouldDistributePartiallyWhenDoNotHaveEnoughClients() {
          setupTopology(3, 0);
  
          createMockTaskManager();
          adminClient = createMockAdminClientForAssignor(getTopicPartitionOffsetsMap(
                  Arrays.asList(APPLICATION_ID + "-store0-changelog", APPLICATION_ID + "-store1-changelog", APPLICATION_ID + "-store2-changelog"),
                  Arrays.asList(3, 3, 3)),
              true);
          configurePartitionAssignorWith(Collections.singletonMap(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 2));
  
          final Map<String, String> clientTags1 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-1"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-1"));
          final Map<String, String> clientTags2 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-1"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-2"));
          final Map<String, String> clientTags3 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-1"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-3"));
          final Map<String, String> clientTags4 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-2"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-1"));
          final Map<String, String> clientTags5 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-2"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-2"));
          final Map<String, String> clientTags6 = mkMap(
              mkEntry(ALL_TAG_KEYS.get(0), "value-0-2"),
              mkEntry(ALL_TAG_KEYS.get(1), "value-1-3"));
  
          final Map<String, Map<String, String>> hostTags = new HashMap<>();
          subscriptions.put(consumer1, getSubscription(PID_1, EMPTY_TASKS, clientTags1));
          hostTags.put(consumer1, clientTags1);
          subscriptions.put(consumer2, getSubscription(PID_2, EMPTY_TASKS, clientTags2));
          hostTags.put(consumer2, clientTags2);
          subscriptions.put(consumer3, getSubscription(PID_3, EMPTY_TASKS, clientTags3));
          hostTags.put(consumer3, clientTags3);
          subscriptions.put(consumer4, getSubscription(PID_4, EMPTY_TASKS, clientTags4));
          hostTags.put(consumer4, clientTags4);
          subscriptions.put(consumer5, getSubscription(PID_5, EMPTY_TASKS, clientTags5));
          hostTags.put(consumer5, clientTags5);
          subscriptions.put(consumer6, getSubscription(PID_6, EMPTY_TASKS, clientTags6));
          hostTags.put(consumer6, clientTags6);
  
          final Map<String, ConsumerPartitionAssignor.Assignment> assignments = partitionAssignor
              .assign(metadata, new ConsumerPartitionAssignor.GroupSubscription(subscriptions))
              .groupAssignment();
  
          verifyIdealTaskDistributionReached(getClientTagDistributions(assignments, hostTags), Collections.singletonList(ALL_TAG_KEYS.get(1)));
          verifyPartialTaskDistributionReached(getClientTagDistributions(assignments, hostTags), Collections.singletonList(ALL_TAG_KEYS.get(0)));
      }
  
      private Map<TaskId, ClientTagDistribution> getClientTagDistributions(final Map<String, ConsumerPartitionAssignor.Assignment> assignments,
                                                                           final Map<String, Map<String, String>> hostTags) {
          final Map<TaskId, ClientTagDistribution> taskClientTags = new HashMap<>();
  
          for (final Map.Entry<String, ConsumerPartitionAssignor.Assignment> entry : assignments.entrySet()) {
              final AssignmentInfo info = AssignmentInfo.decode(entry.getValue().userData());
  
              for (final TaskId activeTaskId : info.activeTasks()) {
                  taskClientTags.putIfAbsent(activeTaskId, new ClientTagDistribution(activeTaskId));
                  final ClientTagDistribution tagDistribution = taskClientTags.get(activeTaskId);
                  tagDistribution.addActiveTags(hostTags.get(entry.getKey()));
              }
  
              for (final TaskId standbyTaskId : info.standbyTasks().keySet()) {
                  taskClientTags.putIfAbsent(standbyTaskId, new ClientTagDistribution(standbyTaskId));
                  final ClientTagDistribution tagDistribution = taskClientTags.get(standbyTaskId);
                  tagDistribution.addStandbyTags(hostTags.get(entry.getKey()));
              }
          }
  
          return taskClientTags;
      }
  
      private void verifyIdealTaskDistributionReached(final Map<TaskId, ClientTagDistribution> taskClientTags,
                                                      final List<String> tagsToCheck) {
          for (final Map.Entry<TaskId, ClientTagDistribution> entry: taskClientTags.entrySet()) {
              if (!tagsAmongStandbysAreDifferent(entry.getValue(), tagsToCheck))
                  throw new AssertionError("task " + entry.getKey() + "'s tag-distribution for " + tagsToCheck +
                      " among standbys is not ideal: " + entry.getValue());
  
              if (!tagsAmongActiveAndAllStandbysAreDifferent(entry.getValue(), tagsToCheck))
                  throw new AssertionError("task " + entry.getKey() + "'s tag-distribution for " + tagsToCheck +
                      " between active and standbys is not ideal: " + entry.getValue());
          }
      }
  
      private void verifyPartialTaskDistributionReached(final Map<TaskId, ClientTagDistribution> taskClientTags,
                                                        final List<String> tagsToCheck) {
          for (final Map.Entry<TaskId, ClientTagDistribution> entry: taskClientTags.entrySet()) {
              if (!tagsAmongActiveAndAtLeastOneStandbyIsDifferent(entry.getValue(), tagsToCheck))
                  throw new AssertionError("task " + entry.getKey() + "'s tag-distribution for " + tagsToCheck +
                      "between active and standbys is not partially ideal: " + entry.getValue());
          }
      }
  
      private static boolean tagsAmongActiveAndAllStandbysAreDifferent(final ClientTagDistribution tagDistribution,
                                                                       final List<String> tagsToCheck) {
          return tagDistribution.standbysClientTags.stream().allMatch(standbyTags ->
              tagsToCheck.stream().noneMatch(tag -> tagDistribution.activeClientTags.get(tag).equals(standbyTags.get(tag))));
      }
  
      private static boolean tagsAmongActiveAndAtLeastOneStandbyIsDifferent(final ClientTagDistribution tagDistribution,
                                                                            final List<String> tagsToCheck) {
          return tagDistribution.standbysClientTags.stream().anyMatch(standbyTags ->
              tagsToCheck.stream().noneMatch(tag -> tagDistribution.activeClientTags.get(tag).equals(standbyTags.get(tag))));
      }
  
      private static boolean tagsAmongStandbysAreDifferent(final ClientTagDistribution tagDistribution,
                                                           final List<String> tagsToCheck) {
          final Map<String, Integer> statistics = new HashMap<>();
  
          for (final Map<String, String> tags : tagDistribution.standbysClientTags) {
              for (final Map.Entry<String, String> tag : tags.entrySet()) {
                  if (tagsToCheck.contains(tag.getKey())) {
                      final String tagValue = tag.getValue();
                      final Integer tagValueOccurrence = statistics.getOrDefault(tagValue, 0);
                      statistics.put(tagValue, tagValueOccurrence + 1);
                  }
              }
          }
  
          return statistics.values().stream().noneMatch(occurrence -> occurrence > 1);
      }
  
      private void setupTopology(final int numOfStatefulTopologies, final int numOfStatelessTopologies) {
          if (numOfStatefulTopologies + numOfStatelessTopologies > 5) {
              throw new IllegalArgumentException("Should not have more than 5 topologies, but have " + numOfStatefulTopologies);
          }
  
          for (int i = 0; i < numOfStatelessTopologies; i++) {
              builder.addSource(null, "source" + i, null, null, null, "topic" + i);
              builder.addProcessor("processor" + i, new MockApiProcessorSupplier<>(), "source" + i);
          }
  
          for (int i = numOfStatelessTopologies; i < numOfStatelessTopologies + numOfStatefulTopologies; i++) {
              builder.addSource(null, "source" + i, null, null, null, "topic" + i);
              builder.addProcessor("processor" + i, new MockApiProcessorSupplier<>(), "source" + i);
              builder.addStateStore(new MockKeyValueStoreBuilder("store" + i, false), "processor" + i);
          }
      }
  
      private static final class ClientTagDistribution {
          private final TaskId taskId;
          private final Map<String, String> activeClientTags;
          private final List<Map<String, String>> standbysClientTags;
  
          ClientTagDistribution(final TaskId taskId) {
              this.taskId = taskId;
              this.activeClientTags = new HashMap<>();
              this.standbysClientTags = new ArrayList<>();
          }
  
          void addActiveTags(final Map<String, String> activeClientTags) {
              if (!this.activeClientTags.isEmpty()) {
                  throw new IllegalStateException("Found multiple active tasks for " + taskId + ", this should not happen");
              }
              this.activeClientTags.putAll(activeClientTags);
          }
  
          void addStandbyTags(final Map<String, String> standbyClientTags) {
              this.standbysClientTags.add(standbyClientTags);
          }
  
          @Override
          public String toString() {
              return "ClientTagDistribution{" +
                  "taskId=" + taskId +
                  ", activeClientTags=" + activeClientTags +
                  ", standbysClientTags=" + standbysClientTags +
                  '}';
          }
      }
  
      /**
       * Helper for building the input to createMockAdminClient in cases where we don't care about the actual offsets
       * @param changelogTopics The names of all changelog topics in the topology
       * @param topicsNumPartitions The number of partitions for the corresponding changelog topic, such that the number
       *            of partitions of the ith topic in changelogTopics is given by the ith element of topicsNumPartitions
       */
      private static Map<TopicPartition, Long> getTopicPartitionOffsetsMap(final List<String> changelogTopics,
                                                                           final List<Integer> topicsNumPartitions) {
          if (changelogTopics.size() != topicsNumPartitions.size()) {
              throw new IllegalStateException("Passed in " + changelogTopics.size() + " changelog topic names, but " +
                      topicsNumPartitions.size() + " different numPartitions for the topics");
          }
          final Map<TopicPartition, Long> changelogEndOffsets = new HashMap<>();
          for (int i = 0; i < changelogTopics.size(); ++i) {
              final String topic = changelogTopics.get(i);
              final int numPartitions = topicsNumPartitions.get(i);
              for (int partition = 0; partition < numPartitions; ++partition) {
                  changelogEndOffsets.put(new TopicPartition(topic, partition), Long.MAX_VALUE);
              }
          }
          return changelogEndOffsets;
      }
  
      private static ConsumerPartitionAssignor.Subscription getSubscription(final ProcessId processId,
                                                                            final Collection<TaskId> prevActiveTasks,
                                                                            final Map<String, String> clientTags) {
          return new ConsumerPartitionAssignor.Subscription(
              singletonList("source1"),
              new SubscriptionInfo(LATEST_SUPPORTED_VERSION, LATEST_SUPPORTED_VERSION, processId, null,
                  getTaskOffsetSums(prevActiveTasks), (byte) 0, 0, clientTags).encode()
          );
      }
  
      // Stub offset sums for when we only care about the prev/standby task sets, not the actual offsets
      private static Map<TaskId, Long> getTaskOffsetSums(final Collection<TaskId> activeTasks) {
          final Map<TaskId, Long> taskOffsetSums = activeTasks.stream().collect(Collectors.toMap(t -> t, t -> Task.LATEST_OFFSET));
          taskOffsetSums.putAll(EMPTY_TASKS.stream().collect(Collectors.toMap(t -> t, t -> 0L)));
          return taskOffsetSums;
      }
  }
  