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

Code before:
        final Map<String, Object> configurationMap = new HashMap<>();
        configurationMap.put(StreamsConfig.APPLICATION_ID_CONFIG, APPLICATION_ID);
        configurationMap.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, USER_END_POINT);
        referenceContainer = new ReferenceContainer();
        referenceContainer.mainConsumer = consumer;
        referenceContainer.adminClient = adminClient;
        referenceContainer.taskManager = taskManager;
        referenceContainer.streamsMetadataState = streamsMetadataState;
        referenceContainer.time = time;
        configurationMap.put(InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
        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(final Set<TaskId> activeTasks) {
        createMockTaskManager(getTaskOffsetSums(activeTasks));
    }

    private void createMockTaskManager(final Map<TaskId, Long> taskOffsetSums) {
        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) {
        final ListOffsetsResult result = mock(ListOffsetsResult.class);
        final KafkaFutureImpl<Map<TopicPartition, ListOffsetsResultInfo>> allFuture = new KafkaFutureImpl<>();
        allFuture.complete(changelogEndOffsets.entrySet().stream().collect(Collectors.toMap(
            Entry::getKey,
            t -> {
                final ListOffsetsResultInfo info = mock(ListOffsetsResultInfo.class);
                when(info.offset()).thenReturn(t.getValue());
                return info;
            }))
        );

        when(adminClient.listOffsets(any())).thenReturn(result);

Review comment (refactor):
This is found that after we refactored with mockito, the passed in task set is not needed any more, so I removed them to eliminate the warnings.

Code after:
        referenceContainer.mainConsumer = consumer;
        referenceContainer.adminClient = adminClient;
        referenceContainer.taskManager = taskManager;
        referenceContainer.streamsMetadataState = streamsMetadataState;
        referenceContainer.time = time;
        configurationMap.put(InternalConfig.REFERENCE_CONTAINER_PARTITION_ASSIGNOR, referenceContainer);
        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();
    }

    private void createMockTaskManager() {
        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 shouldReturnAllActiveTasksToPreviousOwnerRegardlessOfBalanceAndTriggerRebalanceIfEndOffsetFetchFailsAndHighAvailabilityEnabled() {
        final long rebalanceInterval = 5 * 60 * 1000L;

        builder.addSource(null, "source1", null, null, null, "topic1");
        builder.addProcessor("processor1", new MockApiProcessorSupplier<>(), "source1");
        builder.addStateStore(new MockKeyValueStoreBuilder("store1", false), "processor1");
        final Set<TaskId> allTasks = mkSet(TASK_0_0, TASK_0_1, TASK_0_2);