Resolve comments

sdk/identity/azure-identity-extensions/CHANGELOG.md

@@ -2,9 +2,11 @@
 
 ## 1.2.0-beta.2 (Unreleased)
 
+#### Features Added
+- Support cache for token credential object. [#39393](https://github.com/Azure/azure-sdk-for-java/issues/39393).
+
 #### Bugs Fixed
 - Fix the issue where the token acquisition timeout is not set via the property `azure.accessTokenTimeoutInSeconds`. [#43512](https://github.com/Azure/azure-sdk-for-java/issues/43512).
-- Fix inefficiencies: Support cache for token credential. [#39393](https://github.com/Azure/azure-sdk-for-java/issues/39393).
 
 ### Other Changes
 

sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/credential/provider/CachingTokenCredentialProviderTest.java

@@ -5,14 +5,16 @@
 
 import com.azure.core.credential.TokenCredential;
 import com.azure.identity.DefaultAzureCredential;
+import com.azure.identity.ManagedIdentityCredential;
+import com.azure.identity.extensions.implementation.credential.TokenCredentialProviderOptions;
 import org.junit.jupiter.api.Test;
 
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 class CachingTokenCredentialProviderTest {
 
     @Test
-    void cacheTokenCredential() {
+    void cacheTokenCredentialUsingDefaultOptions() {
         DefaultTokenCredentialProvider defaultTokenCredentialProvider1 = new DefaultTokenCredentialProvider(null);
         CachingTokenCredentialProvider provider1
             = new CachingTokenCredentialProvider(null, defaultTokenCredentialProvider1);
@@ -28,4 +30,83 @@ void cacheTokenCredential() {
         assertTrue(tokenCredential1 == tokenCredential2);
     }
 
+    @Test
+    void cacheTokenCredentialUsingCustomOptions() {
+        TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
+
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider1
+            = new DefaultTokenCredentialProvider(customOptions);
+        CachingTokenCredentialProvider cachingProvider1
+            = new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
+        TokenCredential tokenCredential1 = cachingProvider1.get();
+
+        TokenCredentialProviderOptions customOptions2 = getSystemManagedIdentityCredentialProviderOptions();
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider2
+            = new DefaultTokenCredentialProvider(customOptions2);
+        CachingTokenCredentialProvider cachingProvider2
+            = new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
+
+        TokenCredential tokenCredential2 = cachingProvider2.get();
+        assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential1 == tokenCredential2);
+    }
+
+    @Test
+    void cacheTokenCredentialViaDifferentGetMethods() {
+        TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
+
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider1
+            = new DefaultTokenCredentialProvider(customOptions);
+        CachingTokenCredentialProvider cachingProvider1
+            = new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
+        TokenCredential tokenCredential1 = cachingProvider1.get();
+
+        TokenCredentialProviderOptions customOptions2 = getSystemManagedIdentityCredentialProviderOptions();
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider2
+            = new DefaultTokenCredentialProvider(customOptions2);
+        CachingTokenCredentialProvider cachingProvider2
+            = new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
+
+        TokenCredential tokenCredential2 = cachingProvider2.get(customOptions2);
+        assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential1 == tokenCredential2);
+    }
+
+    @Test
+    void cacheTokenCredentialViaDifferentProviderOptions() {
+        TokenCredentialProviderOptions customOptions = getSystemManagedIdentityCredentialProviderOptions();
+
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider1
+            = new DefaultTokenCredentialProvider(customOptions);
+        CachingTokenCredentialProvider cachingProvider1
+            = new CachingTokenCredentialProvider(customOptions, defaultTokenCredentialProvider1);
+        TokenCredential tokenCredential1 = cachingProvider1.get();
+
+        TokenCredentialProviderOptions customOptions2
+            = getUserManagedIdentityCredentialProviderOptions("test-client-id");
+        DefaultTokenCredentialProvider defaultTokenCredentialProvider2
+            = new DefaultTokenCredentialProvider(customOptions2);
+        CachingTokenCredentialProvider cachingProvider2
+            = new CachingTokenCredentialProvider(customOptions2, defaultTokenCredentialProvider2);
+
+        TokenCredential tokenCredential2 = cachingProvider2.get(customOptions2);
+        assertTrue(tokenCredential1 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential2 instanceof ManagedIdentityCredential);
+        assertTrue(tokenCredential1 != tokenCredential2);
+    }
+
+    private static TokenCredentialProviderOptions getUserManagedIdentityCredentialProviderOptions(String clientId) {
+        TokenCredentialProviderOptions customOptions = new TokenCredentialProviderOptions();
+        customOptions.setManagedIdentityEnabled(true);
+        customOptions.setClientId(clientId);
+        return customOptions;
+    }
+
+    private static TokenCredentialProviderOptions getSystemManagedIdentityCredentialProviderOptions() {
+        TokenCredentialProviderOptions customOptions = new TokenCredentialProviderOptions();
+        customOptions.setManagedIdentityEnabled(true);
+        return customOptions;
+    }
 }

sdk/identity/azure-identity-extensions/src/test/java/com/azure/identity/extensions/implementation/template/AzureAuthenticationTemplateTest.java

@@ -91,10 +91,11 @@ void cacheTokenCredential() {
 
         TokenCredentialProviderOptions providerOptions = new TokenCredentialProviderOptions(properties);
 
-        assertNotNull(template.getTokenCredentialProvider().get(providerOptions));
-        assertNotNull(template2.getTokenCredentialProvider().get(providerOptions));
-        assertEquals(template.getTokenCredentialProvider().get(providerOptions),
-            template2.getTokenCredentialProvider().get(providerOptions));
+        TokenCredential tokenCredential = template.getTokenCredentialProvider().get(providerOptions);
+        TokenCredential tokenCredential2 = template2.getTokenCredentialProvider().get(providerOptions);
+        assertNotNull(tokenCredential);
+        assertNotNull(tokenCredential2);
+        assertTrue(tokenCredential == tokenCredential2);
     }
 
     @Test