Fix bug: Multiple TokenCredential in AzureServiceBusMessagingAutoConfiguration

Author: rujche

sdk/spring/spring-cloud-azure-autoconfigure/src/main/java/com/azure/spring/cloud/autoconfigure/implementation/servicebus/AzureServiceBusMessagingAutoConfiguration.java

@@ -34,6 +34,7 @@
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.BeanUtils;
 import org.springframework.beans.factory.ObjectProvider;
+import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
@@ -44,6 +45,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Import;
 
+import static com.azure.spring.cloud.autoconfigure.implementation.context.AzureContextUtils.DEFAULT_TOKEN_CREDENTIAL_BEAN_NAME;
 import static com.azure.spring.cloud.core.implementation.util.AzurePropertiesUtils.copyAzureCommonProperties;
 
 
@@ -92,7 +94,7 @@ ServiceBusProcessorFactory defaultServiceBusNamespaceProcessorFactory(
             NamespaceProperties properties,
             ObjectProvider<PropertiesSupplier<ConsumerIdentifier, ProcessorProperties>> suppliers,
             ObjectProvider<AzureTokenCredentialResolver> tokenCredentialResolvers,
-            ObjectProvider<TokenCredential> defaultTokenCredentials,
+            @Qualifier(DEFAULT_TOKEN_CREDENTIAL_BEAN_NAME) ObjectProvider<TokenCredential> defaultTokenCredentials,
             ObjectProvider<AzureServiceClientBuilderCustomizer<ServiceBusClientBuilder>> clientBuilderCustomizers,
             ObjectProvider<AzureServiceClientBuilderCustomizer<ServiceBusClientBuilder.ServiceBusProcessorClientBuilder>> processorClientBuilderCustomizers,
             ObjectProvider<AzureServiceClientBuilderCustomizer<ServiceBusClientBuilder.ServiceBusSessionProcessorClientBuilder>> sessionProcessorClientBuilderCustomizers) {
@@ -115,7 +117,7 @@ ServiceBusConsumerFactory defaultServiceBusNamespaceConsumerFactory(
             NamespaceProperties properties,
             ObjectProvider<PropertiesSupplier<ConsumerIdentifier, ConsumerProperties>> suppliers,
             ObjectProvider<AzureTokenCredentialResolver> tokenCredentialResolvers,
-            ObjectProvider<TokenCredential> defaultTokenCredentials,
+            @Qualifier(DEFAULT_TOKEN_CREDENTIAL_BEAN_NAME) ObjectProvider<TokenCredential> defaultTokenCredentials,
             ObjectProvider<AzureServiceClientBuilderCustomizer<ServiceBusClientBuilder>> customizers,
             ObjectProvider<AzureServiceClientBuilderCustomizer<ServiceBusClientBuilder.ServiceBusSessionReceiverClientBuilder>> sessionReceiverCustomizers) {
             DefaultServiceBusNamespaceConsumerFactory factory = new DefaultServiceBusNamespaceConsumerFactory(properties, suppliers.getIfAvailable());
@@ -136,7 +138,7 @@ ServiceBusProducerFactory defaultServiceBusNamespaceProducerFactory(
             NamespaceProperties properties,
             ObjectProvider<PropertiesSupplier<String, ProducerProperties>> suppliers,
             ObjectProvider<AzureTokenCredentialResolver> tokenCredentialResolvers,
-            ObjectProvider<TokenCredential> defaultTokenCredentials,
+            @Qualifier(DEFAULT_TOKEN_CREDENTIAL_BEAN_NAME) ObjectProvider<TokenCredential> defaultTokenCredentials,
             ObjectProvider<AzureServiceClientBuilderCustomizer<ServiceBusClientBuilder>> clientBuilderCustomizers,
             ObjectProvider<AzureServiceClientBuilderCustomizer<ServiceBusClientBuilder.ServiceBusSenderClientBuilder>> senderClientBuilderCustomizers) {
             DefaultServiceBusNamespaceProducerFactory factory = new DefaultServiceBusNamespaceProducerFactory(properties, suppliers.getIfAvailable());

sdk/spring/spring-cloud-azure-autoconfigure/src/test/java/com/azure/spring/cloud/autoconfigure/implementation/servicebus/AzureServiceBusMessagingAutoConfigurationTests.java

@@ -3,7 +3,13 @@
 
 package com.azure.spring.cloud.autoconfigure.implementation.servicebus;
 
+import com.azure.core.credential.TokenCredential;
+import com.azure.spring.cloud.autoconfigure.implementation.context.AzureGlobalPropertiesAutoConfiguration;
+import com.azure.spring.cloud.autoconfigure.implementation.context.AzureTokenCredentialAutoConfiguration;
+import com.azure.spring.cloud.autoconfigure.implementation.servicebus.properties.AzureServiceBusProperties;
+import com.azure.spring.cloud.core.credential.AzureCredentialResolver;
 import com.azure.spring.messaging.servicebus.core.ServiceBusProcessorFactory;
+import com.azure.spring.messaging.servicebus.core.ServiceBusProducerFactory;
 import com.azure.spring.messaging.servicebus.core.ServiceBusTemplate;
 import com.azure.spring.messaging.servicebus.implementation.support.converter.ServiceBusMessageConverter;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -12,10 +18,15 @@
 import org.springframework.boot.autoconfigure.jackson.JacksonAutoConfiguration;
 import org.springframework.boot.test.context.FilteredClassLoader;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+import java.lang.reflect.Field;
 
 import static com.azure.spring.cloud.autoconfigure.implementation.util.TestServiceBusUtils.CONNECTION_STRING_FORMAT;
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
+import static org.mockito.Mockito.mock;
 
 class AzureServiceBusMessagingAutoConfigurationTests {
 
@@ -117,7 +128,7 @@ void withUserProvidedObjectMapper() {
             .withPropertyValues("spring.cloud.azure.servicebus.connection-string=" + String.format(CONNECTION_STRING_FORMAT, "test-namespace"),
                 "spring.cloud.azure.message-converter.isolated-object-mapper=false")
             .withUserConfiguration(AzureServiceBusPropertiesTestConfiguration.class)
-            .withBean("userObjectMapper", ObjectMapper.class, () -> new ObjectMapper())
+            .withBean("userObjectMapper", ObjectMapper.class, ObjectMapper::new)
             .withConfiguration(AutoConfigurations.of(JacksonAutoConfiguration.class))
             .run(context -> {
                 assertThat(context).hasBean("userObjectMapper");
@@ -126,4 +137,63 @@ void withUserProvidedObjectMapper() {
             });
     }
 
+    @Test
+    void testCustomTokenCredentialConfiguration() {
+        this.contextRunner
+            .withConfiguration(AutoConfigurations.of(CustomTokenCredentialConfiguration.class,
+                AzureTokenCredentialAutoConfiguration.class,
+                AzureGlobalPropertiesAutoConfiguration.class))
+            .withBean(ServiceBusMessageConverter.class, () -> mock(ServiceBusMessageConverter.class))
+            .withPropertyValues(
+                "spring.cloud.azure.servicebus.connection-string=" + String.format(CONNECTION_STRING_FORMAT, "test-namespace"),
+                "spring.cloud.azure.servicebus.credential.token-credential-bean-name=customTokenCredential"
+            )
+            .withUserConfiguration(AzureServiceBusPropertiesTestConfiguration.class)
+            .run(context -> {
+
+                // Verify that the properties contain the correct credential bean name
+                AzureServiceBusProperties serviceBusProperties = context.getBean(AzureServiceBusProperties.class);
+                assertThat(serviceBusProperties).isNotNull();
+                assertThat(serviceBusProperties.getCredential()).isNotNull();
+                assertThat(serviceBusProperties.getCredential().getTokenCredentialBeanName())
+                    .as("The token-credential-bean-name property should be set to customTokenCredential")
+                    .isEqualTo("customTokenCredential");
+
+                // Verify that the custom token credential bean exists
+                assertThat(context).hasBean("customTokenCredential");
+                TokenCredential customCredential = context.getBean("customTokenCredential", TokenCredential.class);
+                assertThat(customCredential).isNotNull();
+
+                // Verify the ServiceBusProducerFactory has the tokenCredentialResolver configured
+                assertThat(context).hasSingleBean(ServiceBusProducerFactory.class);
+                ServiceBusProducerFactory producerFactory = context.getBean(ServiceBusProducerFactory.class);
+                assertThat(producerFactory).isNotNull();
+
+                // Verify tokenCredentialResolver resolves to the custom credential
+                Field tokenCredentialResolverField =
+                    producerFactory.getClass().getDeclaredField("tokenCredentialResolver");
+                tokenCredentialResolverField.setAccessible(true);
+                Object tokenCredentialResolver = tokenCredentialResolverField.get(producerFactory);
+                assertThat(tokenCredentialResolver)
+                    .as("TokenCredentialResolver should be configured").isNotNull();
+
+                // Cast to AzureCredentialResolver and invoke resolve() to verify it returns customTokenCredential
+                @SuppressWarnings("unchecked")
+                AzureCredentialResolver<TokenCredential> resolver =
+                    (AzureCredentialResolver<TokenCredential>) tokenCredentialResolver;
+                TokenCredential resolvedCredential = resolver.resolve(serviceBusProperties);
+                assertThat(resolvedCredential)
+                    .as("The resolved credential should be the customTokenCredential bean")
+                    .isSameAs(customCredential);
+            });
+    }
+
+    @Configuration
+    public static class CustomTokenCredentialConfiguration {
+        @Bean
+        public TokenCredential customTokenCredential() {
+            return mock(TokenCredential.class);
+        }
+    }
+
 }