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.assignment;
  
  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.TopicPartitionInfo;
  import org.apache.kafka.common.errors.TimeoutException;
  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.AssignmentConfigs;
  import org.apache.kafka.streams.processor.assignment.ProcessId;
  import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
  import org.apache.kafka.test.MockClientSupplier;
  import org.apache.kafka.test.MockInternalTopicManager;
  
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.Arguments;
  import org.junit.jupiter.params.provider.MethodSource;
  import org.junit.jupiter.params.provider.ValueSource;
  import org.mockito.Mockito;
  
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Map.Entry;
  import java.util.Optional;
  import java.util.Set;
  import java.util.SortedMap;
  import java.util.SortedSet;
  import java.util.TreeMap;
  import java.util.TreeSet;
  import java.util.stream.Stream;
  
  import static java.util.Collections.emptyMap;
  import static java.util.Collections.emptySet;
  import static org.apache.kafka.common.utils.Utils.mkEntry;
  import static org.apache.kafka.common.utils.Utils.mkMap;
  import static org.apache.kafka.common.utils.Utils.mkSortedSet;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.CHANGELOG_TP_0_0;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.CHANGELOG_TP_0_NAME;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.EMPTY_CLIENT_TAGS;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NODE_0;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NODE_1;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NODE_2;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.NO_RACK_NODE;
  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.PI_0_0;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.PI_0_1;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.RACK_1;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.RACK_2;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.REPLICA_1;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_0;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_1;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_0_2;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_1_0;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_1_1;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TASK_1_2;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_0_0;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_0_NAME;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.TP_1_0;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertBalancedTasks;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.assertValidAssignment;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.clientTaskCount;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.configProps;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getClusterForAllTopics;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getProcessRacksForAllProcess;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomClientState;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomCluster;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getRandomProcessRacks;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskChangelogMapForAllTasks;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMap;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTaskTopicPartitionMapForAllTasks;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTasksForTopicGroup;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.getTopologyGroupTaskMap;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForChangelog;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.mockInternalTopicManagerForRandomChangelog;
  import static org.apache.kafka.streams.processor.internals.assignment.AssignmentTestUtils.verifyStandbySatisfyRackReplica;
  import static org.hamcrest.MatcherAssert.assertThat;
  import static org.hamcrest.Matchers.greaterThanOrEqualTo;
  import static org.hamcrest.Matchers.lessThanOrEqualTo;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertFalse;
  import static org.junit.jupiter.api.Assertions.assertInstanceOf;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  import static org.junit.jupiter.api.Assertions.assertTrue;
  import static org.mockito.ArgumentMatchers.anySet;
  import static org.mockito.ArgumentMatchers.eq;
  import static org.mockito.Mockito.doReturn;
  import static org.mockito.Mockito.doThrow;
  import static org.mockito.Mockito.never;
  import static org.mockito.Mockito.spy;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  
  public class RackAwareTaskAssignorTest {
  
      private final MockTime time = new MockTime();
      private final StreamsConfig streamsConfig = new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC));
      private final MockClientSupplier mockClientSupplier = new MockClientSupplier();
  
      private int trafficCost;
      private int nonOverlapCost;
  
      private final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
              time,
              streamsConfig,
              mockClientSupplier.restoreConsumer,
              false
      );
  
      static Stream<Arguments> paramStoreType() {
          return Stream.of(
              Arguments.of(true, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY),
              Arguments.of(false, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC),
              Arguments.of(true, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC),
              Arguments.of(false, StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY)
          );
      }
  
      public void setUp(final boolean stateful) {
          if (stateful) {
              trafficCost = 10;
              nonOverlapCost = 1;
          } else {
              trafficCost = 1;
              nonOverlapCost = 0;
          }
      }
  
      private AssignmentConfigs getRackAwareEnabledConfig(final String assignmentStrategy) {
          return new AssignorConfiguration(
              new StreamsConfig(configProps(assignmentStrategy)).originals()).assignmentConfigs();
      }
  
      private AssignmentConfigs getRackAwareEnabledConfigWithStandby(final int replicaNum, final String assignmentStrategy) {
          return new AssignorConfiguration(
              new StreamsConfig(configProps(assignmentStrategy, replicaNum)).originals()).assignmentConfigs();
      }
  
      private AssignmentConfigs getRackAwareDisabledConfig() {
          return new AssignorConfiguration(
              new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_NONE)).originals()).assignmentConfigs();
      }
  
      @ParameterizedTest
      @ValueSource(booleans = {true, false})
      public void shouldDisableAssignorFromConfig(final boolean stateful) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = spy(new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(true),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              mockInternalTopicManager,
              getRackAwareDisabledConfig(),
              time
          ));
  
          // False since partitionWithoutInfo10 is missing in cluster metadata
          assertFalse(assignor.canEnableRackAwareAssignor());
          verify(assignor, never()).populateTopicsToDescribe(anySet(), eq(false));
          verify(assignor, never()).populateTopicsToDescribe(anySet(), eq(true));
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldDisableActiveWhenMissingClusterInfo(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = spy(new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(true),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          ));
  
          // False since partitionWithoutInfo10 is missing in cluster metadata
          assertFalse(assignor.canEnableRackAwareAssignor());
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(false));
          verify(assignor, never()).populateTopicsToDescribe(anySet(), eq(true));
          assertFalse(assignor.populateTopicsToDescribe(new HashSet<>(), false));
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldDisableActiveWhenRackMissingInNode(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = spy(new RackAwareTaskAssignor(
              getClusterWithPartitionMissingRack(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          ));
  
          // False since nodeMissingRack has one node which doesn't have rack
          assertFalse(assignor.canEnableRackAwareAssignor());
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(false));
          verify(assignor, never()).populateTopicsToDescribe(anySet(), eq(true));
          assertFalse(assignor.populateTopicsToDescribe(new HashSet<>(), false));
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldReturnInvalidClientRackWhenRackMissingInClientConsumer(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(true),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
          // False since process1 doesn't have rackId
          assertFalse(assignor.validClientRack());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldReturnFalseWhenRackMissingInProcess(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessWithNoConsumerRacks(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
          // False since process1 doesn't have rackId
          assertFalse(assignor.validClientRack());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldPopulateRacksForProcess(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          // Throws since process1 doesn't have rackId
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
          final Map<ProcessId, String> racksForProcess = assignor.racksForProcess();
          assertEquals(mkMap(mkEntry(PID_1, RACK_1)), racksForProcess);
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldReturnInvalidClientRackWhenRackDiffersInSameProcess(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final Map<ProcessId, Map<String, Optional<String>>> processRacks = new HashMap<>();
  
          // Different consumers in same process have different rack ID. This shouldn't happen.
          // If happens, there's a bug somewhere
          processRacks.computeIfAbsent(PID_1, k -> new HashMap<>()).put("consumer1", Optional.of("rack1"));
          processRacks.computeIfAbsent(PID_1, k -> new HashMap<>()).put("consumer2", Optional.of("rack2"));
  
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              processRacks,
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          assertFalse(assignor.validClientRack());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldEnableRackAwareAssignorWithoutDescribingTopics(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          // partitionWithoutInfo00 has rackInfo in cluster metadata
          assertTrue(assignor.canEnableRackAwareAssignor());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldEnableRackAwareAssignorWithCacheResult(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = spy(new RackAwareTaskAssignor(
              getClusterForTopic0(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          ));
  
          // partitionWithoutInfo00 has rackInfo in cluster metadata
          assertTrue(assignor.canEnableRackAwareAssignor());
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(false));
  
          // Should use cache result
          Mockito.reset(assignor);
          assertTrue(assignor.canEnableRackAwareAssignor());
          verify(assignor, never()).populateTopicsToDescribe(anySet(), eq(false));
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldEnableRackAwareAssignorWithDescribingTopics(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
          doReturn(
              Collections.singletonMap(
                  TP_0_NAME,
                  Collections.singletonList(
                      new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
                  )
              )
          ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME));
  
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterWithNoNode(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              spyTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          assertTrue(assignor.canEnableRackAwareAssignor());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldEnableRackAwareAssignorWithStandbyDescribingTopics(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
          doReturn(
              Collections.singletonMap(
                  TP_0_NAME,
                  Collections.singletonList(
                      new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
                  )
              )
          ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME));
  
          doReturn(
              Collections.singletonMap(
                  CHANGELOG_TP_0_NAME,
                  Collections.singletonList(
                      new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
                  )
              )
          ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(CHANGELOG_TP_0_NAME));
  
          final RackAwareTaskAssignor assignor = spy(new RackAwareTaskAssignor(
              getClusterWithNoNode(),
              getTaskTopicPartitionMapForTask0(),
              getTaskChangeLogTopicPartitionMapForTask0(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              spyTopicManager,
              getRackAwareEnabledConfigWithStandby(1, assignmentStrategy),
              time
          ));
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(false));
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(true));
  
          final Map<TopicPartition, Set<String>> racksForPartition = assignor.racksForPartition();
          final Map<TopicPartition, Set<String>> expected = mkMap(
              mkEntry(TP_0_0, Set.of(RACK_1, RACK_2)),
              mkEntry(CHANGELOG_TP_0_0, Set.of(RACK_1, RACK_2))
          );
          assertEquals(expected, racksForPartition);
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldDisableRackAwareAssignorWithStandbyDescribingTopicsFailure(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
          doReturn(
              Collections.singletonMap(
                  TP_0_NAME,
                  Collections.singletonList(
                      new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_1), Collections.emptyList())
                  )
              )
          ).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(TP_0_NAME));
  
          doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(
              CHANGELOG_TP_0_NAME));
  
          final RackAwareTaskAssignor assignor = spy(new RackAwareTaskAssignor(
              getClusterWithNoNode(),
              getTaskTopicPartitionMapForTask0(),
              getTaskChangeLogTopicPartitionMapForTask0(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              spyTopicManager,
              getRackAwareEnabledConfigWithStandby(1, assignmentStrategy),
              time
          ));
  
          assertFalse(assignor.canEnableRackAwareAssignor());
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(false));
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(true));
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldDisableRackAwareAssignorWithDescribingTopicsFailure(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
          doThrow(new TimeoutException("Timeout describing topic")).when(spyTopicManager).getTopicPartitionInfo(Collections.singleton(
              TP_0_NAME));
  
          final RackAwareTaskAssignor assignor = spy(new RackAwareTaskAssignor(
              getClusterWithNoNode(),
              getTaskTopicPartitionMapForTask0(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForProcess0(),
              spyTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          ));
  
          assertFalse(assignor.canEnableRackAwareAssignor());
          verify(assignor, times(1)).populateTopicsToDescribe(anySet(), eq(false));
          verify(assignor, never()).populateTopicsToDescribe(anySet(), eq(true));
          assertTrue(assignor.populateTopicsToDescribe(new HashSet<>(), false));
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeEmptyActiveTasks(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState1.assignActiveTasks(Set.of(TASK_0_1, TASK_1_1));
  
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_1, clientState1)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet();
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          final long originalCost = assignor.activeTasksCost(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(0, originalCost);
  
          final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(0, cost);
  
          assertEquals(Set.of(TASK_0_1, TASK_1_1), clientState1.activeTasks());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeActiveTasks(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
              mkEntry(new Subtopology(0, null), Set.of(TASK_0_0, TASK_0_1)),
              mkEntry(new Subtopology(1, null), Set.of(TASK_1_0, TASK_1_1))
          );
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              tasksForTopicGroup,
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState3 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState1.assignActiveTasks(Set.of(TASK_0_1, TASK_1_1));
          clientState2.assignActive(TASK_1_0);
          clientState3.assignActive(TASK_0_0);
  
          // task_0_0 has same rack as ProcessId_1
          // task_0_1 has same rack as ProcessId_2 and ProcessId_3
          // task_1_0 has same rack as ProcessId_1 and ProcessId_3
          // task_1_1 has same rack as ProcessId_2
          // Optimal assignment is ProcessId_1: {0_0, 1_0}, ProcessId_2: {1_1}, ProcessId_3: {0_1} which result in no cross rack traffic
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_1, clientState1),
              mkEntry(PID_2, clientState2),
              mkEntry(PID_3, clientState3)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          int expected = stateful ? 40 : 4;
          final long originalCost = assignor.activeTasksCost(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(expected, originalCost);
  
          expected = stateful ? 4 : 0;
          final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(expected, cost);
  
          assertEquals(Set.of(TASK_0_0, TASK_1_0), clientState1.activeTasks());
          assertEquals(Set.of(TASK_1_1), clientState2.activeTasks());
          assertEquals(Set.of(TASK_0_1), clientState3.activeTasks());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeRandomActive(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final int nodeSize = 30;
          final int tpSize = 40;
          final int partitionSize = 3;
          final int clientSize = 30;
          final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize, partitionSize, false);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getRandomCluster(nodeSize, tpSize, partitionSize),
              taskTopicPartitionMap,
              mkMap(),
              getTasksForTopicGroup(tpSize, partitionSize),
              getRandomProcessRacks(clientSize, nodeSize),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
          final SortedMap<ProcessId, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize, partitionSize, 1, taskIds);
  
          final Map<ProcessId, Integer> clientTaskCount = clientTaskCount(clientStateMap, ClientState::activeTaskCount);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          final long originalCost = assignor.activeTasksCost(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertThat(cost, lessThanOrEqualTo(originalCost));
  
          for (final Entry<ProcessId, ClientState> entry : clientStateMap.entrySet()) {
              assertEquals((int) clientTaskCount.get(entry.getKey()), entry.getValue().activeTasks().size());
          }
  
          if (assignmentStrategy.equals(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_BALANCE_SUBTOPOLOGY)) {
              assertBalancedTasks(clientStateMap);
          }
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldMaintainOriginalAssignmentForMinCost(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final int nodeSize = 20;
          final int tpSize = 40;
          final int partitionSize = 3;
          final int clientSize = 30;
          final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(tpSize, partitionSize, false);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getRandomCluster(nodeSize, tpSize, partitionSize),
              taskTopicPartitionMap,
              mkMap(),
              getTasksForTopicGroup(tpSize, partitionSize),
              getRandomProcessRacks(clientSize, nodeSize),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
          final SortedMap<ProcessId, ClientState> clientStateMap = getRandomClientState(clientSize, tpSize, partitionSize, 1, taskIds);
  
          final Map<TaskId, ProcessId> taskClientMap = new HashMap<>();
          for (final Entry<ProcessId, ClientState> entry : clientStateMap.entrySet()) {
              entry.getValue().activeTasks().forEach(t -> taskClientMap.put(t, entry.getKey()));
          }
          assertEquals(taskIds.size(), taskClientMap.size());
  
          // Because trafficCost is 0, original assignment should be maintained
          assertTrue(assignor.canEnableRackAwareAssignor());
          final long originalCost = assignor.activeTasksCost(taskIds, clientStateMap, 0, 1);
          assertEquals(0, originalCost);
  
          final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, 0, 1);
  
          if (assignmentStrategy.equals(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC)) {
              assertEquals(0, cost);
  
              // Make sure assignment doesn't change
              for (final Entry<TaskId, ProcessId> entry : taskClientMap.entrySet()) {
                  final ClientState clientState = clientStateMap.get(entry.getValue());
                  assertTrue(clientState.hasAssignedTask(entry.getKey()));
              }
          } else {
              // Balanced assignment recalculate the assignment using max flow first, so original assignment may not be maintained
              assertValidAssignment(0, taskIds, emptySet(), clientStateMap, new StringBuilder());
          }
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeActiveTasksWithMoreClients(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
              mkEntry(new Subtopology(0, null), Set.of(TASK_0_0)),
              mkEntry(new Subtopology(1, null), Set.of(TASK_1_0))
          );
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              tasksForTopicGroup,
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState3 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState2.assignActive(TASK_1_0);
          clientState3.assignActive(TASK_0_0);
  
          // task_0_0 has same rack as ProcessId_1 and ProcessId_2
          // task_1_0 has same rack as ProcessId_1 and ProcessId_3
          // Optimal assignment is ProcessId_1: {}, ProcessId_2: {0_0}, ProcessId_3: {1_0} which result in no cross rack traffic
          // and keeps ProcessId_1 empty since it was originally empty
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_1, clientState1),
              mkEntry(PID_2, clientState2),
              mkEntry(PID_3, clientState3)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_1_0);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          final long originalCost = assignor.activeTasksCost(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          int expected = stateful ? 20 : 2;
          assertEquals(expected, originalCost);
  
          expected = stateful ? 2 : 0;
          final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(expected, cost);
  
          // ProcessId_1 remains empty
          assertEquals(Set.of(), clientState1.activeTasks());
          assertEquals(Set.of(TASK_0_0), clientState2.activeTasks());
          assertEquals(Set.of(TASK_1_0), clientState3.activeTasks());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeActiveTasksWithMoreClientsWithMoreThanOneTask(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
              mkEntry(new Subtopology(0, null), Set.of(TASK_0_0, TASK_0_1)),
              mkEntry(new Subtopology(1, null), Set.of(TASK_1_0))
          );
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              tasksForTopicGroup,
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState3 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState2.assignActiveTasks(Set.of(TASK_0_1, TASK_1_0));
          clientState3.assignActive(TASK_0_0);
  
          // task_0_0 has same rack as ProcessId_1 and ProcessId_2
          // task_0_1 has same rack as ProcessId_2 and ProcessId_3
          // task_1_0 has same rack as ProcessId_1 and ProcessId_3
          // Optimal assignment is ProcessId_1: {}, ProcessId_2: {0_0, 0_1}, ProcessId_3: {1_0} which result in no cross rack traffic
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_1, clientState1),
              mkEntry(PID_2, clientState2),
              mkEntry(PID_3, clientState3)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          final long originalCost = assignor.activeTasksCost(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          int expected = stateful ? 20 : 2;
          assertEquals(expected, originalCost);
  
          expected = stateful ? 2 : 0;
          final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(expected, cost);
  
          // Because original assignment is not balanced (3 tasks but client 0 has no task), we maintain it
          assertEquals(Set.of(), clientState1.activeTasks());
          assertEquals(Set.of(TASK_0_0, TASK_0_1), clientState2.activeTasks());
          assertEquals(Set.of(TASK_1_0), clientState3.activeTasks());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldBalanceAssignmentWithMoreCost(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = mkMap(
              mkEntry(new Subtopology(0, null), Set.of(TASK_0_0, TASK_0_1)),
              mkEntry(new Subtopology(1, null), Set.of(TASK_1_1))
          );
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              tasksForTopicGroup,
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState1.assignActiveTasks(Set.of(TASK_0_0, TASK_1_1));
          clientState2.assignActive(TASK_0_1);
  
          // task_0_0 has same rack as ProcessId_2
          // task_0_1 has same rack as ProcessId_2
          // task_1_1 has same rack as ProcessId_2
          // ProcessId_5 is not in same rack as any task
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_2, clientState1),
              mkEntry(PID_5, clientState2)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_1);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          final int expectedCost = stateful ? 10 : 1;
          final long originalCost = assignor.activeTasksCost(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(expectedCost, originalCost);
  
          final long cost = assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost);
          assertEquals(expectedCost, cost);
  
          if (stateful || assignmentStrategy.equals(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC)) {
              // Even though assigning all tasks to ProcessId_2 will result in min cost, but it's not balanced
              // assignment. That's why TASK_0_1 is still assigned to ProcessId_5
              assertEquals(Set.of(TASK_0_0, TASK_1_1), clientState1.activeTasks());
              assertEquals(Set.of(TASK_0_1), clientState2.activeTasks());
          } else {
              assertEquals(Set.of(TASK_0_0, TASK_0_1), clientState1.activeTasks());
              assertEquals(Set.of(TASK_1_1), clientState2.activeTasks());
          }
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldThrowIfMissingCallCanEnableRackAwareAssignor(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState1.assignActiveTasks(Set.of(TASK_0_0, TASK_1_1));
          clientState2.assignActive(TASK_0_1);
  
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_2, clientState1),
              mkEntry(PID_5, clientState2)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_1);
          final Exception exception = assertThrows(IllegalStateException.class,
              () -> assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost));
          Assertions.assertEquals("TopicPartition topic0-0 has no rack information. "
              + "Maybe forgot to call canEnableRackAwareAssignor first", exception.getMessage());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldThrowIfTaskInMultipleClients(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState1.assignActiveTasks(Set.of(TASK_0_0, TASK_1_1));
          clientState2.assignActiveTasks(Set.of(TASK_0_1, TASK_1_1));
  
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_2, clientState1),
              mkEntry(PID_5, clientState2)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_1);
          assertTrue(assignor.canEnableRackAwareAssignor());
          final Exception exception = assertThrows(IllegalArgumentException.class,
              () -> assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost));
          Assertions.assertEquals(
              "Task 1_1 assigned to multiple clients 00000000-0000-0000-0000-000000000005, "
                  + "00000000-0000-0000-0000-000000000002", exception.getMessage());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldThrowIfTaskMissingInClients(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForAllProcess(),
              mockInternalTopicManager,
              getRackAwareEnabledConfig(assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1);
  
          clientState1.assignActiveTasks(Set.of(TASK_0_0, TASK_1_1));
          clientState2.assignActive(TASK_0_1);
  
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_2, clientState1),
              mkEntry(PID_5, clientState2)
          ));
          final SortedSet<TaskId> taskIds = mkSortedSet(TASK_0_0, TASK_0_1, TASK_1_0, TASK_1_1);
          assertTrue(assignor.canEnableRackAwareAssignor());
          final Exception exception = assertThrows(IllegalArgumentException.class,
              () -> assignor.optimizeActiveTasks(taskIds, clientStateMap, trafficCost, nonOverlapCost));
          Assertions.assertEquals(
              "Task 1_0 not assigned to any client", exception.getMessage());
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldNotCrashForEmptyStandby(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              mkMap(),
              getTopologyGroupTaskMap(),
              getProcessRacksForAllProcess(),
              mockInternalTopicManagerForChangelog(),
              getRackAwareEnabledConfigWithStandby(1, assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_1
          );
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_2
          );
          final ClientState clientState3 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_3
          );
  
          clientState1.assignActiveTasks(Set.of(TASK_0_1, TASK_1_1));
          clientState2.assignActive(TASK_1_0);
          clientState3.assignActive(TASK_0_0);
  
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_1, clientState1),
              mkEntry(PID_2, clientState2),
              mkEntry(PID_3, clientState3)
          ));
  
          final long originalCost = assignor.standByTasksCost(new TreeSet<>(), clientStateMap, 10, 1);
          assertEquals(0, originalCost);
  
          final long cost = assignor.optimizeStandbyTasks(clientStateMap, 10, 1,
              (source, destination, task, clientStates) -> true);
          assertEquals(0, cost);
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeStandbyTasksWhenTasksAllMovable(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final int replicaCount = 2;
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              getTaskChangelogMapForAllTasks(),
              getTopologyGroupTaskMap(),
              getProcessRacksForAllProcess(),
              mockInternalTopicManagerForChangelog(),
              getRackAwareEnabledConfigWithStandby(replicaCount, assignmentStrategy),
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_1
          );
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_2
          );
          final ClientState clientState3 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_3
          );
          final ClientState clientState4 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_4
          );
          final ClientState clientState5 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_6
          );
          final ClientState clientState6 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_7
          );
  
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_1, clientState1),
              mkEntry(PID_2, clientState2),
              mkEntry(PID_3, clientState3),
              mkEntry(PID_4, clientState4),
              mkEntry(PID_6, clientState5),
              mkEntry(PID_7, clientState6)
          ));
  
          clientState1.assignActive(TASK_0_0);
          clientState2.assignActive(TASK_0_1);
          clientState3.assignActive(TASK_1_0);
          clientState4.assignActive(TASK_1_1);
          clientState5.assignActive(TASK_0_2);
          clientState6.assignActive(TASK_1_2);
  
          clientState1.assignStandbyTasks(Set.of(TASK_0_1, TASK_1_1)); // Cost 10
          clientState2.assignStandbyTasks(Set.of(TASK_0_0, TASK_1_0)); // Cost 10
          clientState3.assignStandbyTasks(Set.of(TASK_0_0, TASK_0_2)); // Cost 20
          clientState4.assignStandbyTasks(Set.of(TASK_0_1, TASK_1_2)); // Cost 10
          clientState5.assignStandbyTasks(Set.of(TASK_1_0, TASK_1_2)); // Cost 10
          clientState6.assignStandbyTasks(Set.of(TASK_0_2, TASK_1_1)); // Cost 10
  
          final SortedSet<TaskId> taskIds = new TreeSet<>(Set.of(TASK_0_0, TASK_0_1, TASK_0_2, TASK_1_0, TASK_1_1, TASK_1_2));
          final Map<ProcessId, Integer> standbyTaskCount = clientTaskCount(clientStateMap, ClientState::standbyTaskCount);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          verifyStandbySatisfyRackReplica(taskIds, assignor.racksForProcess(), clientStateMap, replicaCount, false, null);
  
          final long originalCost = assignor.standByTasksCost(taskIds, clientStateMap, 10, 1);
          assertEquals(60, originalCost);
  
          // Task can be moved anywhere so cost can be reduced to 30 compared to in shouldOptimizeStandbyTasksWithMovingConstraint it
          // can only be reduced to 50 since there are moving constraints
          final long cost = assignor.optimizeStandbyTasks(clientStateMap, 10, 1,
              (source, destination, task, clients) -> true);
          assertEquals(20, cost);
          // Don't validate tasks in different racks after moving
          verifyStandbySatisfyRackReplica(taskIds, assignor.racksForProcess(), clientStateMap, replicaCount, true, standbyTaskCount);
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeStandbyTasksWithMovingConstraint(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final int replicaCount = 2;
          final AssignmentConfigs assignorConfiguration = getRackAwareEnabledConfigWithStandby(replicaCount, assignmentStrategy);
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getClusterForAllTopics(),
              getTaskTopicPartitionMapForAllTasks(),
              getTaskChangelogMapForAllTasks(),
              getTopologyGroupTaskMap(),
              getProcessRacksForAllProcess(),
              mockInternalTopicManagerForChangelog(),
              assignorConfiguration,
              time
          );
  
          final ClientState clientState1 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_1
          );
          final ClientState clientState2 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_2
          );
          final ClientState clientState3 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_3
          );
          final ClientState clientState4 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_4
          );
          final ClientState clientState5 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_6
          );
          final ClientState clientState6 = new ClientState(emptySet(), emptySet(), emptyMap(), EMPTY_CLIENT_TAGS, 1,
              PID_7
          );
  
          final SortedMap<ProcessId, ClientState> clientStateMap = new TreeMap<>(mkMap(
              mkEntry(PID_1, clientState1),
              mkEntry(PID_2, clientState2),
              mkEntry(PID_3, clientState3),
              mkEntry(PID_4, clientState4),
              mkEntry(PID_6, clientState5),
              mkEntry(PID_7, clientState6)
          ));
  
          clientState1.assignActive(TASK_0_0);
          clientState2.assignActive(TASK_0_1);
          clientState3.assignActive(TASK_1_0);
          clientState4.assignActive(TASK_1_1);
          clientState5.assignActive(TASK_0_2);
          clientState6.assignActive(TASK_1_2);
  
          clientState1.assignStandbyTasks(Set.of(TASK_0_1, TASK_1_1)); // Cost 10
          clientState2.assignStandbyTasks(Set.of(TASK_0_0, TASK_1_0)); // Cost 10
          clientState3.assignStandbyTasks(Set.of(TASK_0_0, TASK_0_2)); // Cost 20
          clientState4.assignStandbyTasks(Set.of(TASK_0_1, TASK_1_2)); // Cost 10
          clientState5.assignStandbyTasks(Set.of(TASK_1_0, TASK_1_2)); // Cost 10
          clientState6.assignStandbyTasks(Set.of(TASK_0_2, TASK_1_1)); // Cost 10
  
          final SortedSet<TaskId> taskIds = new TreeSet<>(Set.of(TASK_0_0, TASK_0_1, TASK_0_2, TASK_1_0, TASK_1_1, TASK_1_2));
          final Map<ProcessId, Integer> standbyTaskCount = clientTaskCount(clientStateMap, ClientState::standbyTaskCount);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          verifyStandbySatisfyRackReplica(taskIds, assignor.racksForProcess(), clientStateMap, replicaCount, false, null);
  
          final long originalCost = assignor.standByTasksCost(taskIds, clientStateMap, 10, 1);
          assertEquals(60, originalCost);
  
          final StandbyTaskAssignor standbyTaskAssignor = StandbyTaskAssignorFactory.create(assignorConfiguration, assignor);
          assertInstanceOf(ClientTagAwareStandbyTaskAssignor.class, standbyTaskAssignor);
          final long cost = assignor.optimizeStandbyTasks(clientStateMap, 10, 1,
              standbyTaskAssignor::isAllowedTaskMovement);
          assertEquals(50, cost);
          // Validate tasks in different racks after moving
          verifyStandbySatisfyRackReplica(taskIds, assignor.racksForProcess(), clientStateMap, replicaCount, false, standbyTaskCount);
      }
  
      @ParameterizedTest
      @MethodSource("paramStoreType")
      public void shouldOptimizeRandomStandby(final boolean stateful, final String assignmentStrategy) {
          setUp(stateful);
          final int nodeSize = 50;
          final int tpSize = 60;
          final int partitionSize = 3;
          final int clientSize = 50;
          final int replicaCount = 3;
          final int maxCapacity = 3;
          final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = getTaskTopicPartitionMap(
              tpSize, partitionSize, false);
          final AssignmentConfigs assignorConfiguration = getRackAwareEnabledConfigWithStandby(replicaCount, assignmentStrategy);
  
          final RackAwareTaskAssignor assignor = new RackAwareTaskAssignor(
              getRandomCluster(nodeSize, tpSize, partitionSize),
              taskTopicPartitionMap,
              getTaskTopicPartitionMap(tpSize, partitionSize, true),
              getTopologyGroupTaskMap(),
              getRandomProcessRacks(clientSize, nodeSize),
              mockInternalTopicManagerForRandomChangelog(nodeSize, tpSize, partitionSize),
              assignorConfiguration,
              time
          );
  
          final SortedSet<TaskId> taskIds = (SortedSet<TaskId>) taskTopicPartitionMap.keySet();
          final SortedMap<ProcessId, ClientState> clientStateMap = getRandomClientState(clientSize,
              tpSize, partitionSize, maxCapacity, taskIds);
  
          final StandbyTaskAssignor standbyTaskAssignor = StandbyTaskAssignorFactory.create(
              assignorConfiguration, assignor);
          assertInstanceOf(ClientTagAwareStandbyTaskAssignor.class, standbyTaskAssignor);
          // Get a standby assignment
          standbyTaskAssignor.assign(clientStateMap, taskIds, taskIds, assignorConfiguration);
          final Map<ProcessId, Integer> standbyTaskCount = clientTaskCount(clientStateMap,
              ClientState::standbyTaskCount);
  
          assertTrue(assignor.canEnableRackAwareAssignor());
          verifyStandbySatisfyRackReplica(taskIds, assignor.racksForProcess(), clientStateMap,
              replicaCount, false, null);
  
          final long originalCost = assignor.standByTasksCost(taskIds, clientStateMap, 10, 1);
          assertThat(originalCost, greaterThanOrEqualTo(0L));
  
          final long cost = assignor.optimizeStandbyTasks(clientStateMap, 10, 1,
              standbyTaskAssignor::isAllowedTaskMovement);
          assertThat(cost, lessThanOrEqualTo(originalCost));
          // Validate tasks in different racks after moving
          verifyStandbySatisfyRackReplica(taskIds, assignor.racksForProcess(), clientStateMap,
              replicaCount, false, standbyTaskCount);
      }
  
      private Cluster getClusterForTopic0() {
          return new Cluster(
              "cluster",
              Set.of(NODE_0, NODE_1, NODE_2),
              Set.of(PI_0_0, PI_0_1),
              Collections.emptySet(),
              Collections.emptySet()
          );
      }
  
      private Cluster getClusterWithPartitionMissingRack() {
          final Node[] nodeMissingRack = new Node[]{NODE_0, NO_RACK_NODE};
          final PartitionInfo partitionInfoMissingNode = new PartitionInfo(TP_0_NAME, 0, NODE_0, nodeMissingRack, nodeMissingRack);
          return new Cluster(
              "cluster",
              Set.of(NODE_0, NODE_1, NODE_2),
              Set.of(partitionInfoMissingNode, PI_0_1),
              Collections.emptySet(),
              Collections.emptySet()
          );
      }
  
      private Cluster getClusterWithNoNode() {
          final PartitionInfo noNodeInfo = new PartitionInfo(TP_0_NAME, 0, null, new Node[0], new Node[0]);
  
          return new Cluster(
              "cluster",
              Set.of(NODE_0, NODE_1, NODE_2, Node.noNode()), // mockClientSupplier.setCluster requires noNode
              Collections.singleton(noNodeInfo),
              Collections.emptySet(),
              Collections.emptySet()
          );
      }
  
      private Map<ProcessId, Map<String, Optional<String>>> getProcessRacksForProcess0() {
          return getProcessRacksForProcess0(false);
      }
  
      private Map<ProcessId, Map<String, Optional<String>>> getProcessRacksForProcess0(final boolean missingRack) {
          final Map<ProcessId, Map<String, Optional<String>>> processRacks = new HashMap<>();
          final Optional<String> rack = missingRack ? Optional.empty() : Optional.of("rack1");
          processRacks.put(PID_1, Collections.singletonMap("consumer1", rack));
          return processRacks;
      }
  
      private Map<ProcessId, Map<String, Optional<String>>> getProcessWithNoConsumerRacks() {
          return mkMap(
              mkEntry(PID_1, mkMap())
          );
      }
  
      private Map<TaskId, Set<TopicPartition>> getTaskTopicPartitionMapForTask0() {
          return getTaskTopicPartitionMapForTask0(false);
      }
  
      private Map<TaskId, Set<TopicPartition>> getTaskChangeLogTopicPartitionMapForTask0() {
          return mkMap(
              mkEntry(TASK_0_0, Set.of(CHANGELOG_TP_0_0))
          );
      }
  
      private Map<TaskId, Set<TopicPartition>> getTaskTopicPartitionMapForTask0(final boolean extraTopic) {
          final Set<TopicPartition> topicPartitions = new HashSet<>();
          topicPartitions.add(TP_0_0);
          if (extraTopic) {
              topicPartitions.add(TP_1_0);
          }
          return Collections.singletonMap(TASK_0_0, topicPartitions);
      }
  }
  