content: /*
   * Licensed to the Apache Software Foundation (ASF) under one or more
   * contributor license agreements. See the NOTICE file distributed with
   * this work for additional information regarding copyright ownership.
   * The ASF licenses this file to You under the Apache License, Version 2.0
   * (the "License"); you may not use this file except in compliance with
   * the License. You may obtain a copy of the License at
   *
   *    http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.kafka.coordinator.group.streams;
  
  import org.apache.kafka.common.Uuid;
  import org.apache.kafka.coordinator.common.runtime.CoordinatorRecord;
  import org.apache.kafka.coordinator.group.MetadataImageBuilder;
  import org.apache.kafka.coordinator.group.generated.StreamsGroupMemberMetadataValue;
  import org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.TaskRole;
  import org.apache.kafka.coordinator.group.streams.assignor.AssignmentMemberSpec;
  import org.apache.kafka.coordinator.group.streams.assignor.GroupAssignment;
  import org.apache.kafka.coordinator.group.streams.assignor.GroupSpecImpl;
  import org.apache.kafka.coordinator.group.streams.assignor.MemberAssignment;
  import org.apache.kafka.coordinator.group.streams.assignor.TaskAssignor;
  import org.apache.kafka.coordinator.group.streams.topics.ConfiguredSubtopology;
  import org.apache.kafka.coordinator.group.streams.topics.ConfiguredTopology;
  import org.apache.kafka.image.MetadataImage;
  
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.EnumSource;
  
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  import java.util.Set;
  import java.util.SortedMap;
  import java.util.TreeMap;
  
  import static org.apache.kafka.common.utils.Utils.mkEntry;
  import static org.apache.kafka.common.utils.Utils.mkMap;
  import static org.apache.kafka.coordinator.group.Assertions.assertUnorderedRecordsEquals;
  import static org.apache.kafka.coordinator.group.streams.StreamsCoordinatorRecordHelpers.newStreamsGroupTargetAssignmentEpochRecord;
  import static org.apache.kafka.coordinator.group.streams.StreamsCoordinatorRecordHelpers.newStreamsGroupTargetAssignmentRecord;
  import static org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.createAssignmentMemberSpec;
  import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasks;
  import static org.apache.kafka.coordinator.group.streams.TaskAssignmentTestUtil.mkTasksTuple;
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.mockito.ArgumentMatchers.any;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.times;
  import static org.mockito.Mockito.verify;
  import static org.mockito.Mockito.when;
  
  public class TargetAssignmentBuilderTest {
  
      @Test
      public void testBuildEmptyAssignmentWhenTopologyNotReady() {
          String groupId = "test-group";
          int groupEpoch = 1;
          TaskAssignor assignor = mock(TaskAssignor.class);
          ConfiguredTopology topology = mock(ConfiguredTopology.class);
          Map<String, String> assignmentConfigs = new HashMap<>();
  
          when(topology.isReady()).thenReturn(false);
  
          TargetAssignmentBuilder builder = new TargetAssignmentBuilder(groupId, groupEpoch, assignor, assignmentConfigs)
              .withTopology(topology);
  
          TargetAssignmentBuilder.TargetAssignmentResult result = builder.build();
  
          List<CoordinatorRecord> expectedRecords = List.of(
              StreamsCoordinatorRecordHelpers.newStreamsGroupTargetAssignmentEpochRecord(groupId, groupEpoch)
          );
  
          assertEquals(expectedRecords, result.records());
          assertEquals(Map.of(), result.targetAssignment());
      }
  
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testCreateAssignmentMemberSpec(TaskRole taskRole) {
          String fooSubtopologyId = Uuid.randomUuid().toString();
          String barSubtopologyId = Uuid.randomUuid().toString();
  
          final Map<String, String> clientTags = mkMap(mkEntry("tag1", "value1"), mkEntry("tag2", "value2"));
          StreamsGroupMember member = new StreamsGroupMember.Builder("member-id")
              .setRackId("rackId")
              .setInstanceId("instanceId")
              .setProcessId("processId")
              .setClientTags(clientTags)
              .build();
  
          TasksTuple assignment = mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          );
  
          AssignmentMemberSpec assignmentMemberSpec = createAssignmentMemberSpec(
              member,
              assignment
          );
  
          assertEquals(new AssignmentMemberSpec(
              Optional.of("instanceId"),
              Optional.of("rackId"),
              assignment.activeTasks(),
              assignment.standbyTasks(),
              assignment.warmupTasks(),
              "processId",
              clientTags,
              Map.of(),
              Map.of()
          ), assignmentMemberSpec);
      }
  
      @Test
      public void testEmpty() {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
          assertEquals(List.of(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          )), result.records());
          assertEquals(Map.of(), result.targetAssignment());
      }
  
      
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testAssignmentHasNotChanged(TaskRole taskRole) {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          String fooSubtopologyId = context.addSubtopologyWithSingleSourceTopic("foo", 6);
          String barSubtopologyId = context.addSubtopologyWithSingleSourceTopic("bar", 6);
  
          context.addGroupMember("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
  
          context.addGroupMember("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          context.prepareMemberAssignment("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
  
          context.prepareMemberAssignment("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
  
          assertEquals(List.of(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          )), result.records());
  
          Map<String, TasksTuple> expectedAssignment = new HashMap<>();
          expectedAssignment.put("member-1", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
          expectedAssignment.put("member-2", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          assertEquals(expectedAssignment, result.targetAssignment());
      }
  
      
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testAssignmentSwapped(TaskRole taskRole) {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          String fooSubtopologyId = context.addSubtopologyWithSingleSourceTopic("foo", 6);
          String barSubtopologyId = context.addSubtopologyWithSingleSourceTopic("bar", 6);
  
          context.addGroupMember("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
  
          context.addGroupMember("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          context.prepareMemberAssignment("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
  
          context.prepareMemberAssignment("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
  
          assertEquals(3, result.records().size());
  
          assertUnorderedRecordsEquals(List.of(List.of(
              newStreamsGroupTargetAssignmentRecord("my-group", "member-1", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 4, 5, 6),
                  mkTasks(barSubtopologyId, 4, 5, 6)
              )),
              newStreamsGroupTargetAssignmentRecord("my-group", "member-2", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 1, 2, 3),
                  mkTasks(barSubtopologyId, 1, 2, 3)
              ))
          )), result.records().subList(0, 2));
  
          assertEquals(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          ), result.records().get(2));
  
          Map<String, TasksTuple> expectedAssignment = new HashMap<>();
          expectedAssignment.put("member-2", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
          expectedAssignment.put("member-1", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          assertEquals(expectedAssignment, result.targetAssignment());
      }
  
      
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testNewMember(TaskRole taskRole) {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          String fooSubtopologyId = context.addSubtopologyWithSingleSourceTopic("foo", 6);
          String barSubtopologyId = context.addSubtopologyWithSingleSourceTopic("bar", 6);
  
          context.addGroupMember("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
  
          context.addGroupMember("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          context.updateMemberMetadata("member-3");
  
          context.prepareMemberAssignment("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.prepareMemberAssignment("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          context.prepareMemberAssignment("member-3", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
  
          assertEquals(4, result.records().size());
  
          assertUnorderedRecordsEquals(List.of(List.of(
              newStreamsGroupTargetAssignmentRecord("my-group", "member-1", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 1, 2),
                  mkTasks(barSubtopologyId, 1, 2)
              )),
              newStreamsGroupTargetAssignmentRecord("my-group", "member-2", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 3, 4),
                  mkTasks(barSubtopologyId, 3, 4)
              )),
              newStreamsGroupTargetAssignmentRecord("my-group", "member-3", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 5, 6),
                  mkTasks(barSubtopologyId, 5, 6)
              ))
          )), result.records().subList(0, 3));
  
          assertEquals(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          ), result.records().get(3));
  
          Map<String, TasksTuple> expectedAssignment = new HashMap<>();
          expectedAssignment.put("member-1", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
          expectedAssignment.put("member-2", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
          expectedAssignment.put("member-3", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          assertEquals(expectedAssignment, result.targetAssignment());
      }
  
      
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testUpdateMember(TaskRole taskRole) {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          String fooSubtopologyId = context.addSubtopologyWithSingleSourceTopic("foo", 6);
          String barSubtopologyId = context.addSubtopologyWithSingleSourceTopic("bar", 6);
  
          context.addGroupMember("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.addGroupMember("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          context.addGroupMember("member-3", mkTasksTuple(taskRole,
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          context.updateMemberMetadata(
              "member-3",
              Optional.of("instance-id-3"),
              Optional.of("rack-0")
          );
  
          context.prepareMemberAssignment("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.prepareMemberAssignment("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          context.prepareMemberAssignment("member-3", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
  
          assertEquals(4, result.records().size());
  
          assertUnorderedRecordsEquals(List.of(List.of(
              newStreamsGroupTargetAssignmentRecord("my-group", "member-1", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 1, 2),
                  mkTasks(barSubtopologyId, 1, 2)
              )),
              newStreamsGroupTargetAssignmentRecord("my-group", "member-2", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 3, 4),
                  mkTasks(barSubtopologyId, 3, 4)
              )),
              newStreamsGroupTargetAssignmentRecord("my-group", "member-3", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 5, 6),
                  mkTasks(barSubtopologyId, 5, 6)
              ))
          )), result.records().subList(0, 3));
  
          assertEquals(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          ), result.records().get(3));
  
          Map<String, TasksTuple> expectedAssignment = new HashMap<>();
          expectedAssignment.put("member-1", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
          expectedAssignment.put("member-2", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
          expectedAssignment.put("member-3", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          assertEquals(expectedAssignment, result.targetAssignment());
      }
  
      
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testPartialAssignmentUpdate(TaskRole taskRole) {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          String fooSubtopologyId = context.addSubtopologyWithSingleSourceTopic("foo", 6);
          String barSubtopologyId = context.addSubtopologyWithSingleSourceTopic("bar", 6);
  
          context.addGroupMember("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.addGroupMember("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          context.addGroupMember("member-3", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          context.prepareMemberAssignment("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.prepareMemberAssignment("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 3, 4, 5),
              mkTasks(barSubtopologyId, 3, 4, 5)
          ));
  
          context.prepareMemberAssignment("member-3", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 6),
              mkTasks(barSubtopologyId, 6)
          ));
  
          org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
  
          assertEquals(3, result.records().size());
  
          // Member 1 has no record because its assignment did not change.
          assertUnorderedRecordsEquals(List.of(List.of(
              newStreamsGroupTargetAssignmentRecord("my-group", "member-2", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 3, 4, 5),
                  mkTasks(barSubtopologyId, 3, 4, 5)
              )),
              newStreamsGroupTargetAssignmentRecord("my-group", "member-3", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 6),
                  mkTasks(barSubtopologyId, 6)
              ))
          )), result.records().subList(0, 2));
  
          assertEquals(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          ), result.records().get(2));
  
          Map<String, TasksTuple> expectedAssignment = new HashMap<>();
          expectedAssignment.put("member-1", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
          expectedAssignment.put("member-2", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 3, 4, 5),
              mkTasks(barSubtopologyId, 3, 4, 5)
          ));
          expectedAssignment.put("member-3", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 6),
              mkTasks(barSubtopologyId, 6)
          ));
  
          assertEquals(expectedAssignment, result.targetAssignment());
      }
  
      
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testDeleteMember(TaskRole taskRole) {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          String fooSubtopologyId = context.addSubtopologyWithSingleSourceTopic("foo", 6);
          String barSubtopologyId = context.addSubtopologyWithSingleSourceTopic("bar", 6);
  
          context.addGroupMember("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.addGroupMember("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          context.addGroupMember("member-3", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          context.removeMember("member-3");
  
          context.prepareMemberAssignment("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
  
          context.prepareMemberAssignment("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
  
          assertEquals(3, result.records().size());
  
          assertUnorderedRecordsEquals(List.of(List.of(
              newStreamsGroupTargetAssignmentRecord("my-group", "member-1", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 1, 2, 3),
                  mkTasks(barSubtopologyId, 1, 2, 3)
              )),
              newStreamsGroupTargetAssignmentRecord("my-group", "member-2", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 4, 5, 6),
                  mkTasks(barSubtopologyId, 4, 5, 6)
              ))
          )), result.records().subList(0, 2));
  
          assertEquals(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          ), result.records().get(2));
  
          Map<String, TasksTuple> expectedAssignment = new HashMap<>();
          expectedAssignment.put("member-1", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 1, 2, 3),
              mkTasks(barSubtopologyId, 1, 2, 3)
          ));
          expectedAssignment.put("member-2", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 4, 5, 6),
              mkTasks(barSubtopologyId, 4, 5, 6)
          ));
  
          assertEquals(expectedAssignment, result.targetAssignment());
      }
  
      
      @ParameterizedTest
      @EnumSource(TaskRole.class)
      public void testReplaceStaticMember(TaskRole taskRole) {
          TargetAssignmentBuilderTestContext context = new TargetAssignmentBuilderTestContext(
              "my-group",
              20
          );
  
          String fooSubtopologyId = context.addSubtopologyWithSingleSourceTopic("foo", 6);
          String barSubtopologyId = context.addSubtopologyWithSingleSourceTopic("bar", 6);
  
          context.addGroupMember("member-1", "instance-member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.addGroupMember("member-2", "instance-member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          context.addGroupMember("member-3", "instance-member-3", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          // Static member 3 leaves
          context.removeMember("member-3");
  
          // Another static member joins with the same instance id as the departed one
          context.updateMemberMetadata("member-3-a", Optional.of("instance-member-3"),
              Optional.empty());
  
          context.prepareMemberAssignment("member-1", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
  
          context.prepareMemberAssignment("member-2", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          context.prepareMemberAssignment("member-3-a", mkTasksTuple(taskRole,
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          TargetAssignmentBuilder.TargetAssignmentResult result = context.build();
  
          assertEquals(2, result.records().size());
  
          assertUnorderedRecordsEquals(List.of(List.of(
              newStreamsGroupTargetAssignmentRecord("my-group", "member-3-a", mkTasksTuple(taskRole,
                  mkTasks(fooSubtopologyId, 5, 6),
                  mkTasks(barSubtopologyId, 5, 6)
              ))
          )), result.records().subList(0, 1));
  
          assertEquals(newStreamsGroupTargetAssignmentEpochRecord(
              "my-group",
              20
          ), result.records().get(1));
  
          Map<String, TasksTuple> expectedAssignment = new HashMap<>();
          expectedAssignment.put("member-1", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 1, 2),
              mkTasks(barSubtopologyId, 1, 2)
          ));
          expectedAssignment.put("member-2", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 3, 4),
              mkTasks(barSubtopologyId, 3, 4)
          ));
  
          expectedAssignment.put("member-3-a", mkTasksTuple(taskRole, 
              mkTasks(fooSubtopologyId, 5, 6),
              mkTasks(barSubtopologyId, 5, 6)
          ));
  
          assertEquals(expectedAssignment, result.targetAssignment());
      }
  
      public static class TargetAssignmentBuilderTestContext {
  
          private final String groupId;
          private final int groupEpoch;
          private final TaskAssignor assignor = mock(TaskAssignor.class);
          private final SortedMap<String, ConfiguredSubtopology> subtopologies = new TreeMap<>();
          private final ConfiguredTopology topology = new ConfiguredTopology(0, 0, Optional.of(subtopologies), new HashMap<>(),
              Optional.empty());
          private final Map<String, StreamsGroupMember> members = new HashMap<>();
          private final Map<String, org.apache.kafka.coordinator.group.streams.TopicMetadata> subscriptionMetadata = new HashMap<>();
          private final Map<String, StreamsGroupMember> updatedMembers = new HashMap<>();
          private final Map<String, TasksTuple> targetAssignment = new HashMap<>();
          private final Map<String, MemberAssignment> memberAssignments = new HashMap<>();
          private final Map<String, String> staticMembers = new HashMap<>();
          private MetadataImageBuilder topicsImageBuilder = new MetadataImageBuilder();
  
          public TargetAssignmentBuilderTestContext(
              String groupId,
              int groupEpoch
          ) {
              this.groupId = groupId;
              this.groupEpoch = groupEpoch;
          }
  
          public void addGroupMember(
              String memberId,
              TasksTuple targetTasks
          ) {
              addGroupMember(memberId, null, targetTasks);
          }
  
          private void addGroupMember(
              String memberId,
              String instanceId,
              TasksTuple targetTasks
          ) {
              StreamsGroupMember.Builder memberBuilder = new StreamsGroupMember.Builder(memberId);
              memberBuilder.setProcessId("processId");
              memberBuilder.setClientTags(Map.of());
              memberBuilder.setUserEndpoint(new StreamsGroupMemberMetadataValue.Endpoint().setHost("host").setPort(9090));
  
              if (instanceId != null) {
                  memberBuilder.setInstanceId(instanceId);
                  staticMembers.put(instanceId, memberId);
              } else {
                  memberBuilder.setInstanceId(null);
              }
              memberBuilder.setRackId(null);
              members.put(memberId, memberBuilder.build());
              targetAssignment.put(memberId, targetTasks);
          }
  
          public String addSubtopologyWithSingleSourceTopic(
              String topicName,
              int numTasks
          ) {
              String subtopologyId = Uuid.randomUuid().toString();
              Uuid topicId = Uuid.randomUuid();
              topicsImageBuilder = topicsImageBuilder.addTopic(topicId, topicName, numTasks);
              subtopologies.put(subtopologyId, new ConfiguredSubtopology(numTasks, Set.of(topicId.toString()), Map.of(), Set.of(), Map.of()));
  
              return subtopologyId;
          }
  
          public void updateMemberMetadata(
              String memberId
          ) {
              updateMemberMetadata(
                  memberId,
                  Optional.empty(),
                  Optional.empty()
              );
          }
  
          public void updateMemberMetadata(
              String memberId,
              Optional<String> instanceId,
              Optional<String> rackId
          ) {
              StreamsGroupMember existingMember = members.get(memberId);
              StreamsGroupMember.Builder builder;
              if (existingMember != null) {
                  builder = new StreamsGroupMember.Builder(existingMember);
              } else {
                  builder = new StreamsGroupMember.Builder(memberId);
                  builder.setProcessId("processId");
                  builder.setRackId(null);
                  builder.setInstanceId(null);
                  builder.setClientTags(Map.of());
                  builder.setUserEndpoint(new StreamsGroupMemberMetadataValue.Endpoint().setHost("host").setPort(9090));
              }
              updatedMembers.put(memberId, builder
                  .maybeUpdateInstanceId(instanceId)
                  .maybeUpdateRackId(rackId)
                  .build());
          }
  
          public void removeMember(
              String memberId
          ) {
              this.updatedMembers.put(memberId, null);
          }
  
          public void prepareMemberAssignment(
              String memberId,
              TasksTuple assignment
          ) {
              memberAssignments.put(memberId, new MemberAssignment(assignment.activeTasks(), assignment.standbyTasks(), assignment.warmupTasks()));
          }
  
          public org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult build() {
              // Prepare expected member specs.
              Map<String, AssignmentMemberSpec> memberSpecs = new HashMap<>();
  
              // All the existing members are prepared.
              members.forEach((memberId, member) ->
                  memberSpecs.put(memberId, createAssignmentMemberSpec(
                          member,
                          targetAssignment.getOrDefault(memberId, TasksTuple.EMPTY)
                      )
                  ));
  
              // All the updated are added and all the deleted
              // members are removed.
              updatedMembers.forEach((memberId, updatedMemberOrNull) -> {
                  if (updatedMemberOrNull == null) {
                      memberSpecs.remove(memberId);
                  } else {
                      TasksTuple assignment = targetAssignment.getOrDefault(memberId,
                          TasksTuple.EMPTY);
  
                      // A new static member joins and needs to replace an existing departed one.
                      if (updatedMemberOrNull.instanceId().isPresent()) {
                          String previousMemberId = staticMembers.get(updatedMemberOrNull.instanceId().get());
                          if (previousMemberId != null && !previousMemberId.equals(memberId)) {
                              assignment = targetAssignment.getOrDefault(previousMemberId,
                                  TasksTuple.EMPTY);
                          }
                      }
  
                      memberSpecs.put(memberId, createAssignmentMemberSpec(
                          updatedMemberOrNull,
                          assignment
                      ));
                  }
              });
  
              MetadataImage metadataImage = topicsImageBuilder.build();
  
              // Prepare the expected topology metadata.
              TopologyMetadata topologyMetadata = new TopologyMetadata(metadataImage, subtopologies);
  
              // Prepare the expected assignment spec.
              GroupSpecImpl groupSpec = new GroupSpecImpl(memberSpecs, new HashMap<>());
  
              // We use `any` here to always return an assignment but use `verify` later on
              // to ensure that the input was correct.
              when(assignor.assign(any(), any()))
                  .thenReturn(new GroupAssignment(memberAssignments));
  
              // Create and populate the assignment builder.
              org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder builder = new org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder(
                  groupId, groupEpoch, assignor, Map.of())
                  .withMembers(members)
                  .withTopology(topology)
                  .withStaticMembers(staticMembers)
                  .withMetadataImage(metadataImage)
                  .withTargetAssignment(targetAssignment);
  
              // Add the updated members or delete the deleted members.
              updatedMembers.forEach((memberId, updatedMemberOrNull) -> {
                  if (updatedMemberOrNull != null) {
                      builder.addOrUpdateMember(memberId, updatedMemberOrNull);
                  } else {
                      builder.removeMember(memberId);
                  }
              });
  
              // Execute the builder.
              org.apache.kafka.coordinator.group.streams.TargetAssignmentBuilder.TargetAssignmentResult result = builder.build();
  
              // Verify that the assignor was called once with the expected
              // assignment spec.
              verify(assignor, times(1))
                  .assign(groupSpec, topologyMetadata);
  
              return result;
          }
      }
  }
  