Repository: apache/kafka
File: streams/src/test/java/org/apache/kafka/streams/processor/internals/RackAwarenessStreamsPartitionAssignorTest.java (Java)

Code before:
        configurationMap.put(StreamsConfig.RACK_AWARE_ASSIGNMENT_TAGS_CONFIG, String.join(",", ALL_TAG_KEYS));
        ALL_TAG_KEYS.forEach(key -> configurationMap.put(StreamsConfig.clientTagPrefix(key), "dummy"));
        return configurationMap;
    }

    // Make sure to complete setting up any mocks (such as TaskManager or AdminClient) before configuring the assignor
    private void configurePartitionAssignorWith(final Map<String, Object> props) {
        final Map<String, Object> configMap = configProps();
        configMap.putAll(props);

        streamsConfig = new StreamsConfig(configMap);
        topologyMetadata = new TopologyMetadata(builder, streamsConfig);
        partitionAssignor.configure(configMap);

        overwriteInternalTopicManagerWithMock();
    }

    // Useful for tests that don't care about the task offset sums
    private void createMockTaskManager() {
        taskManager = mock(TaskManager.class);
        when(taskManager.topologyMetadata()).thenReturn(topologyMetadata);
        when(taskManager.processId()).thenReturn(UUID_1);
        topologyMetadata.buildAndRewriteTopology();
    }

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

        final ListOffsetsResult result = mock(ListOffsetsResult.class);
        final KafkaFutureImpl<Map<TopicPartition, ListOffsetsResult.ListOffsetsResultInfo>> allFuture = new KafkaFutureImpl<>();
        allFuture.complete(changelogEndOffsets.entrySet().stream().collect(Collectors.toMap(
                Map.Entry::getKey,
                t -> {
                    final ListOffsetsResult.ListOffsetsResultInfo info = mock(ListOffsetsResult.ListOffsetsResultInfo.class);
                    when(info.offset()).thenReturn(t.getValue());
                    return info;
                }))
        );

        when(adminClient.listOffsets(any())).thenReturn(result);
        when(result.all()).thenReturn(allFuture);
    }

    private void overwriteInternalTopicManagerWithMock() {
        final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
                time,
                streamsConfig,
                mockClientSupplier.restoreConsumer,
                false

Review comment (refactor):
This is piggy-backed as part of incorporating 2): I found that these funcs are duplicated across multiple test classes from `AssignmentTestUtils.createMockAdminClientForAssignor`, so I removed them.

Code after:
    }

    // Make sure to complete setting up any mocks (such as TaskManager or AdminClient) before configuring the assignor
    private void configurePartitionAssignorWith(final Map<String, Object> props) {
        final Map<String, Object> configMap = configProps();
        configMap.putAll(props);

        streamsConfig = new StreamsConfig(configMap);
        topologyMetadata = new TopologyMetadata(builder, streamsConfig);
        partitionAssignor.configure(configMap);

        overwriteInternalTopicManagerWithMock();
    }

    // Useful for tests that don't care about the task offset sums
    private void createMockTaskManager() {
        taskManager = mock(TaskManager.class);
        when(taskManager.topologyMetadata()).thenReturn(topologyMetadata);
        when(taskManager.processId()).thenReturn(UUID_1);
        topologyMetadata.buildAndRewriteTopology();
    }

    private void overwriteInternalTopicManagerWithMock() {
        final MockInternalTopicManager mockInternalTopicManager = new MockInternalTopicManager(
                time,
                streamsConfig,
                mockClientSupplier.restoreConsumer,
                false
        );
        partitionAssignor.setInternalTopicManager(mockInternalTopicManager);
    }

    @Before
    public void setUp() {
        adminClient = createMockAdminClientForAssignor(EMPTY_CHANGELOG_END_OFFSETS);
    }

    @Test
    public void shouldDistributeWithMaximumNumberOfClientTags() {
        setupTopology(3, 2);

        createMockTaskManager();
        adminClient = createMockAdminClientForAssignor(getTopicPartitionOffsetsMap(
            Arrays.asList(APPLICATION_ID + "-store2-changelog", APPLICATION_ID + "-store3-changelog", APPLICATION_ID + "-store4-changelog"),
            Arrays.asList(3, 3, 3)));
        configurePartitionAssignorWith(Collections.singletonMap(StreamsConfig.NUM_STANDBY_REPLICAS_CONFIG, 1));

        final Map<String, String> clientTags1 = new HashMap<>();
        final Map<String, String> clientTags2 = new HashMap<>();

        for (int i = 0; i < ALL_TAG_KEYS.size(); i++) {