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

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

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

        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 sending and before onAcknowledgement
            sendRecord(producer);
        }
        assertNotEquals(MockProducerInterceptor.CLUSTER_ID_BEFORE_ON_ACKNOWLEDGEMENT.get(), MockProducerInterceptor.NO_CLUSTER_ID);
        assertNotNull(MockProducerInterceptor.CLUSTER_META);
        assertEquals(
            MockProducerInterceptor.CLUSTER_ID_BEFORE_ON_ACKNOWLEDGEMENT.get().clusterId(),
            MockProducerInterceptor.CLUSTER_META.get().clusterId()
        );
        isValidClusterId(MockProducerInterceptor.CLUSTER_META.get().clusterId());

        // Make sure the serializer sees Cluster ID before serialize method
        assertNotEquals(MockSerializer.CLUSTER_ID_BEFORE_SERIALIZE.get(), MockSerializer.NO_CLUSTER_ID);
        assertNotNull(MockSerializer.CLUSTER_META);
        isValidClusterId(MockSerializer.CLUSTER_META.get().clusterId());

        assertNotNull(MockProducerMetricsReporter.CLUSTER_META);

Review comment (suggestion):
`assertNotNull(MockBrokerMetricsReporter.CLUSTER_META.get());`

Code after:

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