update tests

Author: g2vinay

sdk/identity/azure-identity/src/main/java/com/azure/identity/OSBrokerCredential.java

@@ -15,8 +15,7 @@
 
 class OSBrokerCredential implements TokenCredential {
     private static final ClientLogger LOGGER = new ClientLogger(OSBrokerCredential.class);
-    private static final String BROKER_BUILDER_CLASS
-        = "com.azure.identity.broker.InteractiveBrowserBrokerCredentialBuilder";
+    private static final String BROKER_BUILDER_CLASS = "com.azure.identity.broker.InteractiveBrowserBrokerCredentialBuilder";
     private final String tenantId;
     private final AtomicReference<TokenCredential> cached = new AtomicReference<>();
 
@@ -28,7 +27,18 @@ class OSBrokerCredential implements TokenCredential {
     public Mono<AccessToken> getToken(TokenRequestContext request) {
         TokenCredential credential = cached.get();
         if (credential == null) {
-            TokenCredential newCredential = createBrokerCredential();
+            TokenCredential newCredential;
+            try {
+                // Create the broker credential dynamically
+                newCredential = createBrokerCredential();
+            } catch (CredentialUnavailableException e) {
+                // If the broker is unavailable, throw the exception
+                return Mono.error(e);
+            } catch (Exception e) {
+                // Log and throw any other exceptions that occur during credential creation
+                return Mono.error(LOGGER.logExceptionAsError(
+                    new CredentialUnavailableException("Failed to create OS Broker credential.", e)));
+            }
             if (cached.compareAndSet(null, newCredential)) {
                 credential = newCredential;
             } else {
@@ -39,12 +49,11 @@ public Mono<AccessToken> getToken(TokenRequestContext request) {
     }
 
     private TokenCredential createBrokerCredential() {
-        final String troubleshoot
-            = " To mitigate this issue, refer to http://aka.ms/azsdk/java/identity/dacbrokerauth/troubleshoot";
+        final String troubleshoot = " To mitigate this issue, refer to http://aka.ms/azsdk/java/identity/dacbrokerauth/troubleshoot";
         if (!IdentityUtil.isBrokerAvailable()) {
             throw LOGGER.logExceptionAsError(
                 new CredentialUnavailableException("azure-identity-broker dependency is not available. "
-                    + "Ensure you have azure-identity-broker dependency added to your application." + troubleshoot));
+                + "Ensure you have azure-identity-broker dependency added to your application." + troubleshoot));
         }
         try {
             Class<?> builderClass = Class.forName(BROKER_BUILDER_CLASS);
@@ -58,14 +67,12 @@ private TokenCredential createBrokerCredential() {
             }
             return browserCredentialBuilder.build();
         } catch (ClassNotFoundException e) {
-            throw LOGGER
-                .logExceptionAsError(new CredentialUnavailableException(
-                    "InteractiveBrowserBrokerCredentialBuilder class not found. "
-                        + "Ensure you have azure-identity-broker dependency added to your application." + troubleshoot,
-                    e));
+            throw LOGGER.logExceptionAsError(
+                new CredentialUnavailableException("InteractiveBrowserBrokerCredentialBuilder class not found. "
+                + "Ensure you have azure-identity-broker dependency added to your application." + troubleshoot, e));
         } catch (Exception e) {
-            throw LOGGER.logExceptionAsError(new CredentialUnavailableException(
-                "Failed to create InteractiveBrowserBrokerCredential dynamically." + troubleshoot, e));
+            throw LOGGER.logExceptionAsError(
+                new CredentialUnavailableException("Failed to create InteractiveBrowserBrokerCredential dynamically." + troubleshoot, e));
         }
     }
 }

sdk/identity/azure-identity/src/test/java/com/azure/identity/DefaultAzureCredentialTest.java

@@ -321,8 +321,12 @@ public void testNoCredentialWorks() {
                 = mockConstruction(IntelliJCredential.class, (intelliJCredential, context) -> {
                     when(intelliJCredential.getToken(request)).thenReturn(
                         Mono.error(new CredentialUnavailableException("Cannot get token from IntelliJ Credential")));
+                });
+            MockedConstruction<OSBrokerCredential> osBrokerCredentialMock
+                = mockConstruction(OSBrokerCredential.class, (osBrokerCredential, context) -> {
+                    when(osBrokerCredential.getToken(request)).thenReturn(
+                        Mono.error(new CredentialUnavailableException("Cannot get token from OS Broker credential")));
                 })) {
-
             // test
             DefaultAzureCredential credential
                 = new DefaultAzureCredentialBuilder().configuration(configuration).build();
@@ -335,6 +339,7 @@ public void testNoCredentialWorks() {
             Assertions.assertNotNull(azureDeveloperCliCredentialMock);
             Assertions.assertNotNull(azurePowerShellCredentialMock);
             Assertions.assertNotNull(intelliJCredentialMock);
+            Assertions.assertNotNull(osBrokerCredentialMock);
         }
     }
 
@@ -368,7 +373,13 @@ public void testCredentialUnavailable() {
                 = mockConstruction(AzureDeveloperCliCredential.class, (AzureDeveloperCliCredential, context) -> {
                     when(AzureDeveloperCliCredential.getToken(request)).thenReturn(Mono.error(
                         new CredentialUnavailableException("Cannot get token from Azure Developer CLI credential")));
+                });
+            MockedConstruction<OSBrokerCredential> osBrokerCredentialMock
+                = mockConstruction(OSBrokerCredential.class, (osBrokerCredential, context) -> {
+                    when(osBrokerCredential.getToken(request)).thenReturn(
+                        Mono.error(new CredentialUnavailableException("Cannot get token from OS Broker credential")));
                 })) {
+
             // test
             DefaultAzureCredential credential
                 = new DefaultAzureCredentialBuilder().configuration(configuration).build();
@@ -381,8 +392,8 @@ public void testCredentialUnavailable() {
             Assertions.assertNotNull(powerShellCredentialMock);
             Assertions.assertNotNull(azureCliCredentialMock);
             Assertions.assertNotNull(azureDeveloperCliCredentialMock);
+            Assertions.assertNotNull(osBrokerCredentialMock);
         }
-
     }
 
     @Test