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

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

    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 {
        testEndToEnd(GroupProtocol.CONSUMER, clusterInstance);
    }

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

    public void testEndToEnd(GroupProtocol groupProtocol, ClusterInstance clusterInstance) throws Exception {
        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());


        Map<String, Object> producerConfig = Map.of(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, MockProducerInterceptor.class.getName(),
            "mock.interceptor.append", "mock",
            ProducerConfig.METRIC_REPORTER_CLASSES_CONFIG, MockProducerMetricsReporter.class.getName(),
            ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, MockSerializer.class.getName(),
            ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, MockSerializer.class.getName());
        try (var producer = clusterInstance.<byte[], byte[]>producer(producerConfig)) {
            // Send one record and make sure clusterId is set after send and before onAcknowledgement
            sendRecords(producer);
        }

        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());

Review comment (suggestion):
```suggestion
    @ClusterTest
    public void testEndToEndWithClassicProtocol() throws Exception {
        testEndToEnd(GroupProtocol.CONSUMER);
    }

    @ClusterTest
    public void testEndToEndWithConsumerProtocol() throws Exception {
        testEndToEnd(GroupProtocol.CONSUMER);
    }

    public void testEndToEnd(GroupProtocol groupProtocol) throws Exception {
```

Code after:
    public void setup() throws InterruptedException {
        this.clusterInstance.createTopic(TOPIC, 2, (short) 1);
        clusterBrokerId = String.valueOf(clusterInstance.brokerIds().iterator().next());
        controllerId = String.valueOf(clusterInstance.controllerIds().iterator().next());
        MockDeserializer.resetStaticVariables();
    }

    public static class MockCommonMetricsReporter extends MockMetricsReporter implements ClusterResourceListener {
        public static final Map<String, ClusterResource> CLUSTER_RESOURCE_MAP = new ConcurrentHashMap<>();
        public String brokerId;
        public String controllerId;

        @Override
        public void configure(Map<String, ?> configs) {
            super.configure(configs);

            String roles = (String) configs.get("process.roles");
            if (roles == null) return;

            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);