Repository: apache/kafka
File: group-coordinator/src/test/java/org/apache/kafka/coordinator/group/streams/assignor/StickyTaskAssignorTest.java (Java)

Code before:
        assertNotNull(testMember);
        assertNotNull(testMember.standbyTasks());
        if (!testMember.standbyTasks().isEmpty()) {
            return testMember.standbyTasks();
        }
        return new HashMap<>();
    }

    private List<Integer> getAllStandbyTaskIds(GroupAssignment result, String... memberIds) {
        List<Integer> res = new ArrayList<>();
        for (String memberId : memberIds) {
            final MemberAssignment testMember = result.members().get(memberId);
            assertNotNull(testMember);
            assertNotNull(testMember.standbyTasks());
            if (!testMember.standbyTasks().isEmpty()) {
                for (Map.Entry<String, Set<Integer>> entry : testMember.standbyTasks().entrySet()) {
                    res.addAll(entry.getValue());
                }
            }
        }
        return res;
    }

    private List<Integer> getAllStandbyTaskIds(GroupAssignment result) {
        String[] memberIds = new String[result.members().size()];
        return getAllStandbyTaskIds(result, result.members().keySet().toArray(memberIds));
    }

    private Map<String, Set<Integer>> mergeAllActiveTasks(GroupAssignment result, String... memberIds) {
        Map<String, Set<Integer>> res = new HashMap<>();
        for (String memberId : memberIds) {
            Map<String, Set<Integer>> memberActiveTasks = getAllActiveTasks(result, memberId);
            res = Stream.of(res, memberActiveTasks)
                .flatMap(map -> map.entrySet().stream())
                .collect(Collectors.toMap(
                    Map.Entry::getKey,
                    Map.Entry::getValue,
                    (v1, v2) -> {
                        v1.addAll(v2);
                        return new HashSet<>(v1);
                    }));
        }
        return res;
    }

    private Map<String, Set<Integer>> mergeAllStandbyTasks(GroupAssignment result, String... memberIds) {
        Map<String, Set<Integer>> res = new HashMap<>();
        for (String memberId : memberIds) {
            Map<String, Set<Integer>> memberStandbyTasks = getAllStandbyTasks(result, memberId);
            res = Stream.of(res, memberStandbyTasks)
                .flatMap(map -> map.entrySet().stream())

Review comment (refactor):
The comment and assertions are duplicated from lines 1158-1167 for active tasks. The comment should be updated to reflect that this section is checking standby tasks, not active tasks.
```suggestion
        // Each client should get one standby task from each subtopology
```

Code after:
            new TopologyDescriberImpl(2, true, Arrays.asList("test-subtopology1", "test-subtopology2"))
        );

        // Each client should get one task from each subtopology
        final MemberAssignment testMember1 = result.members().get("member1");
        assertNotNull(testMember1);
        assertEquals(1, testMember1.activeTasks().get("test-subtopology1").size());
        assertEquals(1, testMember1.activeTasks().get("test-subtopology2").size());
        
        final MemberAssignment testMember2 = result.members().get("member2");
        assertNotNull(testMember2);
        assertEquals(1, testMember2.activeTasks().get("test-subtopology1").size());
        assertEquals(1, testMember2.activeTasks().get("test-subtopology2").size());
        
        // Verify all tasks are assigned exactly once
        final Set<Integer> allSubtopology1Tasks = new HashSet<>();
        allSubtopology1Tasks.addAll(testMember1.activeTasks().get("test-subtopology1"));
        allSubtopology1Tasks.addAll(testMember2.activeTasks().get("test-subtopology1"));
        assertEquals(Sets.newSet(0, 1), allSubtopology1Tasks);
        
        final Set<Integer> allSubtopology2Tasks = new HashSet<>();
        allSubtopology2Tasks.addAll(testMember1.activeTasks().get("test-subtopology2"));
        allSubtopology2Tasks.addAll(testMember2.activeTasks().get("test-subtopology2"));
        assertEquals(Sets.newSet(0, 1), allSubtopology2Tasks);

        // Each client should get one task from each subtopology
        assertNotNull(testMember1);
        assertEquals(1, testMember1.standbyTasks().get("test-subtopology1").size());
        assertEquals(1, testMember1.standbyTasks().get("test-subtopology2").size());

        assertNotNull(testMember2);
        assertEquals(1, testMember2.standbyTasks().get("test-subtopology1").size());
        assertEquals(1, testMember2.standbyTasks().get("test-subtopology2").size());
    }

    @Test
    public void shouldRangeAssignTasksWhenStartingEmpty() {
        // Two clients starting empty (no previous tasks)
        final AssignmentMemberSpec memberSpec1 = createAssignmentMemberSpec("process1");
        final AssignmentMemberSpec memberSpec2 = createAssignmentMemberSpec("process2");
        final Map<String, AssignmentMemberSpec> members = mkMap(
            mkEntry("member1", memberSpec1), mkEntry("member2", memberSpec2));

        // Two subtopologies with 2 tasks each (4 tasks total) with standby replicas enabled
        final GroupAssignment result = assignor.assign(
            new GroupSpecImpl(members, mkMap(mkEntry(NUM_STANDBY_REPLICAS_CONFIG, String.valueOf(1)))),
            new TopologyDescriberImpl(2, true, Arrays.asList("test-subtopology1", "test-subtopology2"))
        );

        // Each client should get one task from each subtopology
        final MemberAssignment testMember1 = result.members().get("member1");