Repository: apache/kafka
File: clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/ConsumerIntegrationTest.java (Java)

Code before:
            // partition 4 -> broker 1
            // partition 5 -> broker 0
            admin.alterPartitionReassignments(Map.of(
                new TopicPartition(topic, 0), Optional.of(new NewPartitionReassignment(List.of(2))),
                new TopicPartition(topic, 1), Optional.of(new NewPartitionReassignment(List.of(2))),
                new TopicPartition(topic, 2), Optional.of(new NewPartitionReassignment(List.of(2))),
                new TopicPartition(topic, 3), Optional.of(new NewPartitionReassignment(List.of(1))),
                new TopicPartition(topic, 4), Optional.of(new NewPartitionReassignment(List.of(1))),
                new TopicPartition(topic, 5), Optional.of(new NewPartitionReassignment(List.of(0)))
            )).all().get();
            TestUtils.waitForCondition(() -> {
                consumer0.poll(Duration.ofMillis(1000));
                consumer1.poll(Duration.ofMillis(1000));
                consumer2.poll(Duration.ofMillis(1000));
                return consumer0.assignment().equals(Set.of(new TopicPartition(topic, 5))) &&
                    consumer1.assignment().equals(Set.of(new TopicPartition(topic, 3), new TopicPartition(topic, 4))) &&
                    consumer2.assignment().equals(Set.of(new TopicPartition(topic, 0), new TopicPartition(topic, 1), new TopicPartition(topic, 2)));
            }, "Consumer with topic partition mapping should be 0 -> 5 | 1 -> 3, 4 | 2 -> 0, 1, 2");
        }
    }

    @SuppressWarnings("unchecked")
    @ClusterTest(
        brokers = 2,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_PARTITIONS_CONFIG, value = "1"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.GROUP_COORDINATOR_APPEND_LINGER_MS_CONFIG, value = "3000")
        }
    )
    public void testSingleCoordinatorOwnershipAfterPartitionReassignment(ClusterInstance clusterInstance) throws InterruptedException, ExecutionException, TimeoutException {
        try (var producer = clusterInstance.<byte[], byte[]>producer()) {
            producer.send(new ProducerRecord<>("topic", "value".getBytes(StandardCharsets.UTF_8)));
        }

        try (var admin = clusterInstance.admin()) {
            admin.createTopics(List.of(new NewTopic(Topic.GROUP_METADATA_TOPIC_NAME, Map.of(0, List.of(0))))).all().get();
        }

        try (var consumer = clusterInstance.consumer(Map.of(ConsumerConfig.GROUP_ID_CONFIG, "test-group"));
             var admin = clusterInstance.admin()) {
            consumer.subscribe(List.of("topic"));
            TestUtils.waitForCondition(() -> consumer.poll(Duration.ofMillis(100)).isEmpty(), "polling to join group");
            // append records to coordinator
            consumer.commitSync();

            // unload the coordinator by changing leader (0 -> 1)
            admin.alterPartitionReassignments(Map.of(new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, 0),
                Optional.of(new NewPartitionReassignment(List.of(1))))).all().get();
        }


Review comment (bug):
Does this test fail before the fix? I guess that it does not because the unloading of the shard actually goes through. It does not hurt to have it though.

Code after:
                new TopicPartition(topic, 1), Optional.of(new NewPartitionReassignment(List.of(2))),
                new TopicPartition(topic, 2), Optional.of(new NewPartitionReassignment(List.of(2))),
                new TopicPartition(topic, 3), Optional.of(new NewPartitionReassignment(List.of(1))),
                new TopicPartition(topic, 4), Optional.of(new NewPartitionReassignment(List.of(1))),
                new TopicPartition(topic, 5), Optional.of(new NewPartitionReassignment(List.of(0)))
            )).all().get();
            TestUtils.waitForCondition(() -> {
                consumer0.poll(Duration.ofMillis(1000));
                consumer1.poll(Duration.ofMillis(1000));
                consumer2.poll(Duration.ofMillis(1000));
                return consumer0.assignment().equals(Set.of(new TopicPartition(topic, 5))) &&
                    consumer1.assignment().equals(Set.of(new TopicPartition(topic, 3), new TopicPartition(topic, 4))) &&
                    consumer2.assignment().equals(Set.of(new TopicPartition(topic, 0), new TopicPartition(topic, 1), new TopicPartition(topic, 2)));
            }, "Consumer with topic partition mapping should be 0 -> 5 | 1 -> 3, 4 | 2 -> 0, 1, 2");
        }
    }

    @ClusterTest(
        brokers = 2,
        types = {Type.KRAFT},
        serverProperties = {
            @ClusterConfigProperty(key = GroupCoordinatorConfig.OFFSETS_TOPIC_PARTITIONS_CONFIG, value = "1"),
            @ClusterConfigProperty(key = GroupCoordinatorConfig.GROUP_COORDINATOR_APPEND_LINGER_MS_CONFIG, value = "3000")
        }
    )
    public void testSingleCoordinatorOwnershipAfterPartitionReassignment(ClusterInstance clusterInstance) throws InterruptedException, ExecutionException, TimeoutException {
        try (var producer = clusterInstance.<byte[], byte[]>producer()) {
            producer.send(new ProducerRecord<>("topic", "value".getBytes(StandardCharsets.UTF_8)));
        }

        try (var admin = clusterInstance.admin()) {
            admin.createTopics(List.of(new NewTopic(Topic.GROUP_METADATA_TOPIC_NAME, Map.of(0, List.of(0))))).all().get();
        }

        try (var consumer = clusterInstance.consumer(Map.of(ConsumerConfig.GROUP_ID_CONFIG, "test-group"));
             var admin = clusterInstance.admin()) {
            consumer.subscribe(List.of("topic"));
            TestUtils.waitForCondition(() -> consumer.poll(Duration.ofMillis(100)).isEmpty(), "polling to join group");
            // Append records to coordinator.
            consumer.commitSync();

            var broker0Metrics = clusterInstance.brokers().get(0).metrics();
            var broker1Metrics = clusterInstance.brokers().get(1).metrics();
            var activeNumPartitions = broker0Metrics.metricName(
                "num-partitions",
                GroupCoordinatorRuntimeMetrics.METRICS_GROUP,
                Map.of("state", "active")
            );

            assertEquals(1L, broker0Metrics.metric(activeNumPartitions).metricValue());
            assertEquals(0L, broker1Metrics.metric(activeNumPartitions).metricValue());