/*
 * 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.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.ListOffsetsResult;
import org.apache.kafka.clients.admin.ListOffsetsResult.ListOffsetsResultInfo;
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.internals.KafkaFutureImpl;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.StreamsConfig.InternalConfig;
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.InternalTopicManager;
import org.apache.kafka.streams.processor.internals.Task;
import org.apache.kafka.streams.processor.internals.TopologyMetadata.Subtopology;
import org.apache.kafka.test.MockClientSupplier;
import org.apache.kafka.test.MockInternalTopicManager;

import org.hamcrest.BaseMatcher;
import org.hamcrest.Description;
import org.hamcrest.Matcher;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Optional;
import java.util.Random;
import java.util.Set;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Function;
import java.util.stream.Collectors;

import static java.util.Collections.emptyMap;
import static java.util.Collections.emptySet;
import static org.apache.kafka.common.utils.Utils.entriesToMap;
import static org.apache.kafka.common.utils.Utils.intersection;
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.StreamsAssignmentProtocolVersions.LATEST_SUPPORTED_VERSION;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.hasItems;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.not;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anySet;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public final class AssignmentTestUtils {

    public static final ProcessId PID_1 = processIdForInt(1);
    public static final ProcessId PID_2 = processIdForInt(2);
    public static final ProcessId PID_3 = processIdForInt(3);
    public static final ProcessId PID_4 = processIdForInt(4);
    public static final ProcessId PID_5 = processIdForInt(5);
    public static final ProcessId PID_6 = processIdForInt(6);
    public static final ProcessId PID_7 = processIdForInt(7);
    public static final ProcessId PID_8 = processIdForInt(8);
    public static final ProcessId PID_9 = processIdForInt(9);

    public static final String RACK_0 = "rack0";
    public static final String RACK_1 = "rack1";
    public static final String RACK_2 = "rack2";
    public static final String RACK_3 = "rack3";
    public static final String RACK_4 = "rack4";

    public static final Node NODE_0 = new Node(0, "node0", 1, RACK_0);
    public static final Node NODE_1 = new Node(1, "node1", 1, RACK_1);
    public static final Node NODE_2 = new Node(2, "node2", 1, RACK_2);
    public static final Node NODE_3 = new Node(3, "node3", 1, RACK_3);
    public static final Node NODE_4 = new Node(4, "node4", 1, RACK_4);
    public static final Node NO_RACK_NODE = new Node(3, "node3", 1);

    public static final Node[] REPLICA_0 = new Node[] {NODE_0, NODE_1};
    public static final Node[] REPLICA_1 = new Node[] {NODE_1, NODE_2};
    public static final Node[] REPLICA_2 = new Node[] {NODE_0, NODE_2};
    public static final Node[] REPLICA_3 = new Node[] {NODE_1, NODE_3};
    public static final Node[] REPLICA_4 = new Node[] {NODE_3, NODE_4};

    public static final String TP_0_NAME = "topic0";
    public static final String TP_1_NAME = "topic1";
    public static final String TP_2_NAME = "topic2";
    public static final String TP_3_NAME = "topic3";

    public static final String CHANGELOG_TP_0_NAME = "store-0-changelog";
    public static final String CHANGELOG_TP_1_NAME = "store-1-changelog";
    public static final String CHANGELOG_TP_2_NAME = "store-2-changelog";
    public static final String CHANGELOG_TP_3_NAME = "store-3-changelog";

    public static final TopicPartition CHANGELOG_TP_0_0 = new TopicPartition(CHANGELOG_TP_0_NAME, 0);
    public static final TopicPartition CHANGELOG_TP_0_1 = new TopicPartition(CHANGELOG_TP_0_NAME, 1);
    public static final TopicPartition CHANGELOG_TP_0_2 = new TopicPartition(CHANGELOG_TP_0_NAME, 2);
    public static final TopicPartition CHANGELOG_TP_0_3 = new TopicPartition(CHANGELOG_TP_0_NAME, 3);
    public static final TopicPartition CHANGELOG_TP_0_4 = new TopicPartition(CHANGELOG_TP_0_NAME, 4);
    public static final TopicPartition CHANGELOG_TP_0_5 = new TopicPartition(CHANGELOG_TP_0_NAME, 5);
    public static final TopicPartition CHANGELOG_TP_0_6 = new TopicPartition(CHANGELOG_TP_0_NAME, 6);
    public static final TopicPartition CHANGELOG_TP_1_0 = new TopicPartition(CHANGELOG_TP_1_NAME, 0);
    public static final TopicPartition CHANGELOG_TP_1_1 = new TopicPartition(CHANGELOG_TP_1_NAME, 1);
    public static final TopicPartition CHANGELOG_TP_1_2 = new TopicPartition(CHANGELOG_TP_1_NAME, 2);
    public static final TopicPartition CHANGELOG_TP_1_3 = new TopicPartition(CHANGELOG_TP_1_NAME, 3);
    public static final TopicPartition CHANGELOG_TP_2_0 = new TopicPartition(CHANGELOG_TP_2_NAME, 0);
    public static final TopicPartition CHANGELOG_TP_2_1 = new TopicPartition(CHANGELOG_TP_2_NAME, 1);
    public static final TopicPartition CHANGELOG_TP_2_2 = new TopicPartition(CHANGELOG_TP_2_NAME, 2);
    public static final TopicPartition CHANGELOG_TP_2_3 = new TopicPartition(CHANGELOG_TP_2_NAME, 3);
    public static final TopicPartition CHANGELOG_TP_3_0 = new TopicPartition(CHANGELOG_TP_3_NAME, 0);
    public static final TopicPartition CHANGELOG_TP_3_1 = new TopicPartition(CHANGELOG_TP_3_NAME, 1);
    public static final TopicPartition CHANGELOG_TP_3_2 = new TopicPartition(CHANGELOG_TP_3_NAME, 2);

    public static final TopicPartition TP_0_0 = new TopicPartition(TP_0_NAME, 0);
    public static final TopicPartition TP_0_1 = new TopicPartition(TP_0_NAME, 1);
    public static final TopicPartition TP_0_2 = new TopicPartition(TP_0_NAME, 2);
    public static final TopicPartition TP_0_3 = new TopicPartition(TP_0_NAME, 3);
    public static final TopicPartition TP_0_4 = new TopicPartition(TP_0_NAME, 4);
    public static final TopicPartition TP_0_5 = new TopicPartition(TP_0_NAME, 5);
    public static final TopicPartition TP_0_6 = new TopicPartition(TP_0_NAME, 6);
    public static final TopicPartition TP_1_0 = new TopicPartition(TP_1_NAME, 0);
    public static final TopicPartition TP_1_1 = new TopicPartition(TP_1_NAME, 1);
    public static final TopicPartition TP_1_2 = new TopicPartition(TP_1_NAME, 2);
    public static final TopicPartition TP_1_3 = new TopicPartition(TP_1_NAME, 3);
    public static final TopicPartition TP_2_0 = new TopicPartition(TP_2_NAME, 0);
    public static final TopicPartition TP_2_1 = new TopicPartition(TP_2_NAME, 1);
    public static final TopicPartition TP_2_2 = new TopicPartition(TP_2_NAME, 2);
    public static final TopicPartition TP_2_3 = new TopicPartition(TP_2_NAME, 3);
    public static final TopicPartition TP_3_0 = new TopicPartition(TP_3_NAME, 0);
    public static final TopicPartition TP_3_1 = new TopicPartition(TP_3_NAME, 1);
    public static final TopicPartition TP_3_2 = new TopicPartition(TP_3_NAME, 2);

    public static final PartitionInfo PI_0_0 = new PartitionInfo(TP_0_NAME, 0, NODE_0, REPLICA_0, REPLICA_0);
    public static final PartitionInfo PI_0_1 = new PartitionInfo(TP_0_NAME, 1, NODE_1, REPLICA_1, REPLICA_1);
    public static final PartitionInfo PI_0_2 = new PartitionInfo(TP_0_NAME, 2, NODE_1, REPLICA_1, REPLICA_1);
    public static final PartitionInfo PI_0_3 = new PartitionInfo(TP_0_NAME, 3, NODE_2, REPLICA_2, REPLICA_2);
    public static final PartitionInfo PI_0_4 = new PartitionInfo(TP_0_NAME, 4, NODE_3, REPLICA_3, REPLICA_3);
    public static final PartitionInfo PI_0_5 = new PartitionInfo(TP_0_NAME, 5, NODE_4, REPLICA_4, REPLICA_4);
    public static final PartitionInfo PI_0_6 = new PartitionInfo(TP_0_NAME, 6, NODE_2, REPLICA_2, REPLICA_2);
    public static final PartitionInfo PI_1_0 = new PartitionInfo(TP_1_NAME, 0, NODE_2, REPLICA_2, REPLICA_2);
    public static final PartitionInfo PI_1_1 = new PartitionInfo(TP_1_NAME, 1, NODE_3, REPLICA_3, REPLICA_3);
    public static final PartitionInfo PI_1_2 = new PartitionInfo(TP_1_NAME, 2, NODE_0, REPLICA_0, REPLICA_0);
    public static final PartitionInfo PI_1_3 = new PartitionInfo(TP_1_NAME, 3, NODE_1, REPLICA_1, REPLICA_1);
    public static final PartitionInfo PI_2_0 = new PartitionInfo(TP_2_NAME, 0, NODE_4, REPLICA_4, REPLICA_4);
    public static final PartitionInfo PI_2_1 = new PartitionInfo(TP_2_NAME, 1, NODE_3, REPLICA_3, REPLICA_3);
    public static final PartitionInfo PI_2_2 = new PartitionInfo(TP_2_NAME, 2, NODE_1, REPLICA_4, REPLICA_4);
    public static final PartitionInfo PI_2_3 = new PartitionInfo(TP_2_NAME, 3, NODE_0, REPLICA_0, REPLICA_0);
    public static final PartitionInfo PI_3_0 = new PartitionInfo(TP_3_NAME, 0, NODE_2, REPLICA_2, REPLICA_2);
    public static final PartitionInfo PI_3_1 = new PartitionInfo(TP_3_NAME, 1, NODE_3, REPLICA_3, REPLICA_3);
    public static final PartitionInfo PI_3_2 = new PartitionInfo(TP_3_NAME, 2, NODE_4, REPLICA_4, REPLICA_4);

    public static final TaskId TASK_0_0 = new TaskId(0, 0);
    public static final TaskId TASK_0_1 = new TaskId(0, 1);
    public static final TaskId TASK_0_2 = new TaskId(0, 2);
    public static final TaskId TASK_0_3 = new TaskId(0, 3);
    public static final TaskId TASK_0_4 = new TaskId(0, 4);
    public static final TaskId TASK_0_5 = new TaskId(0, 5);
    public static final TaskId TASK_0_6 = new TaskId(0, 6);
    public static final TaskId TASK_1_0 = new TaskId(1, 0);
    public static final TaskId TASK_1_1 = new TaskId(1, 1);
    public static final TaskId TASK_1_2 = new TaskId(1, 2);
    public static final TaskId TASK_1_3 = new TaskId(1, 3);
    public static final TaskId TASK_2_0 = new TaskId(2, 0);
    public static final TaskId TASK_2_1 = new TaskId(2, 1);
    public static final TaskId TASK_2_2 = new TaskId(2, 2);
    public static final TaskId TASK_2_3 = new TaskId(2, 3);
    public static final TaskId TASK_3_0 = new TaskId(3, 0);
    public static final TaskId TASK_3_1 = new TaskId(3, 1);
    public static final TaskId TASK_3_2 = new TaskId(3, 2);

    public static final TaskId NAMED_TASK_T0_0_0 = new TaskId(0, 0, "topology0");
    public static final TaskId NAMED_TASK_T0_0_1 = new TaskId(0, 1, "topology0");
    public static final TaskId NAMED_TASK_T0_1_0 = new TaskId(1, 0, "topology0");
    public static final TaskId NAMED_TASK_T0_1_1 = new TaskId(1, 1, "topology0");
    public static final TaskId NAMED_TASK_T1_0_0 = new TaskId(0, 0, "topology1");
    public static final TaskId NAMED_TASK_T1_0_1 = new TaskId(0, 1, "topology1");
    public static final TaskId NAMED_TASK_T2_0_0 = new TaskId(0, 0, "topology2");
    public static final TaskId NAMED_TASK_T2_2_0 = new TaskId(2, 0, "topology2");

    public static final Subtopology SUBTOPOLOGY_0 = new Subtopology(0, null);
    public static final Subtopology SUBTOPOLOGY_1 = new Subtopology(1, null);
    public static final Subtopology SUBTOPOLOGY_2 = new Subtopology(2, null);

    public static final Set<TaskId> EMPTY_TASKS = emptySet();
    public static final Map<TopicPartition, Long> EMPTY_CHANGELOG_END_OFFSETS = new HashMap<>();
    public static final List<String> EMPTY_RACK_AWARE_ASSIGNMENT_TAGS = Collections.emptyList();
    public static final Map<String, String> EMPTY_CLIENT_TAGS = Collections.emptyMap();

    private static final String USER_END_POINT = "localhost:8080";
    private static final String APPLICATION_ID = "stream-partition-assignor-test";
    private static Random random;
    public static final String TOPIC_PREFIX = "topic";
    public static final String CHANGELOG_TOPIC_PREFIX = "changelog-topic";
    public static final String RACK_PREFIX = "rack";

    private AssignmentTestUtils() {}

    static Map<ProcessId, ClientState> getClientStatesMap(final ClientState... states) {
        final Map<ProcessId, ClientState> clientStates = new HashMap<>();
        int nthState = 1;
        for (final ClientState state : states) {
            clientStates.put(processIdForInt(nthState), state);
            ++nthState;
        }
        return clientStates;
    }

    // If you don't care about setting the end offsets for each specific topic partition, the helper method
    // getTopicPartitionOffsetMap is useful for building this input map for all partitions
    public static AdminClient createMockAdminClientForAssignor(final Map<TopicPartition, Long> changelogEndOffsets, final boolean mockListOffsets) {
        final AdminClient adminClient = mock(AdminClient.class);

        final ListOffsetsResult result = mock(ListOffsetsResult.class);
        if (mockListOffsets) {
            when(adminClient.listOffsets(any())).thenReturn(result);
        }
        for (final Map.Entry<TopicPartition, Long> entry : changelogEndOffsets.entrySet()) {
            final KafkaFutureImpl<ListOffsetsResultInfo> partitionFuture = new KafkaFutureImpl<>();
            final ListOffsetsResultInfo info = mock(ListOffsetsResultInfo.class);
            lenient().when(info.offset()).thenReturn(entry.getValue());
            partitionFuture.complete(info);
            lenient().when(result.partitionResult(entry.getKey())).thenReturn(partitionFuture);
        }

        return adminClient;
    }

    public static SubscriptionInfo getInfo(final ProcessId processId,
                                           final Set<TaskId> prevTasks,
                                           final Set<TaskId> standbyTasks) {
        return new SubscriptionInfo(
            LATEST_SUPPORTED_VERSION, LATEST_SUPPORTED_VERSION, processId, null, getTaskOffsetSums(prevTasks, standbyTasks), (byte) 0, 0, EMPTY_CLIENT_TAGS);
    }

    public static SubscriptionInfo getInfo(final ProcessId processId,
                                           final Set<TaskId> prevTasks,
                                           final Set<TaskId> standbyTasks,
                                           final String userEndPoint) {
        return new SubscriptionInfo(
            LATEST_SUPPORTED_VERSION, LATEST_SUPPORTED_VERSION, processId, userEndPoint, getTaskOffsetSums(prevTasks, standbyTasks), (byte) 0, 0, EMPTY_CLIENT_TAGS);
    }

    public static SubscriptionInfo getInfo(final ProcessId processId,
                                           final Set<TaskId> prevTasks,
                                           final Set<TaskId> standbyTasks,
                                           final byte uniqueField) {
        return new SubscriptionInfo(
            LATEST_SUPPORTED_VERSION, LATEST_SUPPORTED_VERSION, processId, null, getTaskOffsetSums(prevTasks, standbyTasks), uniqueField, 0, EMPTY_CLIENT_TAGS);
    }

    public static SubscriptionInfo getInfo(final ProcessId processId,
                                           final Set<TaskId> prevTasks,
                                           final Set<TaskId> standbyTasks,
                                           final byte uniqueField,
                                           final Map<String, String> clientTags) {
        return new SubscriptionInfo(
            LATEST_SUPPORTED_VERSION, LATEST_SUPPORTED_VERSION, processId, null, getTaskOffsetSums(prevTasks, standbyTasks), uniqueField, 0, clientTags);
    }

    // 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 Collection<TaskId> standbyTasks) {
        final Map<TaskId, Long> taskOffsetSums = activeTasks.stream().collect(Collectors.toMap(t -> t, t -> Task.LATEST_OFFSET));
        taskOffsetSums.putAll(standbyTasks.stream().collect(Collectors.toMap(t -> t, t -> 0L)));
        return taskOffsetSums;
    }

    /**
     * Builds a ProcessId with a UUID by repeating the given number n. For valid n, it is guaranteed that the returned UUIDs satisfy
     * the same relation relative to others as their parameter n does: iff n < m, then uuidForInt(n) < uuidForInt(m)
     */
    public static ProcessId processIdForInt(final int n) {
        return new ProcessId(new UUID(0, n));
    }

    static void assertValidAssignment(final int numStandbyReplicas,
                                      final Set<TaskId> statefulTasks,
                                      final Set<TaskId> statelessTasks,
                                      final Map<ProcessId, ClientState> assignedStates,
                                      final StringBuilder failureContext) {
        assertValidAssignment(
            numStandbyReplicas,
            0,
            statefulTasks,
            statelessTasks,
            assignedStates,
            failureContext
        );
    }

    static void assertValidAssignment(final int numStandbyReplicas,
                                      final int maxWarmupReplicas,
                                      final Set<TaskId> statefulTasks,
                                      final Set<TaskId> statelessTasks,
                                      final Map<ProcessId, ClientState> assignedStates,
                                      final StringBuilder failureContext) {
        final Map<TaskId, Set<ProcessId>> assignments = new TreeMap<>();
        for (final TaskId taskId : statefulTasks) {
            assignments.put(taskId, new TreeSet<>());
        }
        for (final TaskId taskId : statelessTasks) {
            assignments.put(taskId, new TreeSet<>());
        }
        for (final Map.Entry<ProcessId, ClientState> entry : assignedStates.entrySet()) {
            validateAndAddActiveAssignments(statefulTasks, statelessTasks, failureContext, assignments, entry);
            validateAndAddStandbyAssignments(statefulTasks, statelessTasks, failureContext, assignments, entry);
        }

        final AtomicInteger remainingWarmups = new AtomicInteger(maxWarmupReplicas);

        final TreeMap<TaskId, Set<ProcessId>> misassigned =
            assignments
                .entrySet()
                .stream()
                .filter(entry -> {
                    final int expectedActives = 1;
                    final boolean isStateless = statelessTasks.contains(entry.getKey());
                    final int expectedStandbys = isStateless ? 0 : numStandbyReplicas;
                    // We'll never assign even the expected number of standbys if they don't actually fit in the cluster
                    final int expectedAssignments = Math.min(
                        assignedStates.size(),
                        expectedActives + expectedStandbys
                    );
                    final int actualAssignments = entry.getValue().size();
                    if (actualAssignments == expectedAssignments) {
                        return false; // not misassigned
                    } else {
                        if (actualAssignments == expectedAssignments + 1 && remainingWarmups.get() > 0) {
                            remainingWarmups.getAndDecrement();
                            return false; // it's a warmup, so it's fine
                        } else {
                            return true; // misassigned
                        }
                    }
                })
                .collect(entriesToMap(TreeMap::new));

        if (!misassigned.isEmpty()) {
            assertThat("Found some over- or under-assigned tasks in the final assignment with " + numStandbyReplicas +
                    " and max warmups " + maxWarmupReplicas + " standby replicas, stateful tasks:" + statefulTasks +
                    ", and stateless tasks:" + statelessTasks + failureContext, misassigned, is(emptyMap()));
        }
    }

    private static void validateAndAddStandbyAssignments(final Set<TaskId> statefulTasks,
                                                         final Set<TaskId> statelessTasks,
                                                         final StringBuilder failureContext,
                                                         final Map<TaskId, Set<ProcessId>> assignments,
                                                         final Map.Entry<ProcessId, ClientState> entry) {
        for (final TaskId standbyTask : entry.getValue().standbyTasks()) {
            if (statelessTasks.contains(standbyTask)) {
                throw new AssertionError("Found a standby task for stateless task " + standbyTask + " on client " +
                        entry + " stateless tasks:" + statelessTasks + failureContext);
            } else if (assignments.containsKey(standbyTask)) {
                assignments.get(standbyTask).add(entry.getKey());
            } else {
                throw new AssertionError("Found an extra standby task " + standbyTask + " on client " +
                        entry + " but expected stateful tasks:" + statefulTasks + failureContext);
            }
        }
    }

    private static void validateAndAddActiveAssignments(final Set<TaskId> statefulTasks,
                                                        final Set<TaskId> statelessTasks,
                                                        final StringBuilder failureContext,
                                                        final Map<TaskId, Set<ProcessId>> assignments,
                                                        final Map.Entry<ProcessId, ClientState> entry) {
        for (final TaskId activeTask : entry.getValue().activeTasks()) {
            if (assignments.containsKey(activeTask)) {
                assignments.get(activeTask).add(entry.getKey());
            } else {
                throw new AssertionError("Found an extra active task " + activeTask + " on client " + entry + " but expected stateful tasks:" + statefulTasks + " and stateless tasks:" + statelessTasks + failureContext);
            }
        }
    }

    static void assertBalancedStatefulAssignment(final Set<TaskId> allStatefulTasks,
                                                 final Map<ProcessId, ClientState> clientStates,
                                                 final StringBuilder failureContext) {
        double maxStateful = Double.MIN_VALUE;
        double minStateful = Double.MAX_VALUE;
        for (final ClientState clientState : clientStates.values()) {
            final Set<TaskId> statefulTasks =
                intersection(HashSet::new, clientState.assignedTasks(), allStatefulTasks);
            final double statefulTaskLoad = 1.0 * statefulTasks.size() / clientState.capacity();
            maxStateful = Math.max(maxStateful, statefulTaskLoad);
            minStateful = Math.min(minStateful, statefulTaskLoad);
        }
        final double statefulDiff = maxStateful - minStateful;

        if (statefulDiff > 1.0) {
            final StringBuilder builder = new StringBuilder()
                .append("detected a stateful assignment balance factor violation: ")
                .append(statefulDiff)
                .append(">")
                .append(1.0)
                .append(" in: ");
            appendClientStates(builder, clientStates);
            fail(builder.append(failureContext).toString());
        }
    }

    static void assertBalancedActiveAssignment(final Map<ProcessId, ClientState> clientStates,
                                               final StringBuilder failureContext) {
        double maxActive = Double.MIN_VALUE;
        double minActive = Double.MAX_VALUE;
        for (final ClientState clientState : clientStates.values()) {
            final double activeTaskLoad = clientState.activeTaskLoad();
            maxActive = Math.max(maxActive, activeTaskLoad);
            minActive = Math.min(minActive, activeTaskLoad);
        }
        final double activeDiff = maxActive - minActive;
        if (activeDiff > 1.0) {
            final StringBuilder builder = new StringBuilder()
                .append("detected an active assignment balance factor violation: ")
                .append(activeDiff)
                .append(">")
                .append(1.0)
                .append(" in: ");
            appendClientStates(builder, clientStates);
            fail(builder.append(failureContext).toString());
        }
    }

    static void assertBalancedTasks(final Map<ProcessId, ClientState> clientStates) {
        assertBalancedTasks(clientStates, 1);
    }

    static void assertBalancedTasks(final Map<ProcessId, ClientState> clientStates, final int skewThreshold) {
        final TaskSkewReport taskSkewReport = analyzeTaskAssignmentBalance(clientStates, skewThreshold);
        if (taskSkewReport.totalSkewedTasks() > 0) {
            fail("Expected a balanced task assignment, but was: " + taskSkewReport);
        }
    }

    static TaskSkewReport analyzeTaskAssignmentBalance(final Map<ProcessId, ClientState> clientStates, final int skewThreshold) {
        final Function<Integer, Map<ProcessId, AtomicInteger>> initialClientCounts =
            i -> clientStates.keySet().stream().collect(Collectors.toMap(c -> c, c -> new AtomicInteger(0)));

        final Map<Integer, Map<ProcessId, AtomicInteger>> subtopologyToClientsWithPartition = new TreeMap<>();
        for (final Map.Entry<ProcessId, ClientState> entry : clientStates.entrySet()) {
            final ProcessId client = entry.getKey();
            final ClientState clientState = entry.getValue();
            for (final TaskId task : clientState.activeTasks()) {
                final int subtopology = task.subtopology();
                subtopologyToClientsWithPartition
                    .computeIfAbsent(subtopology, initialClientCounts)
                    .get(client)
                    .incrementAndGet();
            }
        }

        int maxTaskSkew = 0;
        final Set<Integer> skewedSubtopologies = new TreeSet<>();

        for (final Map.Entry<Integer, Map<ProcessId, AtomicInteger>> entry : subtopologyToClientsWithPartition.entrySet()) {
            final Map<ProcessId, AtomicInteger> clientsWithPartition = entry.getValue();
            int max = Integer.MIN_VALUE;
            int min = Integer.MAX_VALUE;
            for (final AtomicInteger count : clientsWithPartition.values()) {
                max = Math.max(max, count.get());
                min = Math.min(min, count.get());
            }
            final int taskSkew = max - min;
            maxTaskSkew = Math.max(maxTaskSkew, taskSkew);
            if (taskSkew > skewThreshold) {
                skewedSubtopologies.add(entry.getKey());
            }
        }

        return new TaskSkewReport(maxTaskSkew, skewedSubtopologies, subtopologyToClientsWithPartition);
    }

    static Matcher<ClientState> hasAssignedTasks(final int taskCount) {
        return hasProperty("assignedTasks", ClientState::assignedTaskCount, taskCount);
    }

    static Matcher<ClientState> hasActiveTasks(final int taskCount) {
        return hasProperty("activeTasks", ClientState::activeTaskCount, taskCount);
    }

    static Matcher<ClientState> hasStandbyTasks(final int taskCount) {
        return hasProperty("standbyTasks", ClientState::standbyTaskCount, taskCount);
    }

    static <V> Matcher<ClientState> hasProperty(final String propertyName,
                                                final Function<ClientState, V> propertyExtractor,
                                                final V propertyValue) {
        return new BaseMatcher<>() {
            @Override
            public void describeTo(final Description description) {
                description.appendText(propertyName).appendText(":").appendValue(propertyValue);
            }

            @Override
            public boolean matches(final Object actual) {
                if (actual instanceof ClientState) {
                    return Objects.equals(propertyExtractor.apply((ClientState) actual), propertyValue);
                } else {
                    return false;
                }
            }
        };
    }

    static void appendClientStates(final StringBuilder stringBuilder,
                                   final Map<ProcessId, ClientState> clientStates) {
        stringBuilder.append('{').append('\n');
        for (final Map.Entry<ProcessId, ClientState> entry : clientStates.entrySet()) {
            stringBuilder.append("  ").append(entry.getKey()).append(": ").append(entry.getValue()).append('\n');
        }
        stringBuilder.append('}').append('\n');
    }

    static final class TaskSkewReport {
        private final int maxTaskSkew;
        private final Set<Integer> skewedSubtopologies;
        private final Map<Integer, Map<ProcessId, AtomicInteger>> subtopologyToClientsWithPartition;

        private TaskSkewReport(final int maxTaskSkew,
                               final Set<Integer> skewedSubtopologies,
                               final Map<Integer, Map<ProcessId, AtomicInteger>> subtopologyToClientsWithPartition) {
            this.maxTaskSkew = maxTaskSkew;
            this.skewedSubtopologies = skewedSubtopologies;
            this.subtopologyToClientsWithPartition = subtopologyToClientsWithPartition;
        }

        int totalSkewedTasks() {
            return skewedSubtopologies.size();
        }

        Set<Integer> skewedSubtopologies() {
            return skewedSubtopologies;
        }

        @Override
        public String toString() {
            return "TaskSkewReport{" +
                "maxTaskSkew=" + maxTaskSkew +
                ", skewedSubtopologies=" + skewedSubtopologies +
                ", subtopologyToClientsWithPartition=" + subtopologyToClientsWithPartition +
                '}';
        }
    }

    static List<Node> getRandomNodes(final int nodeSize) {
        final List<Node> nodeList = new ArrayList<>(nodeSize);
        for (int i = 0; i < nodeSize; i++) {
            nodeList.add(new Node(i, "node" + i, 1, RACK_PREFIX + i));
        }
        final Random rand = getRandom();
        Collections.shuffle(nodeList, rand);
        return nodeList;
    }

    static Node[] getRandomReplica(final List<Node> nodeList, final int index, final int partition) {
        final Node firstNode = nodeList.get((index * partition) % nodeList.size());
        final Node secondNode = nodeList.get((index * partition + 1) % nodeList.size());
        return new Node[] {firstNode, secondNode};
    }

    static Cluster getRandomCluster(final int nodeSize, final int tpSize, final int partitionSize) {
        final List<Node> nodeList = getRandomNodes(nodeSize);
        final Set<PartitionInfo> partitionInfoSet = new HashSet<>();
        for (int i = 0; i < tpSize; i++) {
            for (int j = 0; j < partitionSize; j++) {
                final Node[] replica = getRandomReplica(nodeList, i, j);
                partitionInfoSet.add(
                    new PartitionInfo(TOPIC_PREFIX + i, j, replica[0], replica, replica));
            }
        }

        return new Cluster(
            "cluster",
            new HashSet<>(nodeList),
            partitionInfoSet,
            Collections.emptySet(),
            Collections.emptySet()
        );
    }

    static Map<ProcessId, Map<String, Optional<String>>> getRandomProcessRacks(final int clientSize, final int nodeSize) {
        final List<String> racks = new ArrayList<>(nodeSize);
        for (int i = 0; i < nodeSize; i++) {
            racks.add(RACK_PREFIX + i);
        }
        final Random rand = getRandom();
        Collections.shuffle(racks, rand);
        final Map<ProcessId, Map<String, Optional<String>>> processRacks = new HashMap<>();
        for (int i = 1; i <= clientSize; i++) {
            final String rack = racks.get(i % nodeSize);
            processRacks.put(processIdForInt(i), mkMap(mkEntry("1", Optional.of(rack))));
        }
        return processRacks;
    }

    static SortedMap<TaskId, Set<TopicPartition>> getTaskTopicPartitionMap(final int tpSize, final int partitionSize, final boolean changelog) {
        final SortedMap<TaskId, Set<TopicPartition>> taskTopicPartitionMap = new TreeMap<>();
        final String topicName = changelog ? CHANGELOG_TOPIC_PREFIX : TOPIC_PREFIX;
        for (int i = 0; i < tpSize; i++) {
            for (int j = 0; j < partitionSize; j++) {
                taskTopicPartitionMap.put(new TaskId(i, j), new HashSet<>(List.of(
                    new TopicPartition(topicName + i, j),
                    new TopicPartition(topicName + ((i + 1) % tpSize), j))
                ));
            }
        }
        return taskTopicPartitionMap;
    }

    static Map<Subtopology, Set<TaskId>> getTasksForTopicGroup(final int tpSize, final int partitionSize) {
        final Map<Subtopology, Set<TaskId>> tasksForTopicGroup = new HashMap<>();
        for (int i = 0; i < tpSize; i++) {
            for (int j = 0; j < partitionSize; j++) {
                final Subtopology subtopology = new Subtopology(i, null);
                tasksForTopicGroup.computeIfAbsent(subtopology, k -> new HashSet<>()).add(new TaskId(i, j));
            }
        }
        return tasksForTopicGroup;
    }

    static Map<String, Object> configProps(final String rackAwareConfig) {
        return configProps(rackAwareConfig, 0);
    }

    static Map<String, Object> configProps(final String rackAwareConfig, final int replicaNum) {
        final Map<String, Object> configurationMap = new HashMap<>();
        configurationMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
        configurationMap.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, USER_END_POINT);
        configurationMap.put(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, replicaNum);
        configurationMap.put(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_CONFIG, rackAwareConfig);

        final ReferenceContainer referenceContainer = new ReferenceContainer();
        configurationMap.put(InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
        return configurationMap;
    }

    static InternalTopicManager mockInternalTopicManagerForRandomChangelog(final int nodeSize, final int tpSize, final int partitionSize) {
        final MockTime time = new MockTime();
        final StreamsConfig streamsConfig = new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC));
        final MockClientSupplier mockClientSupplier = new MockClientSupplier();
        final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
            time,
            streamsConfig,
            mockClientSupplier.restoreConsumer,
            false
        );

        final Set<String> changelogNames = new HashSet<>();
        final List<Node> nodeList = getRandomNodes(nodeSize);
        final Map<String, List<TopicPartitionInfo>> topicPartitionInfo = new HashMap<>();
        for (int i = 0; i < tpSize; i++) {
            final String topicName = CHANGELOG_TOPIC_PREFIX + i;
            changelogNames.add(topicName);
            for (int j = 0; j < partitionSize; j++) {

                final Node[] replica = getRandomReplica(nodeList, i, j);
                final TopicPartitionInfo info = new TopicPartitionInfo(j, replica[0],
                    Arrays.asList(replica), Arrays.asList(replica));

                topicPartitionInfo.computeIfAbsent(topicName, tp -> new ArrayList<>()).add(info);
            }
        }

        final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
        doReturn(topicPartitionInfo).when(spyTopicManager).getTopicPartitionInfo(changelogNames);
        return spyTopicManager;
    }

    static SortedMap<ProcessId, ClientState> getRandomClientState(final int clientSize, final int tpSize, final int partitionSize, final int maxCapacity, final Set<TaskId> statefulTasks) {
        return getRandomClientState(clientSize, tpSize, partitionSize, maxCapacity, true, statefulTasks);
    }

    static List<Set<TaskId>> getRandomSubset(final Set<TaskId> taskIds, final int listSize) {
        final Random rand = getRandom();
        final List<TaskId> taskIdList = new ArrayList<>(taskIds);
        Collections.shuffle(taskIdList, rand);
        int start = 0;
        final List<Set<TaskId>> subSets = new ArrayList<>(listSize);
        for (int i = 0; i < listSize; i++) {
            final int remaining = taskIdList.size() - start;
            final Set<TaskId> subset = new HashSet<>();
            if (remaining != 0) {
                // In last round, get all tasks
                final int subSetSize = (i == listSize - 1) ? remaining : rand.nextInt(remaining) + 1;
                for (int j = 0; j < subSetSize; j++) {
                    subset.add(taskIdList.get(start + j));
                }
                start += subSetSize;
            }
            subSets.add(subset);
        }
        return subSets;
    }

    static SortedMap<ProcessId, ClientState> getRandomClientState(final int clientSize, final int tpSize, final int partitionSize, final int maxCapacity, final boolean initialAssignment, final Set<TaskId> statefulTasks) {
        final SortedMap<ProcessId, ClientState> clientStates = new TreeMap<>();
        final Map<TaskId, Long> taskLags = statefulTasks.stream().collect(Collectors.toMap(taskId -> taskId, taskId -> 0L));
        final Set<TaskId> taskIds = new HashSet<>();
        for (int i = 0; i < tpSize; i++) {
            for (int j = 0; j < partitionSize; j++) {
                taskIds.add(new TaskId(i, j));
            }
        }

        final Set<TaskId> missingTaskIds = taskLags.keySet().stream().filter(id -> !taskIds.contains(id)).collect(
            Collectors.toSet());
        if (!missingTaskIds.isEmpty()) {
            throw new IllegalArgumentException(missingTaskIds + " missing in all task ids " + taskIds);
        }

        final List<Set<TaskId>> previousActives = getRandomSubset(taskIds, clientSize);
        final List<Set<TaskId>> previousStandbys = getRandomSubset(statefulTasks, clientSize);

        final Random rand = getRandom();

        for (int i = 1; i <= clientSize; i++) {
            final int capacity = rand.nextInt(maxCapacity) + 1;
            final ProcessId processId = processIdForInt(i);
            final ClientState clientState = new ClientState(previousActives.get(i - 1), previousStandbys.get(i - 1), taskLags, EMPTY_CLIENT_TAGS, capacity, processId);
            clientStates.put(processId, clientState);
        }

        if (initialAssignment) {
            Iterator<Entry<ProcessId, ClientState>> iterator = clientStates.entrySet().iterator();
            final List<TaskId> taskIdList = new ArrayList<>(taskIds);
            Collections.shuffle(taskIdList, rand);
            for (final TaskId taskId : taskIdList) {
                if (!iterator.hasNext()) {
                    iterator = clientStates.entrySet().iterator();
                }
                iterator.next().getValue().assignActive(taskId);
            }
        }
        return clientStates;
    }

    static Cluster getClusterForAllTopics() {
        return new Cluster(
            "cluster",
            Set.of(NODE_0, NODE_1, NODE_2, NODE_3, NODE_4),
            Set.of(
                PI_0_0,
                PI_0_1,
                PI_0_2,
                PI_0_3,
                PI_0_4,
                PI_0_5,
                PI_0_6,
                PI_1_0,
                PI_1_1,
                PI_1_2,
                PI_1_3,
                PI_2_0,
                PI_2_1,
                PI_2_2,
                PI_2_3,
                PI_3_0,
                PI_3_1,
                PI_3_2
            ),
            Collections.emptySet(),
            Collections.emptySet()
        );
    }

    static Map<TaskId, Set<TopicPartition>> getTaskTopicPartitionMapForAllTasks() {
        return mkMap(
            mkEntry(TASK_0_0, Set.of(TP_0_0)),
            mkEntry(TASK_0_1, Set.of(TP_0_1)),
            mkEntry(TASK_0_2, Set.of(TP_0_2)),
            mkEntry(TASK_0_3, Set.of(TP_0_3)),
            mkEntry(TASK_0_4, Set.of(TP_0_4)),
            mkEntry(TASK_0_5, Set.of(TP_0_5)),
            mkEntry(TASK_0_6, Set.of(TP_0_6)),
            mkEntry(TASK_1_0, Set.of(TP_1_0)),
            mkEntry(TASK_1_1, Set.of(TP_1_1)),
            mkEntry(TASK_1_2, Set.of(TP_1_2)),
            mkEntry(TASK_1_3, Set.of(TP_1_3)),
            mkEntry(TASK_2_0, Set.of(TP_2_0)),
            mkEntry(TASK_2_1, Set.of(TP_2_1)),
            mkEntry(TASK_2_2, Set.of(TP_2_2)),
            mkEntry(TASK_2_3, Set.of(TP_2_3)),
            mkEntry(TASK_3_0, Set.of(TP_3_0)),
            mkEntry(TASK_3_1, Set.of(TP_3_1)),
            mkEntry(TASK_3_2, Set.of(TP_3_2))
        );
    }

    static Map<Subtopology, Set<TaskId>> getTasksForTopicGroup() {
        return mkMap(
            mkEntry(new Subtopology(0, null), Set.of(TASK_0_0, TASK_0_1, TASK_0_2, TASK_0_3, TASK_0_4, TASK_0_5, TASK_0_6)),
            mkEntry(new Subtopology(1, null), Set.of(TASK_1_0, TASK_1_1, TASK_1_2, TASK_1_3)),
            mkEntry(new Subtopology(2, null), Set.of(TASK_2_0, TASK_2_1, TASK_2_2, TASK_2_3)),
            mkEntry(new Subtopology(3, null), Set.of(TASK_3_0, TASK_3_1, TASK_3_2))
        );
    }

    static Map<TaskId, Set<TopicPartition>> getTaskChangelogMapForAllTasks() {
        return mkMap(
            mkEntry(TASK_0_0, Set.of(CHANGELOG_TP_0_0)),
            mkEntry(TASK_0_1, Set.of(CHANGELOG_TP_0_1)),
            mkEntry(TASK_0_2, Set.of(CHANGELOG_TP_0_2)),
            mkEntry(TASK_0_3, Set.of(CHANGELOG_TP_0_3)),
            mkEntry(TASK_0_4, Set.of(CHANGELOG_TP_0_4)),
            mkEntry(TASK_0_5, Set.of(CHANGELOG_TP_0_5)),
            mkEntry(TASK_0_6, Set.of(CHANGELOG_TP_0_6)),
            mkEntry(TASK_1_0, Set.of(CHANGELOG_TP_1_0)),
            mkEntry(TASK_1_1, Set.of(CHANGELOG_TP_1_1)),
            mkEntry(TASK_1_2, Set.of(CHANGELOG_TP_1_2)),
            mkEntry(TASK_1_3, Set.of(CHANGELOG_TP_1_3)),
            mkEntry(TASK_2_0, Set.of(CHANGELOG_TP_2_0)),
            mkEntry(TASK_2_1, Set.of(CHANGELOG_TP_2_1)),
            mkEntry(TASK_2_2, Set.of(CHANGELOG_TP_2_2)),
            mkEntry(TASK_2_3, Set.of(CHANGELOG_TP_2_3)),
            mkEntry(TASK_3_0, Set.of(CHANGELOG_TP_3_0)),
            mkEntry(TASK_3_1, Set.of(CHANGELOG_TP_3_1)),
            mkEntry(TASK_3_2, Set.of(CHANGELOG_TP_3_2))
        );
    }

    static InternalTopicManager mockInternalTopicManagerForChangelog() {
        final MockTime time = new MockTime();
        final StreamsConfig streamsConfig = new StreamsConfig(configProps(StreamsConfig.RACK_AWARE_ASSIGNMENT_STRATEGY_MIN_TRAFFIC));
        final MockClientSupplier mockClientSupplier = new MockClientSupplier();
        final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
            time,
            streamsConfig,
            mockClientSupplier.restoreConsumer,
            false
        );

        final MockInternalTopicManager spyTopicManager = spy(mockInternalTopicManager);
        doReturn(
            mkMap(
                mkEntry(
                    CHANGELOG_TP_0_NAME, Arrays.asList(
                        new TopicPartitionInfo(0, NODE_0, Arrays.asList(REPLICA_0), Collections.emptyList()),
                        new TopicPartitionInfo(1, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList()),
                        new TopicPartitionInfo(2, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList()),
                        new TopicPartitionInfo(3, NODE_2, Arrays.asList(REPLICA_2), Collections.emptyList()),
                        new TopicPartitionInfo(4, NODE_3, Arrays.asList(REPLICA_3), Collections.emptyList()),
                        new TopicPartitionInfo(5, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList()),
                        new TopicPartitionInfo(6, NODE_0, Arrays.asList(REPLICA_0), Collections.emptyList())
                    )
                ),
                mkEntry(
                    CHANGELOG_TP_1_NAME, Arrays.asList(
                        new TopicPartitionInfo(0, NODE_2, Arrays.asList(REPLICA_2), Collections.emptyList()),
                        new TopicPartitionInfo(1, NODE_3, Arrays.asList(REPLICA_3), Collections.emptyList()),
                        new TopicPartitionInfo(2, NODE_0, Arrays.asList(REPLICA_0), Collections.emptyList()),
                        new TopicPartitionInfo(3, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList())
                    )
                ),
                mkEntry(
                    CHANGELOG_TP_2_NAME, Arrays.asList(
                        new TopicPartitionInfo(0, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList()),
                        new TopicPartitionInfo(1, NODE_2, Arrays.asList(REPLICA_2), Collections.emptyList()),
                        new TopicPartitionInfo(2, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList()),
                        new TopicPartitionInfo(3, NODE_3, Arrays.asList(REPLICA_3), Collections.emptyList())
                    )
                ),
                mkEntry(
                    CHANGELOG_TP_3_NAME, Arrays.asList(
                        new TopicPartitionInfo(0, NODE_4, Arrays.asList(REPLICA_4), Collections.emptyList()),
                        new TopicPartitionInfo(1, NODE_3, Arrays.asList(REPLICA_3), Collections.emptyList()),
                        new TopicPartitionInfo(2, NODE_1, Arrays.asList(REPLICA_1), Collections.emptyList())
                    )
                )
            )
        ).when(spyTopicManager).getTopicPartitionInfo(anySet());
        return spyTopicManager;
    }

    static Map<Subtopology, Set<TaskId>> getTopologyGroupTaskMap() {
        return Collections.singletonMap(SUBTOPOLOGY_0, Collections.singleton(new TaskId(1, 1)));
    }

    static void verifyStandbySatisfyRackReplica(
        final Set<TaskId> taskIds,
        final Map<ProcessId, String> racksForProcess,
        final Map<ProcessId, ClientState> clientStateMap,
        final Integer replica,
        final boolean relaxRackCheck,
        final Map<ProcessId, Integer> standbyTaskCount
    ) {
        if (standbyTaskCount != null) {
            for (final Entry<ProcessId, ClientState> entry : clientStateMap.entrySet()) {
                final int expected = standbyTaskCount.get(entry.getKey());
                final int actual = entry.getValue().standbyTaskCount();
                assertEquals(expected, actual, "StandbyTaskCount for " + entry.getKey() + " doesn't match");
            }
        }
        for (final TaskId taskId : taskIds) {
            int activeCount = 0;
            int standbyCount = 0;
            final Map<String, ProcessId> racks = new HashMap<>();
            for (final Map.Entry<ProcessId, ClientState> entry : clientStateMap.entrySet()) {
                final ProcessId processId = entry.getKey();
                final ClientState clientState = entry.getValue();

                if (!relaxRackCheck && clientState.hasAssignedTask(taskId)) {
                    final String rack = racksForProcess.get(processId);
                    assertThat("Task " + taskId + " appears in both " + processId + " and " + racks.get(rack), racks.keySet(), not(hasItems(rack)));
                    racks.put(rack, processId);
                }

                boolean hasActive = false;
                if (clientState.hasActiveTask(taskId)) {
                    activeCount++;
                    hasActive = true;
                }

                boolean hasStandby = false;
                if (clientState.hasStandbyTask(taskId)) {
                    standbyCount++;
                    hasStandby = true;
                }

                assertFalse(hasActive && hasStandby, clientState + " has both active and standby task " + taskId);
            }

            assertEquals(1, activeCount, "Task " + taskId + " should have 1 active task");
            if (replica != null) {
                assertEquals(replica.intValue(), standbyCount, "Task " + taskId + " has wrong replica count");
            }
        }
    }

    static Map<ProcessId, Integer> clientTaskCount(final Map<ProcessId, ClientState> clientStateMap,
        final Function<ClientState, Integer> taskFunc) {
        return clientStateMap.entrySet().stream().collect(Collectors.toMap(Entry::getKey, v -> taskFunc.apply(v.getValue())));
    }

    static Map<ProcessId, Map<String, Optional<String>>> getProcessRacksForAllProcess() {
        return mkMap(
            mkEntry(PID_1, mkMap(mkEntry("1", Optional.of(RACK_0)))),
            mkEntry(PID_2, mkMap(mkEntry("1", Optional.of(RACK_1)))),
            mkEntry(PID_3, mkMap(mkEntry("1", Optional.of(RACK_2)))),
            mkEntry(PID_4, mkMap(mkEntry("1", Optional.of(RACK_3)))),
            mkEntry(PID_5, mkMap(mkEntry("1", Optional.of(RACK_4)))),
            mkEntry(PID_6, mkMap(mkEntry("1", Optional.of(RACK_0)))),
            mkEntry(PID_7, mkMap(mkEntry("1", Optional.of(RACK_1))))
        );
    }

    static RackAwareTaskAssignor getRackAwareTaskAssignor(final AssignmentConfigs configs) {
        return getRackAwareTaskAssignor(configs, mkMap());
    }

    static RackAwareTaskAssignor getRackAwareTaskAssignor(final AssignmentConfigs configs, final Map<Subtopology, Set<TaskId>> taskForTopicGroup) {
        return spy(
            new RackAwareTaskAssignor(
                getClusterForAllTopics(),
                getTaskTopicPartitionMapForAllTasks(),
                getTaskChangelogMapForAllTasks(),
                taskForTopicGroup,
                getProcessRacksForAllProcess(),
                mockInternalTopicManagerForChangelog(),
                configs,
                new MockTime()
            )
        );
    }

    static void verifyTaskPlacementWithRackAwareAssignor(final RackAwareTaskAssignor rackAwareTaskAssignor,
                                                         final Set<TaskId> allTaskIds,
                                                         final Map<ProcessId, ClientState> clientStates,
                                                         final boolean hasStandby,
                                                         final boolean enableRackAwareTaskAssignor) {
        // Verifies active and standby are in different clients
        verifyStandbySatisfyRackReplica(allTaskIds, rackAwareTaskAssignor.racksForProcess(), clientStates, null, true, null);

        if (enableRackAwareTaskAssignor) {
            verify(rackAwareTaskAssignor, times(2)).optimizeActiveTasks(any(), any(), anyInt(), anyInt());
            verify(rackAwareTaskAssignor, hasStandby ? times(1) : never()).optimizeStandbyTasks(any(), anyInt(), anyInt(), any());
        } else {
            verify(rackAwareTaskAssignor, never()).optimizeActiveTasks(any(), any(), anyInt(), anyInt());
            verify(rackAwareTaskAssignor, never()).optimizeStandbyTasks(any(), anyInt(), anyInt(), any());
        }
    }

    static SortedMap<ProcessId, ClientState> copyClientStateMap(final Map<ProcessId, ClientState> originalMap) {
        return new TreeMap<>(originalMap
            .entrySet()
            .stream()
            .collect(
                Collectors.toMap(
                    Entry::getKey,
                    entry -> new ClientState(entry.getValue())
                )
            )
        );
    }

    static synchronized Random getRandom() {
        if (random == null) {
            final long seed = System.currentTimeMillis();
            System.out.println("seed for getRandom: " + seed);
            random = new Random(seed);
        }
        return random;
    }
}
