Repository: apache/kafka
File: server/src/test/java/org/apache/kafka/api/EndToEndClusterIdTest.java (Java)

Code before:
    public static class MockBrokerMetricsReporter extends MockMetricsReporter implements ClusterResourceListener {
        public static final AtomicReference<ClusterResource> CLUSTER_META = new AtomicReference<>();

        @Override
        public void onUpdate(ClusterResource clusterMetadata) {
            CLUSTER_META.set(clusterMetadata);
        }
    }

    @ClusterTest
    public void testEndToEndWithClassicProtocol(ClusterInstance clusterInstance) throws Exception {
        GroupProtocol groupProtocol = GroupProtocol.CLASSIC;
        MockConsumerInterceptor.resetCounters();
        MockProducerInterceptor.resetCounters();

        // Should have cluster id at broker metrics reporter after startup
        Assertions.assertNotNull(MockBrokerMetricsReporter.CLUSTER_META);
        isValidClusterId(MockBrokerMetricsReporter.CLUSTER_META.get().clusterId());


        // --- PRODUCER CONFIGURATION ---
        try (var producer = clusterInstance.<byte[], byte[]>producer(PRODUCER_CONFIG)) {
            // Send one record and make sure clusterId is set after send and before onAcknowledgement
            sendRecords(producer, NUM_RECORDS, TP);
        }

        Assertions.assertNotEquals(MockProducerInterceptor.CLUSTER_ID_BEFORE_ON_ACKNOWLEDGEMENT.get(), MockProducerInterceptor.NO_CLUSTER_ID);
        Assertions.assertNotNull(MockProducerInterceptor.CLUSTER_META);
        Assertions.assertEquals(
            MockProducerInterceptor.CLUSTER_ID_BEFORE_ON_ACKNOWLEDGEMENT.get().clusterId(),
            MockProducerInterceptor.CLUSTER_META.get().clusterId()
        );
        isValidClusterId(MockProducerInterceptor.CLUSTER_META.get().clusterId());

        // Ensure the serializer sees Cluster ID before serialization
        Assertions.assertNotEquals(MockSerializer.CLUSTER_ID_BEFORE_SERIALIZE.get(), MockSerializer.NO_CLUSTER_ID);
        Assertions.assertNotNull(MockSerializer.CLUSTER_META);
        isValidClusterId(MockSerializer.CLUSTER_META.get().clusterId());

        // Producer metric reporter receives cluster id
        Assertions.assertNotNull(MockProducerMetricsReporter.CLUSTER_META);
        isValidClusterId(MockProducerMetricsReporter.CLUSTER_META.get().clusterId());

        // --- CONSUMER CONFIGURATION ---
        Map<String, Object> consumerConfigWithGroupProtocol = new HashMap<>(CONSUMER_CONFIG);
        consumerConfigWithGroupProtocol.put(ConsumerConfig.GROUP_PROTOCOL_CONFIG, groupProtocol.name());
        try (var consumer = clusterInstance.<byte[], byte[]>consumer(consumerConfigWithGroupProtocol)) {
            consumer.assign(Collections.singletonList(TP));
            consumer.seek(TP, 0);
            // Consume one record
            consumeRecords(consumer, NUM_RECORDS);

Review comment (refactor):
Could you extract the same code into a test method?

FYI: https://github.com/apache/kafka/blob/trunk/clients/clients-integration-tests/src/test/java/org/apache/kafka/clients/consumer/PlaintextConsumerCallbackTest.java

Code after:

            String id = (String) configs.get(ServerConfigs.BROKER_ID_CONFIG);
            controllerId = roles.contains("controller") ? id : null;
            brokerId    = roles.contains("broker")    ? id : null;
        }

        @Override
        public void onUpdate(ClusterResource clusterMetadata) {
            if (clientId != null) CLUSTER_RESOURCE_MAP.put(clientId, clusterMetadata);
            if (brokerId != null) CLUSTER_RESOURCE_MAP.put(brokerId, clusterMetadata);
            if (controllerId != null) CLUSTER_RESOURCE_MAP.put(controllerId, clusterMetadata);
        }
    }

    @ClusterTest
    public void testEndToEndWithClassicProtocol() throws Exception {
        testEndToEnd(GroupProtocol.CLASSIC);
    }

    @ClusterTest
    public void testEndToEndWithConsumerProtocol() throws Exception {
        testEndToEnd(GroupProtocol.CONSUMER);
    }

    public void testEndToEnd(GroupProtocol groupProtocol) throws Exception {
        MockConsumerInterceptor.resetCounters();
        MockProducerInterceptor.resetCounters();

        ClusterResource brokerClusterResource = MockCommonMetricsReporter.CLUSTER_RESOURCE_MAP.get(clusterBrokerId);
        assertNotNull(brokerClusterResource);
        isValidClusterId(brokerClusterResource.clusterId());
        ClusterResource controllerClusterResource = MockCommonMetricsReporter.CLUSTER_RESOURCE_MAP.get(controllerId);
        assertNotNull(controllerClusterResource);
        isValidClusterId(controllerClusterResource.clusterId());

        Map<String, Object> producerConfig = Map.of(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, MockProducerInterceptor.class.getName(),
            "mock.interceptor.append", "mock",
            ProducerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockCommonMetricsReporter.class.getName(),
            ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, MockSerializer.class.getName(),
            ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, MockSerializer.class.getName(),
            ProducerConfig.CLIENT_ID_CONFIG, PRODUCER_CLIENT_ID);
        try (var producer = clusterInstance.<String, String>producer(producerConfig)) {
            // Send one record and make sure clusterId is set after sending and before onAcknowledgement
            sendRecord(producer);
        }
        assertNotEquals(MockProducerInterceptor.CLUSTER_ID_BEFORE_ON_ACKNOWLEDGEMENT.get(), MockProducerInterceptor.NO_CLUSTER_ID);
        assertNotNull(MockProducerInterceptor.CLUSTER_META.get());
        assertEquals(
            MockProducerInterceptor.CLUSTER_ID_BEFORE_ON_ACKNOWLEDGEMENT.get().clusterId(),
            MockProducerInterceptor.CLUSTER_META.get().clusterId()
        );