migrate azure messaging servicebus tests to use token auth.

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/IntegrationTestBase.java

@@ -6,7 +6,6 @@
 import com.azure.core.amqp.AmqpTransportType;
 import com.azure.core.amqp.ProxyAuthenticationType;
 import com.azure.core.amqp.ProxyOptions;
-import com.azure.core.amqp.implementation.ConnectionStringProperties;
 import com.azure.core.amqp.models.AmqpMessageBody;
 import com.azure.core.credential.TokenCredential;
 import com.azure.core.experimental.util.tracing.LoggingTracerProvider;
@@ -43,6 +42,7 @@
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -72,6 +72,7 @@ public abstract class IntegrationTestBase extends TestBase {
     private List<AutoCloseable> toClose = new ArrayList<>();
     private String testName;
     private final Scheduler scheduler = Schedulers.parallel();
+    private final AtomicReference<TokenCredential> credentialCached = new AtomicReference<>();
 
     protected static final byte[] CONTENTS_BYTES = "Some-contents".getBytes(StandardCharsets.UTF_8);
     protected String sessionId;
@@ -90,7 +91,7 @@ public void setupTest(TestContextManager testContextManager) {
 
         logger.info("========= SET-UP [{}] =========", testName);
 
-        assumeTrue(getTestMode() == TestMode.RECORD);
+        assertRunnable();
 
         toClose = new ArrayList<>();
         optionsWithTracing = new ClientOptions().setTracingOptions(new LoggingTracerProvider.LoggingTracingOptions());
@@ -108,39 +109,6 @@ public void teardownTest() {
         dispose();
     }
 
-    /**
-     * Gets the test mode for this API test. If AZURE_TEST_MODE equals {@link TestMode#RECORD} and Event Hubs connection
-     * string is set, then we return {@link TestMode#RECORD}. Otherwise, {@link TestMode#PLAYBACK} is returned.
-     */
-    @Override
-    public TestMode getTestMode() {
-        if (super.getTestMode() == TestMode.PLAYBACK) {
-            return TestMode.PLAYBACK;
-        }
-
-        return CoreUtils.isNullOrEmpty(getConnectionString()) ? TestMode.PLAYBACK : TestMode.RECORD;
-    }
-
-    public static String getConnectionString() {
-        return TestUtils.getConnectionString(false);
-    }
-
-    public static String getConnectionString(boolean withSas) {
-        return TestUtils.getConnectionString(withSas);
-    }
-
-    protected static ConnectionStringProperties getConnectionStringProperties() {
-        return new ConnectionStringProperties(getConnectionString(false));
-    }
-
-    protected static ConnectionStringProperties getConnectionStringProperties(boolean withSas) {
-        return new ConnectionStringProperties(getConnectionString(withSas));
-    }
-
-    public String getFullyQualifiedDomainName() {
-        return TestUtils.getFullyQualifiedDomainName();
-    }
-
     /**
      * Gets the name of the queue.
      *
@@ -204,46 +172,67 @@ public ProxyOptions getProxyConfiguration() {
     }
 
     /**
-     * Creates a new instance of {@link ServiceBusClientBuilder} with the default integration test settings and uses a
-     * connection string to authenticate.
+     * Creates a new instance of {@link ServiceBusClientBuilder} with authentication set up in {@link TestMode#LIVE} and
+     * {@link TestMode#RECORD} modes.
+     *
+     * @return the builder with authentication set up.
+     * @throws org.opentest4j.TestAbortedException if the test mode is {@link TestMode#PLAYBACK}.
      */
-    protected ServiceBusClientBuilder getBuilder() {
-        return getBuilder(false);
+    protected ServiceBusClientBuilder getAuthenticatedBuilder() {
+        final TestMode mode = super.getTestMode();
+        final ServiceBusClientBuilder builder = new ServiceBusClientBuilder();
+        if (mode == TestMode.LIVE) {
+            final String fullyQualifiedDomainName = TestUtils.getFullyQualifiedDomainName(false);
+            assumeTrue(!CoreUtils.isNullOrEmpty(fullyQualifiedDomainName), "FullyQualifiedDomainName is not set.");
+            final TokenCredential credential = TestUtils.getPipelineCredential(credentialCached);
+            return builder.credential(fullyQualifiedDomainName, credential);
+        } else if (mode == TestMode.RECORD) {
+            final String connectionString = TestUtils.getConnectionString(false);
+            if (CoreUtils.isNullOrEmpty(connectionString)) {
+                final String fullyQualifiedDomainName = TestUtils.getFullyQualifiedDomainName(false);
+                assumeTrue(!CoreUtils.isNullOrEmpty(fullyQualifiedDomainName), "FullyQualifiedDomainName is not set.");
+                final TokenCredential credential = new DefaultAzureCredentialBuilder().build();
+                return builder.credential(fullyQualifiedDomainName, credential);
+            } else {
+                return builder.connectionString(connectionString);
+            }
+        } else {
+            // Throws org.opentest4j.TestAbortedException exception.
+            assumeTrue(false, "Integration tests are not enabled in playback mode.");
+            return null;
+        }
     }
 
     /**
-     * Creates a new instance of {@link ServiceBusClientBuilder} with the default integration test settings and uses a
-     * connection string to authenticate if {@code useCredentials} is false. Otherwise, uses a service principal through
-     * {@link com.azure.identity.ClientSecretCredential}.
+     * Creates a new instance of {@link ServiceBusClientBuilder} with the default integration test settings.
      */
-    protected ServiceBusClientBuilder getBuilder(boolean useCredentials) {
-        final ServiceBusClientBuilder builder = new ServiceBusClientBuilder()
+    protected ServiceBusClientBuilder getBuilder() {
+        return getAuthenticatedBuilder()
             .proxyOptions(ProxyOptions.SYSTEM_DEFAULTS)
             .retryOptions(RETRY_OPTIONS)
             .clientOptions(optionsWithTracing)
             .transportType(AmqpTransportType.AMQP)
             .scheduler(scheduler)
             .configuration(v1OrV2(true));
+    }
 
-        logger.info("Getting Builder using credentials : [{}] ", useCredentials);
-        if (useCredentials) {
-            final String fullyQualifiedDomainName = getFullyQualifiedDomainName();
-
-            assumeTrue(fullyQualifiedDomainName != null && !fullyQualifiedDomainName.isEmpty(),
-                "AZURE_SERVICEBUS_FULLY_QUALIFIED_DOMAIN_NAME variable needs to be set when using credentials.");
-
-            final TokenCredential tokenCredential = new DefaultAzureCredentialBuilder().build();
-
-            return builder.credential(fullyQualifiedDomainName, tokenCredential);
+    protected ServiceBusClientBuilder getBuilder(boolean sharedConnection) {
+        ServiceBusClientBuilder builder;
+        if (sharedConnection && sharedBuilder == null) {
+            sharedBuilder = getBuilder();
+            builder = sharedBuilder;
+        } else if (sharedConnection) {
+            builder = sharedBuilder;
         } else {
-            return builder.connectionString(getConnectionString());
+            builder = getBuilder();
         }
+        return builder;
     }
 
-    protected ServiceBusSenderClientBuilder getSenderBuilder(boolean useCredentials, MessagingEntityType entityType,
+    protected ServiceBusSenderClientBuilder getSenderBuilder(MessagingEntityType entityType,
         int entityIndex, boolean isSessionAware, boolean sharedConnection) {
 
-        ServiceBusClientBuilder builder = getBuilder(useCredentials, sharedConnection);
+        ServiceBusClientBuilder builder = getBuilder(sharedConnection);
         switch (entityType) {
             case QUEUE:
                 final String queueName = isSessionAware ? getSessionQueueName(entityIndex) : getQueueName(entityIndex);
@@ -260,10 +249,10 @@ protected ServiceBusSenderClientBuilder getSenderBuilder(boolean useCredentials,
         }
     }
 
-    protected ServiceBusReceiverClientBuilder getReceiverBuilder(boolean useCredentials, MessagingEntityType entityType,
+    protected ServiceBusReceiverClientBuilder getReceiverBuilder(MessagingEntityType entityType,
         int entityIndex, boolean sharedConnection) {
 
-        ServiceBusClientBuilder builder = getBuilder(useCredentials, sharedConnection);
+        ServiceBusClientBuilder builder = getBuilder(sharedConnection);
         switch (entityType) {
             case QUEUE:
                 final String queueName = getQueueName(entityIndex);
@@ -283,10 +272,10 @@ protected ServiceBusReceiverClientBuilder getReceiverBuilder(boolean useCredenti
         }
     }
 
-    protected ServiceBusSessionReceiverClientBuilder getSessionReceiverBuilder(boolean useCredentials,
-        MessagingEntityType entityType, int entityIndex, boolean sharedConnection, AmqpRetryOptions retryOptions) {
+    protected ServiceBusSessionReceiverClientBuilder getSessionReceiverBuilder(MessagingEntityType entityType,
+        int entityIndex, boolean sharedConnection, AmqpRetryOptions retryOptions) {
 
-        ServiceBusClientBuilder builder = getBuilder(useCredentials, sharedConnection);
+        ServiceBusClientBuilder builder = getBuilder(sharedConnection);
 
         switch (entityType) {
             case QUEUE:
@@ -453,19 +442,6 @@ protected void assertMessageEquals(ServiceBusReceivedMessage message, String mes
         }
     }
 
-    protected ServiceBusClientBuilder getBuilder(boolean useCredentials, boolean sharedConnection) {
-        ServiceBusClientBuilder builder;
-        if (sharedConnection && sharedBuilder == null) {
-            sharedBuilder = getBuilder(useCredentials);
-            builder = sharedBuilder;
-        } else if (sharedConnection) {
-            builder = sharedBuilder;
-        } else {
-            builder = getBuilder(useCredentials);
-        }
-        return builder;
-    }
-
     protected final Configuration v1OrV2(boolean isV2) {
         final TestConfigurationSource configSource = new TestConfigurationSource();
         if (isV2) {
@@ -486,4 +462,38 @@ protected final Configuration v1OrV2(boolean isV2) {
         return new ConfigurationBuilder(configSource)
             .build();
     }
+
+    /**
+     * Asserts that if the integration tests can be run. This method is expected to be called at the beginning of each
+     * test run.
+     *
+     * @throws org.opentest4j.TestAbortedException if the integration tests cannot be run.
+     */
+    protected void assertRunnable() {
+        final TestMode mode = super.getTestMode();
+        if (mode == TestMode.PLAYBACK) {
+            // AMQP traffic never gets recorded so there is no PLAYBACK supported.
+            assumeTrue(false, "Skipping integration tests in playback mode.");
+            return;
+        }
+
+        if (mode == TestMode.RECORD) {
+            // RECORD mode used in SDK-dev setup not on CI pipeline.
+            if (!CoreUtils.isNullOrEmpty(TestUtils.getConnectionString(false))) {
+                // integration tests are runnable using the connection string.
+                return;
+            }
+            if (!CoreUtils.isNullOrEmpty(TestUtils.getFullyQualifiedDomainName(false))) {
+                // best effort check:
+                // integration tests are potentially runnable using DefaultAzureCredential. Here we're assuming that
+                // in RECORD mode with FullyQualifiedDomainName set, the dev environment is also set up for one of the
+                // token credential type in DefaultAzureCredential (all token credentials requires FullyQualifiedDomainName).
+                return;
+            }
+            assumeTrue(false, "Not running integration in record mode (missing authentication set up).");
+            return;
+        }
+        // The CI pipeline is expected to have federated identity configured, so tests are runnable.
+        assert mode == TestMode.LIVE;
+    }
 }

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ProxyReceiveTest.java

@@ -79,21 +79,18 @@ public void receiveMessage() {
         final String messageTracking = UUID.randomUUID().toString();
 
         final List<ServiceBusMessage> messages = TestUtils.getServiceBusMessages(NUMBER_OF_EVENTS, messageTracking);
-        final ServiceBusSenderAsyncClient sender = new ServiceBusClientBuilder()
+        final ServiceBusSenderAsyncClient sender = getAuthenticatedBuilder()
             .transportType(AmqpTransportType.AMQP_WEB_SOCKETS)
             .verifyMode(SslDomain.VerifyMode.ANONYMOUS_PEER)
-            .connectionString(getConnectionString())
-
             .sender()
             .queueName(queueName)
             .buildAsyncClient();
 
         toClose(sender);
 
-        final ServiceBusReceiverAsyncClient receiver = new ServiceBusClientBuilder()
+        final ServiceBusReceiverAsyncClient receiver = getAuthenticatedBuilder()
             .transportType(AmqpTransportType.AMQP_WEB_SOCKETS)
             .verifyMode(SslDomain.VerifyMode.ANONYMOUS_PEER)
-            .connectionString(getConnectionString())
             .receiver()
             .receiveMode(ServiceBusReceiveMode.RECEIVE_AND_DELETE)
             .queueName(queueName)

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ProxySelectorTest.java

@@ -68,8 +68,7 @@ public void connectFailed(URI uri, SocketAddress sa, IOException ioe) {
         });
 
         final ServiceBusMessage message = new ServiceBusMessage(BinaryData.fromString("Hello"));
-        final ServiceBusSenderAsyncClient sender = new ServiceBusClientBuilder()
-            .connectionString(getConnectionString())
+        final ServiceBusSenderAsyncClient sender = getAuthenticatedBuilder()
             .transportType(AmqpTransportType.AMQP_WEB_SOCKETS)
             .retryOptions(new AmqpRetryOptions().setTryTimeout(Duration.ofSeconds(10)))
             .sender()

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ProxySendTest.java

@@ -77,8 +77,7 @@ public void sendEvents() {
         final String messageId = UUID.randomUUID().toString();
 
         final List<ServiceBusMessage> messages = TestUtils.getServiceBusMessages(NUMBER_OF_EVENTS, messageId);
-        final ServiceBusSenderAsyncClient sender = new ServiceBusClientBuilder()
-            .connectionString(getConnectionString())
+        final ServiceBusSenderAsyncClient sender = getAuthenticatedBuilder()
             .transportType(AmqpTransportType.AMQP_WEB_SOCKETS)
             .verifyMode(SslDomain.VerifyMode.ANONYMOUS_PEER)
             .retryOptions(new AmqpRetryOptions().setTryTimeout(Duration.ofSeconds(10)))

sdk/servicebus/azure-messaging-servicebus/src/test/java/com/azure/messaging/servicebus/ServiceBusClientBuilderTest.java

@@ -6,7 +6,6 @@
 import com.azure.core.amqp.AmqpTransportType;
 import com.azure.core.amqp.ProxyAuthenticationType;
 import com.azure.core.amqp.ProxyOptions;
-import com.azure.core.amqp.implementation.ConnectionStringProperties;
 import com.azure.core.credential.AzureNamedKeyCredential;
 import com.azure.core.credential.AzureSasCredential;
 import com.azure.core.util.Configuration;
@@ -21,7 +20,6 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
-import reactor.test.StepVerifier;
 
 import java.net.InetSocketAddress;
 import java.net.Proxy;
@@ -34,8 +32,6 @@
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
-import static java.nio.charset.StandardCharsets.UTF_8;
-
 class ServiceBusClientBuilderTest extends IntegrationTestBase {
     private static final String NAMESPACE_NAME = "dummyNamespaceName";
     private static final String DEFAULT_DOMAIN_NAME = TestUtils.getEndpoint().substring(1) + "/";
@@ -266,55 +262,6 @@ public void testConnectionStringWithSas() {
                 .connectionString("Endpoint=sb://sb-name" + TestUtils.getEndpoint() + "/;EntityPath=sb-name"));
     }
 
-    @Test
-    public void testBatchSendEventByAzureNameKeyCredential() {
-        ConnectionStringProperties properties = getConnectionStringProperties();
-        String fullyQualifiedNamespace = getFullyQualifiedDomainName();
-        String sharedAccessKeyName = properties.getSharedAccessKeyName();
-        String sharedAccessKey = properties.getSharedAccessKey();
-        String queueName = getQueueName(TestUtils.USE_CASE_DEFAULT);
-
-        final ServiceBusMessage testData = new ServiceBusMessage(TEST_MESSAGE.getBytes(UTF_8));
-
-        ServiceBusSenderAsyncClient senderAsyncClient = toClose(new ServiceBusClientBuilder()
-            .credential(fullyQualifiedNamespace, new AzureNamedKeyCredential(sharedAccessKeyName, sharedAccessKey))
-            .sender()
-            .queueName(queueName)
-            .buildAsyncClient());
-        StepVerifier.create(
-            senderAsyncClient.createMessageBatch().flatMap(batch -> {
-                assertTrue(batch.tryAddMessage(testData));
-                return senderAsyncClient.sendMessages(batch);
-            }))
-            .expectComplete()
-            .verify(TIMEOUT);
-    }
-
-    @Test
-    public void testBatchSendEventByAzureSasCredential() {
-        ConnectionStringProperties properties = getConnectionStringProperties(true);
-        String fullyQualifiedNamespace = getFullyQualifiedDomainName();
-        String sharedAccessSignature = properties.getSharedAccessSignature();
-        String queueName = getQueueName(TestUtils.USE_CASE_DEFAULT);
-
-        final ServiceBusMessage testData = new ServiceBusMessage(TEST_MESSAGE.getBytes(UTF_8));
-
-        ServiceBusSenderAsyncClient senderAsyncClient = toClose(new ServiceBusClientBuilder()
-            .credential(fullyQualifiedNamespace,
-                new AzureSasCredential(sharedAccessSignature))
-            .sender()
-            .queueName(queueName)
-            .buildAsyncClient());
-
-        StepVerifier.create(
-            senderAsyncClient.createMessageBatch().flatMap(batch -> {
-                assertTrue(batch.tryAddMessage(testData));
-                return senderAsyncClient.sendMessages(batch);
-            }))
-            .expectComplete()
-            .verify(TIMEOUT);
-    }
-
     @Test
     public void testConnectionWithAzureNameKeyCredential() {
         String fullyQualifiedNamespace = "sb-name" + TestUtils.getEndpoint();