/*
 * 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.connect.runtime.distributed;

import org.apache.kafka.clients.consumer.internals.RequestFuture;
import org.apache.kafka.common.message.JoinGroupResponseData;
import org.apache.kafka.common.utils.LogContext;
import org.apache.kafka.common.utils.MockTime;
import org.apache.kafka.common.utils.Time;
import org.apache.kafka.connect.runtime.TargetState;
import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.ConnectorsAndTasks;
import org.apache.kafka.connect.storage.AppliedConnectorConfig;
import org.apache.kafka.connect.storage.ClusterConfigState;
import org.apache.kafka.connect.util.ConnectUtils;
import org.apache.kafka.connect.util.ConnectorTaskId;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.junit.jupiter.MockitoExtension;
import org.mockito.junit.jupiter.MockitoSettings;
import org.mockito.quality.Strictness;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeAssignor.ClusterAssignment;
import static org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.WorkerLoad;
import static org.apache.kafka.connect.util.ConnectUtils.transformValues;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.notNull;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
@MockitoSettings(strictness = Strictness.STRICT_STUBS)
public class IncrementalCooperativeAssignorTest {

    // Offset isn't used in most tests but is required for creating a config snapshot object,
    // so just use some arbitrary constant for that
    private static final long CONFIG_OFFSET = 618;

    private Map<String, Integer> connectors;
    private Time time;
    private int rebalanceDelay;
    private IncrementalCooperativeAssignor assignor;
    private int generationId;
    private ClusterAssignment returnedAssignments;
    private Map<String, ConnectorsAndTasks> memberAssignments;

    @BeforeEach
    public void setup() {
        generationId = 1000;
        time = Time.SYSTEM;
        rebalanceDelay = DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT;
        connectors = new HashMap<>();
        addNewConnector("connector1", 4);
        addNewConnector("connector2", 4);
        memberAssignments = new HashMap<>();
        addNewEmptyWorkers("worker1");
        initAssignor();
    }

    public void initAssignor() {
        assignor = new IncrementalCooperativeAssignor(new LogContext(), time, rebalanceDelay);
        assignor.previousGenerationId = generationId;
    }

    @Test
    public void testTaskAssignmentWhenWorkerJoins() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();
        // First assignment with 1 worker and 2 connectors configured but not yet assigned
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1");
        assertConnectorAllocations(2);
        assertTaskAllocations(8);
        assertBalancedAndCompleteAllocation();

        // Second assignment with a second worker joining and all connectors running on previous worker
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 1);
        assertTaskAllocations(0, 4);

        // Third assignment after revocations
        performStandardRebalance();
        assertNoRevocations();
        assertDelay(0);
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        time.sleep(assignor.delay);
        // A fourth rebalance after delay should not change assignments
        performStandardRebalance();
        assertDelay(0);
        assertEmptyAssignment();
    }

    @Test
    public void testAssignmentsWhenWorkersJoinAfterRevocations()  {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        addNewConnector("connector3", 4);
        // First assignment with 1 worker and 3 connectors configured but not yet assigned
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1");
        assertConnectorAllocations(3);
        assertTaskAllocations(12);
        assertBalancedAndCompleteAllocation();

        // Second assignment with a second worker joining and all connectors running on previous worker
        // We should revoke.
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 2);
        assertTaskAllocations(0, 6);

        // Third assignment immediately after revocations, and a third worker joining.
        // This is a successive revoking rebalance. We should not perform any revocations
        // in this round
        addNewEmptyWorkers("worker3");
        performStandardRebalance();
        assertTrue(assignor.delay > 0);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 2);
        assertTaskAllocations(3, 3, 6);

        // Fourth assignment and a fourth worker joining
        // after first revoking rebalance is expired. We should revoke.
        time.sleep(assignor.delay);
        addNewEmptyWorkers("worker4");
        performStandardRebalance();
        assertWorkers("worker1", "worker2", "worker3", "worker4");
        assertConnectorAllocations(0, 0, 1, 1);
        assertTaskAllocations(0, 3, 3, 3);

        // Fifth assignment and a fifth worker joining after a revoking rebalance.
        // We shouldn't revoke and set a delay > initial interval
        addNewEmptyWorkers("worker5");
        performStandardRebalance();
        assertTrue(assignor.delay > 40);
        assertWorkers("worker1", "worker2", "worker3", "worker4", "worker5");
        assertConnectorAllocations(0, 0, 1, 1, 1);
        assertTaskAllocations(1, 2, 3, 3, 3);

        // Sixth assignment with sixth worker joining after the expiry.
        // Should revoke
        time.sleep(assignor.delay);
        addNewEmptyWorkers("worker6");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2", "worker3", "worker4", "worker5", "worker6");
        assertConnectorAllocations(0, 0, 0, 1, 1, 1);
        assertTaskAllocations(0, 1, 2, 2, 2, 2);

        // Follow up rebalance since there were revocations
        performStandardRebalance();
        assertWorkers("worker1", "worker2", "worker3", "worker4", "worker5", "worker6");
        assertConnectorAllocations(0, 0, 0, 1, 1, 1);
        assertTaskAllocations(2, 2, 2, 2, 2, 2);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testImmediateRevocationsWhenMaxDelayIs0()  {
        // Customize assignor for this test case
        rebalanceDelay = 0;
        time = new MockTime();
        initAssignor();

        addNewConnector("connector3", 4);
        // First assignment with 1 worker and 3 connectors configured but not yet assigned
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1");
        assertConnectorAllocations(3);
        assertTaskAllocations(12);
        assertBalancedAndCompleteAllocation();

        // Second assignment with a second worker joining and all connectors running on previous worker
        // We should revoke.
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 2);
        assertTaskAllocations(0, 6);

        // Third assignment immediately after revocations, and a third worker joining.
        // This is a successive revoking rebalance but we should still revoke as rebalance delay is 0
        addNewEmptyWorkers("worker3");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(3, 3, 4);

        // Follow up rebalance post revocations
        performStandardRebalance();
        assertWorkers("worker1", "worker2", "worker3");
        assertNoRevocations();
        assertConnectorAllocations(1, 1, 1);
        assertTaskAllocations(4, 4, 4);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testSuccessiveRevocationsWhenMaxDelayIsEqualToExpBackOffInitialInterval() {
        // Customize assignor for this test case
        rebalanceDelay = 1;
        initAssignor();

        addNewConnector("connector3", 4);
        // First assignment with 1 worker and 3 connectors configured but not yet assigned
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1");
        assertConnectorAllocations(3);
        assertTaskAllocations(12);
        assertBalancedAndCompleteAllocation();

        // Second assignment with a second worker joining and all connectors running on previous worker
        // We should revoke.
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 2);
        assertTaskAllocations(0, 6);

        // Third assignment immediately after revocations, and a third worker joining.
        // This is a successive revoking rebalance. We shouldn't revoke as maxDelay is 1 ms
        addNewEmptyWorkers("worker3");
        performStandardRebalance();
        assertDelay(1);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 2);
        assertTaskAllocations(3, 3, 6);
    }

    @Test
    public void testWorkerJoiningDuringDelayedRebalance() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        addNewConnector("connector3", 4);
        // First assignment with 1 worker and 3 connectors configured but not yet assigned
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1");
        assertConnectorAllocations(3);
        assertTaskAllocations(12);
        assertBalancedAndCompleteAllocation();

        // Second assignment with a second worker joining and all connectors running on previous worker
        // We should revoke.
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 2);
        assertTaskAllocations(0, 6);

        // Third assignment immediately after revocations, and a third worker joining.
        // This is a successive revoking rebalance. We should not perform any revocations
        // in this round, but can allocate the connectors and tasks revoked previously
        addNewEmptyWorkers("worker3");
        performStandardRebalance();
        assertTrue(assignor.delay > 0);
        assertWorkers("worker1", "worker2", "worker3");
        assertNoRevocations();
        assertConnectorAllocations(0, 1, 2);
        assertTaskAllocations(3, 3, 6);

        // Fourth assignment and a fourth worker joining
        // while delayed rebalance is active. We should not revoke
        time.sleep(assignor.delay / 2);
        addNewEmptyWorkers("worker4");
        performStandardRebalance();
        assertTrue(assignor.delay > 0);
        assertWorkers("worker1", "worker2", "worker3", "worker4");
        assertNoRevocations();
        assertConnectorAllocations(0, 0, 1, 2);
        assertTaskAllocations(0, 3, 3, 6);

        // Fifth assignment and a fifth worker joining
        // after the delay has expired. We should perform load-balancing
        // revocations
        time.sleep(assignor.delay);
        addNewEmptyWorkers("worker5");
        performStandardRebalance();
        assertWorkers("worker1", "worker2", "worker3", "worker4", "worker5");
        assertDelay(0);
        assertConnectorAllocations(0, 0, 0, 1, 1);
        assertTaskAllocations(0, 0, 2, 3, 3);

        // Sixth and final rebalance, as a follow-up to the revocations in the previous round.
        // Should allocate all previously-revoked connectors and tasks evenly across the cluster
        performStandardRebalance();
        assertDelay(0);
        assertNoRevocations();
        assertConnectorAllocations(0, 0, 1, 1, 1);
        assertTaskAllocations(2, 2, 2, 3, 3);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenWorkerLeavesPermanently() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        // First assignment with 2 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Second assignment with only one worker remaining in the group. The worker that left the
        // group was a follower. No re-assignments take place immediately and the count
        // down for the rebalance delay starts
        removeWorkers("worker2");
        performStandardRebalance();
        assertDelay(rebalanceDelay);
        assertWorkers("worker1");
        assertEmptyAssignment();

        time.sleep(rebalanceDelay / 2);

        // Third (incidental) assignment with still only one worker in the group. Max delay has not
        // been reached yet
        performStandardRebalance();
        assertDelay(rebalanceDelay / 2);
        assertEmptyAssignment();

        time.sleep(rebalanceDelay / 2 + 1);

        // Fourth assignment after delay expired
        performStandardRebalance();
        assertDelay(0);
        assertConnectorAllocations(2);
        assertTaskAllocations(8);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenWorkerBounces() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        // First assignment with 2 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Second assignment with only one worker remaining in the group. The worker that left the
        // group was a follower. No re-assignments take place immediately and the count
        // down for the rebalance delay starts
        removeWorkers("worker2");
        performStandardRebalance();
        assertDelay(rebalanceDelay);
        assertWorkers("worker1");
        assertEmptyAssignment();

        time.sleep(rebalanceDelay / 2);

        // Third (incidental) assignment with still only one worker in the group. Max delay has not
        // been reached yet
        performStandardRebalance();
        assertDelay(rebalanceDelay / 2);
        assertEmptyAssignment();

        time.sleep(rebalanceDelay / 4);

        // Fourth assignment with the second worker returning before the delay expires
        // Since the delay is still active, lost assignments are not reassigned yet
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(rebalanceDelay / 4);
        assertWorkers("worker1", "worker2");
        assertEmptyAssignment();

        time.sleep(rebalanceDelay / 4);

        // Fifth assignment with the same two workers. The delay has expired, so there
        // should be revocations giving back the assignments to the reappearing worker
        performStandardRebalance();
        assertDelay(0);
        assertNoRevocations();
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenLeaderLeavesPermanently() {
        // First assignment with 3 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2", "worker3");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(2, 3, 3);
        assertBalancedAndCompleteAllocation();

        // Second assignment with two workers remaining in the group. The worker that left the
        // group was the leader. The new leader has no previous assignments and is not tracking a
        // delay upon a leader's exit
        removeWorkers("worker1");
        // The fact that the leader bounces means that the assignor starts from a clean slate
        initAssignor();

        // Capture needs to be reset to point to the new assignor
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker2", "worker3");
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Third (incidental) assignment with still only two workers in the group.
        performStandardRebalance();
        assertDelay(0);
        assertEmptyAssignment();
    }

    @Test
    public void testTaskAssignmentWhenLeaderBounces() {
        // First assignment with 3 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2", "worker3");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(2, 3, 3);
        assertBalancedAndCompleteAllocation();

        // Second assignment with two workers remaining in the group. The worker that left the
        // group was the leader. The new leader has no previous assignments and is not tracking a
        // delay upon a leader's exit
        removeWorkers("worker1");
        // The fact that the leader bounces means that the assignor starts from a clean slate
        initAssignor();

        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker2", "worker3");
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Third assignment with the previous leader returning as a follower. In this case, the
        // arrival of the previous leader is treated as an arrival of a new worker. Reassignment
        // happens immediately, first with a revocation
        addNewEmptyWorkers("worker1");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(0, 3, 3);

        // Fourth assignment after revocations
        performStandardRebalance();
        assertDelay(0);
        assertNoRevocations();
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(2, 3, 3);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenFirstAssignmentAttemptFails() {
        // First assignment with 2 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2");
        performFailedRebalance();
        // This was the assignment that should have been sent, but didn't make it all the way
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 0);
        assertTaskAllocations(0, 0);

        // Second assignment happens with members returning the same assignments (memberConfigs)
        // as the first time. The assignor detects that the number of members did not change and
        // avoids the rebalance delay, treating the lost assignments as new assignments.
        performStandardRebalance();
        assertDelay(0);
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenSubsequentAssignmentAttemptFails() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        // First assignment with 2 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Second assignment triggered by a third worker joining. The computed assignment should
        // revoke tasks from the existing group. But the assignment won't be correctly delivered.
        addNewEmptyWorkers("worker3");
        performFailedRebalance();
        // This was the assignment that should have been sent, but didn't make it all the way
        assertDelay(0);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(0, 4, 4);

        // Third assignment happens with members returning the same assignments (memberConfigs)
        // as the first time. Since this is a consecutive revoking rebalance, delay should be non-zero
        // but no revoking rebalances
        performStandardRebalance();
        assertTrue(assignor.delay > 0);
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(0, 4, 4);

        // Wait for delay ms before triggering another rebalance
        time.sleep(assignor.delay);
        // Fourth assignment after revocations.
        performStandardRebalance();
        assertDelay(0);
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(0, 3, 3);

        // Fourth assignment after revocations
        performStandardRebalance();
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(2, 3, 3);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenSubsequentAssignmentAttemptFailsOutsideTheAssignor() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        // First assignment with 2 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Second assignment triggered by a third worker joining. The computed assignment should
        // revoke tasks from the existing group. But the assignment won't be correctly delivered
        // and sync group with fail on the leader worker.
        addNewEmptyWorkers("worker3");
        performFailedRebalance();
        // This was the assignment that should have been sent, but didn't make it all the way
        assertDelay(0);
        assertWorkers("worker1", "worker2", "worker3");
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(0, 4, 4);

        // Third assignment happens with members returning the same assignments (memberConfigs)
        // as the first time. Since this is a consecutive revoking rebalance, there should be delay
        // not leading to revocations.
        performRebalanceWithMismatchedGeneration();
        assertTrue(assignor.delay > 0);
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(0, 4, 4);

        // Wait for delay to get revoking rebalances
        time.sleep(assignor.delay);
        // Fourth assignment after revocations
        performStandardRebalance();
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(0, 3, 3);

        // Fifth and final rebalance
        performStandardRebalance();
        assertConnectorAllocations(0, 1, 1);
        assertTaskAllocations(2, 3, 3);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenConnectorsAreDeleted() {
        addNewConnector("connector3", 4);

        // First assignment with 1 worker and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(1, 2);
        assertTaskAllocations(6, 6);
        assertBalancedAndCompleteAllocation();

        // Second assignment with an updated config state that reflects removal of a connector
        removeConnector("connector3");
        performStandardRebalance();
        assertDelay(0);
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testAssignConnectorsWhenBalanced() {
        int num = 2;
        List<WorkerLoad> existingAssignment = IntStream.range(0, 3)
                .mapToObj(i -> workerLoad("worker" + i, i * num, num, i * num, num))
                .collect(Collectors.toList());

        List<WorkerLoad> expectedAssignment = existingAssignment.stream()
                .map(wl -> new WorkerLoad.Builder(wl.worker()).withCopies(wl.connectors(), wl.tasks()).build())
                .collect(Collectors.toList());
        expectedAssignment.get(0).connectors().addAll(List.of("connector6", "connector9"));
        expectedAssignment.get(1).connectors().addAll(List.of("connector7", "connector10"));
        expectedAssignment.get(2).connectors().add("connector8");

        List<String> newConnectors = newConnectors(6, 11);
        assignor.assignConnectors(existingAssignment, newConnectors);
        assertEquals(expectedAssignment, existingAssignment);
    }

    @Test
    public void testAssignTasksWhenBalanced() {
        int num = 2;
        List<WorkerLoad> existingAssignment = IntStream.range(0, 3)
                .mapToObj(i -> workerLoad("worker" + i, i * num, num, i * num, num))
                .collect(Collectors.toList());

        List<WorkerLoad> expectedAssignment = existingAssignment.stream()
                .map(wl -> new WorkerLoad.Builder(wl.worker()).withCopies(wl.connectors(), wl.tasks()).build())
                .collect(Collectors.toList());

        expectedAssignment.get(0).connectors().addAll(List.of("connector6", "connector9"));
        expectedAssignment.get(1).connectors().addAll(List.of("connector7", "connector10"));
        expectedAssignment.get(2).connectors().add("connector8");

        expectedAssignment.get(0).tasks().addAll(List.of(new ConnectorTaskId("task", 6), new ConnectorTaskId("task", 9)));
        expectedAssignment.get(1).tasks().addAll(List.of(new ConnectorTaskId("task", 7), new ConnectorTaskId("task", 10)));
        expectedAssignment.get(2).tasks().add(new ConnectorTaskId("task", 8));

        List<String> newConnectors = newConnectors(6, 11);
        assignor.assignConnectors(existingAssignment, newConnectors);
        List<ConnectorTaskId> newTasks = newTasks(6, 11);
        assignor.assignTasks(existingAssignment, newTasks);
        assertEquals(expectedAssignment, existingAssignment);
    }

    @Test
    public void testAssignConnectorsWhenImbalanced() {
        List<WorkerLoad> existingAssignment = new ArrayList<>();
        existingAssignment.add(workerLoad("worker0", 0, 2, 0, 2));
        existingAssignment.add(workerLoad("worker1", 2, 3, 2, 3));
        existingAssignment.add(workerLoad("worker2", 5, 4, 5, 4));
        existingAssignment.add(emptyWorkerLoad("worker3"));

        List<String> newConnectors = newConnectors(9, 24);
        List<ConnectorTaskId> newTasks = newTasks(9, 24);
        assignor.assignConnectors(existingAssignment, newConnectors);
        assignor.assignTasks(existingAssignment, newTasks);
        for (WorkerLoad worker : existingAssignment) {
            assertEquals(6, worker.connectorsSize());
            assertEquals(6, worker.tasksSize());
        }
    }

    @Test
    public void testLostAssignmentHandlingWhenWorkerBounces() {
        time = new MockTime();
        initAssignor();

        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        Map<String, WorkerLoad> configuredAssignment = new HashMap<>();
        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));

        // No lost assignments
        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
                new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());

        String flakyWorker = "worker1";
        WorkerLoad lostLoad = configuredAssignment.remove(flakyWorker);
        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
                .with(lostLoad.connectors(), lostLoad.tasks()).build();

        // Lost assignments detected - No candidate worker has appeared yet (worker with no assignments)
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());
        time.sleep(rebalanceDelay / 2);
        rebalanceDelay /= 2;

        // A new worker (probably returning worker) has joined
        configuredAssignment.put(flakyWorker, new WorkerLoad.Builder(flakyWorker).build());
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(flakyWorker),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());
        time.sleep(rebalanceDelay);

        // The new worker has still no assignments
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertTrue(configuredAssignment.getOrDefault(flakyWorker, new WorkerLoad.Builder(flakyWorker).build())
                                .connectors()
                                .containsAll(lostAssignments.connectors()),
            "Wrong assignment of lost connectors");
        assertTrue(configuredAssignment.getOrDefault(flakyWorker, new WorkerLoad.Builder(flakyWorker).build())
                                .tasks()
                                .containsAll(lostAssignments.tasks()),
            "Wrong assignment of lost tasks");
        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);
    }

    @Test
    public void testLostAssignmentHandlingWhenWorkerLeavesPermanently() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        Map<String, WorkerLoad> configuredAssignment = new HashMap<>();
        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));

        // No lost assignments
        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
                new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());

        String removedWorker = "worker1";
        WorkerLoad lostLoad = configuredAssignment.remove(removedWorker);
        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
                .with(lostLoad.connectors(), lostLoad.tasks()).build();

        // Lost assignments detected - No candidate worker has appeared yet (worker with no assignments)
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        assignor.previousMembers = new HashSet<>(memberAssignments.keySet());
        time.sleep(rebalanceDelay / 2);
        rebalanceDelay /= 2;

        // No new worker has joined
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        time.sleep(rebalanceDelay);

        ConnectorsAndTasks.Builder lostAssignmentsToReassign = new ConnectorsAndTasks.Builder();
        assignor.handleLostAssignments(lostAssignments, lostAssignmentsToReassign,
                new ArrayList<>(configuredAssignment.values()));

        assertTrue(lostAssignmentsToReassign.build().connectors().containsAll(lostAssignments.connectors()),
            "Wrong assignment of lost connectors");
        assertTrue(lostAssignmentsToReassign.build().tasks().containsAll(lostAssignments.tasks()),
            "Wrong assignment of lost tasks");
        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);
    }

    @Test
    public void testLostAssignmentHandlingWithMoreThanOneCandidates() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        Map<String, WorkerLoad> configuredAssignment = new HashMap<>();
        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));

        // No lost assignments
        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
                new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());

        String flakyWorker = "worker1";
        WorkerLoad lostLoad = configuredAssignment.remove(flakyWorker);
        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
                .with(lostLoad.connectors(), lostLoad.tasks()).build();

        String newWorker = "worker3";
        configuredAssignment.put(newWorker, new WorkerLoad.Builder(newWorker).build());

        // Lost assignments detected - A new worker also has joined that is not the returning worker
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(newWorker),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());
        time.sleep(rebalanceDelay / 2);
        rebalanceDelay /= 2;

        // Now two new workers have joined
        configuredAssignment.put(flakyWorker, new WorkerLoad.Builder(flakyWorker).build());
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        Set<String> expectedWorkers = Set.of(newWorker, flakyWorker);
        assertEquals(expectedWorkers,
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());
        time.sleep(rebalanceDelay);

        // The new workers have new assignments, other than the lost ones
        configuredAssignment.put(flakyWorker, workerLoad(flakyWorker, 6, 2, 8, 4));
        configuredAssignment.put(newWorker, workerLoad(newWorker, 8, 2, 12, 4));
        // we don't reflect these new assignments in memberConfigs currently because they are not
        // used in handleLostAssignments method
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        // both the newWorkers would need to be considered for re assignment of connectors and tasks
        List<String> listOfConnectorsInLast2Workers = new ArrayList<>();
        listOfConnectorsInLast2Workers.addAll(configuredAssignment.getOrDefault(newWorker, new WorkerLoad.Builder(flakyWorker).build())
            .connectors());
        listOfConnectorsInLast2Workers.addAll(configuredAssignment.getOrDefault(flakyWorker, new WorkerLoad.Builder(flakyWorker).build())
            .connectors());
        List<ConnectorTaskId> listOfTasksInLast2Workers = new ArrayList<>();
        listOfTasksInLast2Workers.addAll(configuredAssignment.getOrDefault(newWorker, new WorkerLoad.Builder(flakyWorker).build())
            .tasks());
        listOfTasksInLast2Workers.addAll(configuredAssignment.getOrDefault(flakyWorker, new WorkerLoad.Builder(flakyWorker).build())
            .tasks());
        assertTrue(listOfConnectorsInLast2Workers.containsAll(lostAssignments.connectors()),
            "Wrong assignment of lost connectors");
        assertTrue(listOfTasksInLast2Workers.containsAll(lostAssignments.tasks()),
            "Wrong assignment of lost tasks");
        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);
    }

    @Test
    public void testLostAssignmentHandlingWhenWorkerBouncesBackButFinallyLeaves() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        Map<String, WorkerLoad> configuredAssignment = new HashMap<>();
        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));

        // No lost assignments
        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
                new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());

        String veryFlakyWorker = "worker1";
        WorkerLoad lostLoad = configuredAssignment.remove(veryFlakyWorker);
        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
                .with(lostLoad.connectors(), lostLoad.tasks()).build();

        // Lost assignments detected - No candidate worker has appeared yet (worker with no assignments)
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());
        time.sleep(rebalanceDelay / 2);
        rebalanceDelay /= 2;

        // A new worker (probably returning worker) has joined
        configuredAssignment.put(veryFlakyWorker, new WorkerLoad.Builder(veryFlakyWorker).build());
        assignor.handleLostAssignments(lostAssignments, new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(veryFlakyWorker),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(time.milliseconds() + rebalanceDelay, assignor.scheduledRebalance);
        assertEquals(rebalanceDelay, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());
        time.sleep(rebalanceDelay);

        // The returning worker leaves permanently after joining briefly during the delay
        configuredAssignment.remove(veryFlakyWorker);
        ConnectorsAndTasks.Builder lostAssignmentsToReassign = new ConnectorsAndTasks.Builder();
        assignor.handleLostAssignments(lostAssignments, lostAssignmentsToReassign,
                new ArrayList<>(configuredAssignment.values()));

        assertTrue(lostAssignmentsToReassign.build().connectors().containsAll(lostAssignments.connectors()),
            "Wrong assignment of lost connectors");
        assertTrue(lostAssignmentsToReassign.build().tasks().containsAll(lostAssignments.tasks()),
            "Wrong assignment of lost tasks");
        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);
    }

    @Test
    public void testLostAssignmentHandlingWhenScheduledDelayIsDisabled() {
        // Customize assignor for this test case
        rebalanceDelay = 0;
        time = new MockTime();
        initAssignor();

        assertTrue(assignor.candidateWorkersForReassignment.isEmpty());
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        Map<String, WorkerLoad> configuredAssignment = new HashMap<>();
        configuredAssignment.put("worker0", workerLoad("worker0", 0, 2, 0, 4));
        configuredAssignment.put("worker1", workerLoad("worker1", 2, 2, 4, 4));
        configuredAssignment.put("worker2", workerLoad("worker2", 4, 2, 8, 4));

        // No lost assignments
        assignor.handleLostAssignments(new ConnectorsAndTasks.Builder().build(),
                new ConnectorsAndTasks.Builder(),
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);

        assignor.previousMembers = new HashSet<>(configuredAssignment.keySet());

        String veryFlakyWorker = "worker1";
        WorkerLoad lostLoad = configuredAssignment.remove(veryFlakyWorker);
        ConnectorsAndTasks lostAssignments = new ConnectorsAndTasks.Builder()
                .with(lostLoad.connectors(), lostLoad.tasks()).build();

        // Lost assignments detected - Immediately reassigned
        ConnectorsAndTasks.Builder lostAssignmentsToReassign = new ConnectorsAndTasks.Builder();
        assignor.handleLostAssignments(lostAssignments, lostAssignmentsToReassign,
                new ArrayList<>(configuredAssignment.values()));

        assertEquals(Set.of(),
            assignor.candidateWorkersForReassignment,
            "Wrong set of workers for reassignments");
        assertEquals(0, assignor.scheduledRebalance);
        assertEquals(0, assignor.delay);
        assertEquals(lostAssignments.connectors(),
            lostAssignmentsToReassign.build().connectors(), "Wrong assignment of lost connectors");
        assertEquals(lostAssignments.tasks(),
            lostAssignmentsToReassign.build().tasks(), "Wrong assignment of lost tasks");
    }

    @Test
    public void testScheduledDelayIsDisabled() {
        // Customize assignor for this test case
        rebalanceDelay = 0;
        time = new MockTime();
        initAssignor();

        // First assignment with 2 workers and 2 connectors configured but not yet assigned
        addNewEmptyWorkers("worker2");
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Second assignment with only one worker remaining in the group. The worker that left the
        // group was a follower. Re-assignments take place immediately
        removeWorkers("worker2");
        performStandardRebalance();
        assertDelay(rebalanceDelay);
        assertWorkers("worker1");
        assertConnectorAllocations(2);
        assertTaskAllocations(8);
        assertBalancedAndCompleteAllocation();
    }

    @Test
    public void testTaskAssignmentWhenTasksDuplicatedInWorkerAssignment() {
        // First assignment with 1 worker and 2 connectors configured but not yet assigned
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1");
        assertConnectorAllocations(2);
        assertTaskAllocations(8);
        assertBalancedAndCompleteAllocation();

        // Second assignment with a second worker with duplicate assignment joining and all connectors running on previous worker
        addNewWorker("worker2", newConnectors(1, 2), newTasks("connector1", 0, 4));
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 1);
        assertTaskAllocations(0, 4);

        // Third assignment after revocations
        performStandardRebalance();
        assertDelay(0);
        assertNoRevocations();
        assertConnectorAllocations(1, 1);
        assertTaskAllocations(4, 4);
        assertBalancedAndCompleteAllocation();

        // Fifth rebalance should not change assignments
        performStandardRebalance();
        assertDelay(0);
        assertEmptyAssignment();
    }

    @Test
    public void testDuplicatedAssignmentHandleWhenTheDuplicatedAssignmentsDeleted() {
        // Customize assignor for this test case
        time = new MockTime();
        initAssignor();

        // First assignment with 1 worker and 2 connectors configured but not yet assigned
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1");
        assertConnectorAllocations(2);
        assertTaskAllocations(8);
        assertBalancedAndCompleteAllocation();

        // Delete connector1
        removeConnector("connector1");

        // Second assignment with a second worker with duplicate assignment joining and the duplicated assignment is deleted at the same time
        addNewWorker("worker2", newConnectors(1, 2), newTasks("connector1", 0, 4));
        performStandardRebalance();
        assertDelay(0);
        assertWorkers("worker1", "worker2");
        assertConnectorAllocations(0, 1);
        assertTaskAllocations(0, 2);

        // Third assignment after revocations
        performStandardRebalance();
        assertDelay(0);
        assertNoRevocations();
        assertConnectorAllocations(0, 1);
        assertTaskAllocations(2, 2);
        assertBalancedAndCompleteAllocation();

        time.sleep(assignor.delay);
        // Fifth rebalance after delay should not change assignments
        performStandardRebalance();
        assertDelay(0);
        assertEmptyAssignment();
    }

    @Test
    public void testLeaderStateUpdated() {
        // Sanity test to make sure that the coordinator's leader state is actually updated after a rebalance
        connectors.clear();
        String leader = "followMe";
        Map<String, ExtendedWorkerState> workerStates = new HashMap<>();
        workerStates.put(leader, new ExtendedWorkerState("followMe:618", CONFIG_OFFSET, ExtendedAssignment.empty()));
        WorkerCoordinator coordinator = mock(WorkerCoordinator.class);
        when(coordinator.configSnapshot()).thenReturn(configState());
        assignor.performTaskAssignment(
                leader,
                CONFIG_OFFSET,
                workerStates,
                coordinator,
                IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2
        );
        verify(coordinator).leaderState(notNull());
    }

    @Test
    public void testProtocolV1() {
        // Sanity test to make sure that the right protocol is chosen during the assignment
        connectors.clear();
        String leader = "followMe";
        List<JoinGroupResponseData.JoinGroupResponseMember> memberMetadata = new ArrayList<>();
        ExtendedAssignment leaderAssignment = new ExtendedAssignment(
                IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1,
                ConnectProtocol.Assignment.NO_ERROR,
                leader,
                "followMe:618",
                CONFIG_OFFSET,
                Set.of(),
                Set.of(),
                Set.of(),
                Set.of(),
                0
        );
        ExtendedWorkerState leaderState = new ExtendedWorkerState("followMe:618", CONFIG_OFFSET, leaderAssignment);
        JoinGroupResponseData.JoinGroupResponseMember leaderMetadata = new JoinGroupResponseData.JoinGroupResponseMember()
                .setMemberId(leader)
                .setMetadata(IncrementalCooperativeConnectProtocol.serializeMetadata(leaderState, false).array());
        memberMetadata.add(leaderMetadata);
        WorkerCoordinator coordinator = mock(WorkerCoordinator.class);
        when(coordinator.configSnapshot()).thenReturn(configState());
        Map<String, ByteBuffer> serializedAssignments = assignor.performAssignment(
                leader,
                ConnectProtocolCompatibility.COMPATIBLE,
                memberMetadata,
                coordinator
        );
        serializedAssignments.forEach((worker, serializedAssignment) -> {
            ExtendedAssignment assignment = IncrementalCooperativeConnectProtocol.deserializeAssignment(serializedAssignment);
            assertEquals(
                    IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V1,
                    assignment.version(),
                    "Incorrect protocol version in assignment for worker " + worker
            );
        });
    }

    @Test
    public void testProtocolV2() {
        // Sanity test to make sure that the right protocol is chosen during the assignment
        connectors.clear();
        String leader = "followMe";
        List<JoinGroupResponseData.JoinGroupResponseMember> memberMetadata = new ArrayList<>();
        ExtendedAssignment leaderAssignment = new ExtendedAssignment(
                IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2,
                ConnectProtocol.Assignment.NO_ERROR,
                leader,
                "followMe:618",
                CONFIG_OFFSET,
                Set.of(),
                Set.of(),
                Set.of(),
                Set.of(),
                0
        );
        ExtendedWorkerState leaderState = new ExtendedWorkerState("followMe:618", CONFIG_OFFSET, leaderAssignment);
        JoinGroupResponseData.JoinGroupResponseMember leaderMetadata = new JoinGroupResponseData.JoinGroupResponseMember()
                .setMemberId(leader)
                .setMetadata(IncrementalCooperativeConnectProtocol.serializeMetadata(leaderState, true).array());
        memberMetadata.add(leaderMetadata);
        WorkerCoordinator coordinator = mock(WorkerCoordinator.class);
        when(coordinator.configSnapshot()).thenReturn(configState());
        Map<String, ByteBuffer> serializedAssignments = assignor.performAssignment(
                leader,
                ConnectProtocolCompatibility.SESSIONED,
                memberMetadata,
                coordinator
        );
        serializedAssignments.forEach((worker, serializedAssignment) -> {
            ExtendedAssignment assignment = IncrementalCooperativeConnectProtocol.deserializeAssignment(serializedAssignment);
            assertEquals(
                    IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2,
                    assignment.version(),
                    "Incorrect protocol version in assignment for worker " + worker
            );
        });
    }

    private void performStandardRebalance() {
        performRebalance(false, false);
    }

    private void performFailedRebalance() {
        performRebalance(true, false);
    }

    private void performRebalanceWithMismatchedGeneration() {
        performRebalance(false, true);
    }

    private void performRebalance(boolean assignmentFailure, boolean generationMismatch) {
        generationId++;
        int lastCompletedGenerationId = generationMismatch ? generationId - 2 : generationId - 1;
        try {
            Map<String, ConnectorsAndTasks> memberAssignmentsCopy = memberAssignments.entrySet().stream()
                    .collect(Collectors.toMap(
                            Map.Entry::getKey,
                            e -> new ConnectorsAndTasks.Builder().with(e.getValue().connectors(), e.getValue().tasks()).build()
                    ));
            returnedAssignments = assignor.performTaskAssignment(configState(), lastCompletedGenerationId, generationId, memberAssignmentsCopy);
        } catch (RuntimeException e) {
            if (assignmentFailure) {
                RequestFuture.failure(e);
            } else {
                throw e;
            }
        }
        assertNoRedundantAssignments();
        if (!assignmentFailure) {
            applyAssignments();
        }
    }

    private void addNewEmptyWorkers(String... workers) {
        for (String worker : workers) {
            addNewWorker(worker, List.of(), List.of());
        }
    }

    private void addNewWorker(String worker, List<String> connectors, List<ConnectorTaskId> tasks) {
        ConnectorsAndTasks assignment = new ConnectorsAndTasks.Builder().with(connectors, tasks).build();
        assertNull(
                memberAssignments.put(worker, assignment),
                "Worker " + worker + " already exists"
        );
    }

    private void removeWorkers(String... workers) {
        for (String worker : workers) {
            assertNotNull(
                    memberAssignments.remove(worker),
                    "Worker " + worker + " does not exist"
            );
        }
    }

    private static WorkerLoad emptyWorkerLoad(String worker) {
        return new WorkerLoad.Builder(worker).build();
    }

    private static WorkerLoad workerLoad(String worker, int connectorStart, int connectorNum,
                                  int taskStart, int taskNum) {
        return new WorkerLoad.Builder(worker).with(
                newConnectors(connectorStart, connectorStart + connectorNum),
                newTasks(taskStart, taskStart + taskNum)).build();
    }

    private static List<String> newConnectors(int start, int end) {
        return IntStream.range(start, end)
                .mapToObj(i -> "connector" + i)
                .collect(Collectors.toList());
    }

    private static List<ConnectorTaskId> newTasks(int start, int end) {
        return newTasks("task", start, end);
    }

    private static List<ConnectorTaskId> newTasks(String connectorName, int start, int end) {
        return IntStream.range(start, end)
                .mapToObj(i -> new ConnectorTaskId(connectorName, i))
                .collect(Collectors.toList());
    }

    private void addNewConnector(String connector, int taskCount) {
        assertNull(
                connectors.put(connector, taskCount),
                "Connector " + connector + " already exists"
        );
    }

    private void removeConnector(String connector) {
        assertNotNull(
                connectors.remove(connector),
                "Connector " + connector + " does not exist"
        );
    }

    private ClusterConfigState configState() {
        Map<String, Integer> taskCounts = new HashMap<>(connectors);
        Map<String, Map<String, String>> connectorConfigs = transformValues(taskCounts, c -> Map.of());
        Map<String, TargetState> targetStates = transformValues(taskCounts, c -> TargetState.STARTED);
        Map<ConnectorTaskId, Map<String, String>> taskConfigs = taskCounts.entrySet().stream()
                .flatMap(e -> IntStream.range(0, e.getValue()).mapToObj(i -> new ConnectorTaskId(e.getKey(), i)))
                .collect(Collectors.toMap(
                        Function.identity(),
                        connectorTaskId -> Map.of()
                ));
        Map<String, AppliedConnectorConfig> appliedConnectorConfigs = connectorConfigs.entrySet().stream()
                .collect(Collectors.toMap(
                        Map.Entry::getKey,
                        e -> new AppliedConnectorConfig(e.getValue())
                ));
        return new ClusterConfigState(
                CONFIG_OFFSET,
                null,
                taskCounts,
                connectorConfigs,
                targetStates,
                taskConfigs,
                Map.of(),
                Map.of(),
                appliedConnectorConfigs,
                Set.of(),
                Set.of());
    }

    private void applyAssignments() {
        returnedAssignments.allWorkers().forEach(worker -> {
            ConnectorsAndTasks workerAssignment = memberAssignments.computeIfAbsent(worker, ignored -> new ConnectorsAndTasks.Builder().build());

            workerAssignment.connectors().removeAll(returnedAssignments.newlyRevokedConnectors(worker));
            workerAssignment.connectors().addAll(returnedAssignments.newlyAssignedConnectors(worker));
            workerAssignment.tasks().removeAll(returnedAssignments.newlyRevokedTasks(worker));
            workerAssignment.tasks().addAll(returnedAssignments.newlyAssignedTasks(worker));

            assertEquals(new HashSet<>(workerAssignment.connectors()),
                new HashSet<>(returnedAssignments.allAssignedConnectors().get(worker)),
                "Complete connector assignment for worker " + worker + " does not match expectations " +
                    "based on prior assignment and new revocations and assignments");
            assertEquals(new HashSet<>(workerAssignment.tasks()),
                new HashSet<>(returnedAssignments.allAssignedTasks().get(worker)),
                "Complete task assignment for worker " + worker + " does not match expectations " +
                    "based on prior assignment and new revocations and assignments");
        });
    }

    private void assertEmptyAssignment() {
        assertEquals(List.of(),
            ConnectUtils.combineCollections(returnedAssignments.newlyAssignedConnectors().values()),
            "No connectors should have been newly assigned during this round");
        assertEquals(List.of(),
            ConnectUtils.combineCollections(returnedAssignments.newlyAssignedTasks().values()),
            "No tasks should have been newly assigned during this round");
        assertEquals(List.of(),
            ConnectUtils.combineCollections(returnedAssignments.newlyRevokedConnectors().values()),
            "No connectors should have been revoked during this round");
        assertEquals(List.of(),
            ConnectUtils.combineCollections(returnedAssignments.newlyRevokedTasks().values()),
            "No tasks should have been revoked during this round");
    }

    private void assertWorkers(String... workers) {
        assertEquals(Set.of(workers), returnedAssignments.allWorkers(), "Wrong set of workers");
    }

    /**
     * Assert that the connector counts for each worker in the cluster match the expected counts.
     * For example, calling {@code assertConnectorAllocations(0, 0, 2, 3)} ensures that there are two
     * workers in the cluster that are assigned no connectors, one worker that is assigned two connectors,
     * and one worker that is assigned three connectors.
     */
    private void assertConnectorAllocations(int... connectorCounts) {
        assertAllocations("connectors", ConnectorsAndTasks::connectors, connectorCounts);
    }

    /**
     * Assert that the task counts for each worker in the cluster match the expected counts.
     * For example, calling {@code assertTaskAllocations(0, 0, 2, 3)} ensures that there are two
     * workers in the cluster that are assigned no tasks, one worker that is assigned two tasks,
     * and one worker that is assigned three tasks.
     */
    private void assertTaskAllocations(int... taskCounts) {
        assertAllocations("tasks", ConnectorsAndTasks::tasks, taskCounts);
    }

    private void assertAllocations(String allocated, Function<ConnectorsAndTasks, ? extends Collection<?>> allocation, int... rawExpectedAllocations) {
        List<Integer> expectedAllocations = IntStream.of(rawExpectedAllocations)
                .boxed()
                .sorted()
                .collect(Collectors.toList());
        List<Integer> actualAllocations = allocations(allocation);
        assertEquals(expectedAllocations,
            actualAllocations,
            "Allocation of assigned " + allocated + " across cluster does not match expected counts");
    }

    private List<Integer> allocations(Function<ConnectorsAndTasks, ? extends Collection<?>> allocation) {
        return memberAssignments.values().stream()
                .map(allocation)
                .map(Collection::size)
                .sorted()
                .collect(Collectors.toList());
    }

    private void assertNoRevocations() {
        returnedAssignments.newlyRevokedConnectors().forEach((worker, revocations) ->
                assertEquals(
                    Set.of(),
                    new HashSet<>(revocations),
                    "Expected no revocations to take place during this round, but connector revocations were issued for worker " + worker
                                )
        );
        returnedAssignments.newlyRevokedTasks().forEach((worker, revocations) ->
                assertEquals(
                    Set.of(),
                    new HashSet<>(revocations),
                    "Expected no revocations to take place during this round, but task revocations were issued for worker " + worker
                                )
        );
    }

    private void assertDelay(int expectedDelay) {
        assertEquals(
                expectedDelay,
                assignor.delay,
                "Wrong rebalance delay"
        );
    }

    /**
     * Ensure that no connectors or tasks that were already assigned during the previous round are newly assigned in this round,
     * and that each newly-assigned connector and task is only assigned to a single worker.
     */
    private void assertNoRedundantAssignments() {
        List<String> existingConnectors = ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::connectors);
        List<String> newConnectors = ConnectUtils.combineCollections(returnedAssignments.newlyAssignedConnectors().values());
        List<ConnectorTaskId> existingTasks = ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::tasks);
        List<ConnectorTaskId> newTasks = ConnectUtils.combineCollections(returnedAssignments.newlyAssignedTasks().values());

        assertNoDuplicates(
                newConnectors,
                "Connectors should be unique in assignments but duplicates were found; the set of newly-assigned connectors is " + newConnectors
        );
        assertNoDuplicates(
                newTasks,
                "Tasks should be unique in assignments but duplicates were found; the set of newly-assigned tasks is " + newTasks
        );

        existingConnectors.retainAll(newConnectors);
        assertEquals(List.of(),
            existingConnectors,
            "Found connectors in new assignment that already exist in current assignment");
        existingTasks.retainAll(newTasks);
        assertEquals(List.of(),
            existingConnectors,
            "Found tasks in new assignment that already exist in current assignment");
    }

    private void assertBalancedAndCompleteAllocation() {
        assertBalancedAllocation();
        assertCompleteAllocation();
    }

    private void assertBalancedAllocation() {
        List<Integer> connectorCounts = allocations(ConnectorsAndTasks::connectors);
        List<Integer> taskCounts = allocations(ConnectorsAndTasks::tasks);

        int minConnectors = connectorCounts.get(0);
        int maxConnectors = connectorCounts.get(connectorCounts.size() - 1);

        int minTasks = taskCounts.get(0);
        int maxTasks = taskCounts.get(taskCounts.size() - 1);

        assertTrue(maxConnectors - minConnectors <= 1,
            "Assignments are imbalanced. The spread of connectors across each worker is: " + connectorCounts);
        assertTrue(maxTasks - minTasks <= 1,
            "Assignments are imbalanced. The spread of tasks across each worker is: " + taskCounts);
    }

    private void assertCompleteAllocation() {
        List<String> allAssignedConnectors = ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::connectors);
        assertEquals(connectors.keySet(),
            new HashSet<>(allAssignedConnectors),
            "The set of connectors assigned across the cluster does not match the set of connectors in the config topic");

        Map<String, List<ConnectorTaskId>> allAssignedTasks = ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::tasks)
                .stream()
                .collect(Collectors.groupingBy(ConnectorTaskId::connector, Collectors.toList()));

        connectors.forEach((connector, taskCount) -> {
            Set<ConnectorTaskId> expectedTasks = IntStream.range(0, taskCount)
                    .mapToObj(i -> new ConnectorTaskId(connector, i))
                    .collect(Collectors.toSet());
            assertEquals(expectedTasks,
                new HashSet<>(allAssignedTasks.get(connector)),
                "The set of tasks assigned across the cluster for connector " + connector + " does not match the set of tasks in the config topic");
        });
    }

    private static <T> void assertNoDuplicates(List<T> collection, String assertionMessage) {
        assertEquals(new HashSet<>(collection).size(), collection.size(), assertionMessage);
    }

}
