Add ResetStateForTest API for test (#5489)

* Add ResetStateForTest API for test

* Keep same type

* Address PR comments

Author: bgavrilMS

src/client/Microsoft.Identity.Client/ApplicationBase.cs

@@ -6,8 +6,17 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Identity.Client.ApiConfig.Parameters;
+using Microsoft.Identity.Client.AuthScheme.PoP;
+using Microsoft.Identity.Client.Http;
+using Microsoft.Identity.Client.Instance;
+using Microsoft.Identity.Client.Instance.Discovery;
+using Microsoft.Identity.Client.Instance.Oidc;
 using Microsoft.Identity.Client.Internal;
 using Microsoft.Identity.Client.Internal.Requests;
+using Microsoft.Identity.Client.ManagedIdentity;
+using Microsoft.Identity.Client.OAuth2.Throttling;
+using Microsoft.Identity.Client.PlatformsCommon.Shared;
+using Microsoft.Identity.Client.Region;
 
 namespace Microsoft.Identity.Client
 {
@@ -74,5 +83,24 @@ internal static void GuardMobileFrameworks()
                 "See https://aka.ms/msal-net-confidential-availability and https://aka.ms/msal-net-managed-identity for details.");
 #endif
         }
+
+        /// <summary>
+        /// Resets the SDKs internal state, such as static caches, to facilitate testing. 
+        /// This API is meant to be used by other SDKs that build on top of MSAL, and only by test code.
+        /// </summary>
+        public static void ResetStateForTest()
+        {
+            NetworkCacheMetadataProvider.ResetStaticCacheForTest();
+            RegionManager.ResetStaticCacheForTest();
+            OidcRetrieverWithCache.ResetCacheForTest();
+            AuthorityManager.ClearValidationCache();
+            SingletonThrottlingManager.GetInstance().ResetCache();
+            ManagedIdentityClient.ResetSourceForTest();
+            AuthorityManager.ClearValidationCache();
+            PoPCryptoProviderFactory.Reset();
+
+            InMemoryPartitionedAppTokenCacheAccessor.ClearStaticCacheForTest();
+            InMemoryPartitionedUserTokenCacheAccessor.ClearStaticCacheForTest();
+        }
     }
 }

src/client/Microsoft.Identity.Client/Instance/Discovery/INetworkCacheMetadataProvider.cs

@@ -9,6 +9,5 @@ internal interface INetworkCacheMetadataProvider
     {
         void AddMetadata(string environment, InstanceDiscoveryMetadataEntry entry);
         InstanceDiscoveryMetadataEntry GetMetadata(string environment, ILoggerAdapter logger);
-        void /* for test purposes */ Clear();
     }
 }

src/client/Microsoft.Identity.Client/Instance/Discovery/InstanceDiscoveryManager.cs

@@ -37,12 +37,10 @@ internal class InstanceDiscoveryManager : IInstanceDiscoveryManager
 
         public InstanceDiscoveryManager(
           IHttpManager httpManager,
-          bool /* for test */ shouldClearCaches,
           InstanceDiscoveryResponse userProvidedInstanceDiscoveryResponse = null,
           Uri userProvidedInstanceDiscoveryUri = null) :
             this(
                 httpManager,                
-                shouldClearCaches,
                 userProvidedInstanceDiscoveryResponse != null ? new UserMetadataProvider(userProvidedInstanceDiscoveryResponse) : null,
                 userProvidedInstanceDiscoveryUri,
                 null, null, null, null)
@@ -51,7 +49,6 @@ public InstanceDiscoveryManager(
 
         public /* public for test */ InstanceDiscoveryManager(
             IHttpManager httpManager,
-            bool shouldClearCaches,
             IUserMetadataProvider userMetadataProvider = null,
             Uri userProvidedInstanceDiscoveryUri = null,
             IKnownMetadataProvider knownMetadataProvider = null,
@@ -72,12 +69,9 @@ public InstanceDiscoveryManager(
                     userProvidedInstanceDiscoveryUri);
 
             _regionDiscoveryProvider = regionDiscoveryProvider ??
-                new RegionAndMtlsDiscoveryProvider(_httpManager, shouldClearCaches);
+                new RegionAndMtlsDiscoveryProvider(_httpManager);
 
-            if (shouldClearCaches)
-            {
-                _networkCacheMetadataProvider.Clear();
-            }
+            
         }
 
         public InstanceDiscoveryMetadataEntry GetMetadataEntryAvoidNetwork(

src/client/Microsoft.Identity.Client/Instance/Discovery/NetworkCacheMetadataProvider.cs

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using System;
 using System.Collections.Concurrent;
 using Microsoft.Identity.Client.Core;
 
@@ -25,7 +26,7 @@ public void AddMetadata(string environment, InstanceDiscoveryMetadataEntry entry
             s_cache.AddOrUpdate(environment, entry, (_, _) => entry);
         }
 
-        public void Clear()
+        internal static void ResetStaticCacheForTest()
         {
             s_cache.Clear();
         }

src/client/Microsoft.Identity.Client/Instance/Discovery/RegionAndMtlsDiscoveryProvider.cs

@@ -16,9 +16,9 @@ internal class RegionAndMtlsDiscoveryProvider : IRegionDiscoveryProvider
         public const string PublicEnvForRegional = "login.microsoft.com";
         public const string PublicEnvForRegionalMtlsAuth = "mtlsauth.microsoft.com";
 
-        public RegionAndMtlsDiscoveryProvider(IHttpManager httpManager, bool clearCache)
+        public RegionAndMtlsDiscoveryProvider(IHttpManager httpManager)
         {
-            _regionManager = new RegionManager(httpManager, shouldClearStaticCache: clearCache);
+            _regionManager = new RegionManager(httpManager);
         }
 
         public async Task<InstanceDiscoveryMetadataEntry> GetMetadataAsync(Uri authority, RequestContext requestContext)
@@ -55,6 +55,7 @@ public async Task<InstanceDiscoveryMetadataEntry> GetMetadataAsync(Uri authority
             string regionalEnv = GetRegionalizedEnvironment(authority, region, requestContext);
             return CreateEntry(authority.Host, regionalEnv);
         }
+        
 
         private static InstanceDiscoveryMetadataEntry CreateEntry(string originalEnv, string regionalEnv)
         {
@@ -68,7 +69,6 @@ private static InstanceDiscoveryMetadataEntry CreateEntry(string originalEnv, st
 
         private static string GetRegionalizedEnvironment(Uri authority, string region, RequestContext requestContext)
         {
-
             string host = authority.Host;
 
             if (KnownMetadataProvider.IsPublicEnvironment(host))

src/client/Microsoft.Identity.Client/Instance/Region/RegionManager.cs

@@ -49,20 +49,13 @@ public RegionInfo(string region, RegionAutodetectionSource regionSource, string
 
         public RegionManager(
             IHttpManager httpManager,
-            int imdsCallTimeout = 2000,
-            bool shouldClearStaticCache = false) // for test
+            int imdsCallTimeout = 2000) // for test
         {
             _httpManager = httpManager;
             _imdsCallTimeoutMs = imdsCallTimeout;
-
-            if (shouldClearStaticCache)
-            {
-                s_failedAutoDiscovery = false;
-                s_autoDiscoveredRegion = null;
-                s_regionDiscoveryDetails = null;
-            }
         }
 
+        
         public async Task<string> GetAzureRegionAsync(RequestContext requestContext)
         {
             string azureRegionConfig = requestContext.ServiceBundle.Config.AzureRegion;
@@ -107,6 +100,13 @@ public async Task<string> GetAzureRegionAsync(RequestContext requestContext)
             return azureRegionConfig;
         }
 
+        internal static void ResetStaticCacheForTest()
+        {
+            s_failedAutoDiscovery = false;
+            s_autoDiscoveredRegion = null;
+            s_regionDiscoveryDetails = null;
+        }
+
         private static bool IsAutoDiscoveryRequested(string azureRegionConfig)
         {
             return string.Equals(azureRegionConfig, ConfidentialClientApplication.AttemptRegionDiscovery);

src/client/Microsoft.Identity.Client/Internal/ServiceBundle.cs

@@ -21,8 +21,7 @@ namespace Microsoft.Identity.Client.Internal
     internal class ServiceBundle : IServiceBundle
     {
         internal ServiceBundle(
-            ApplicationConfiguration config,
-            bool shouldClearCaches = false)
+            ApplicationConfiguration config)
         {
             Config = config;
 
@@ -38,20 +37,13 @@ internal ServiceBundle(
             HttpTelemetryManager = new HttpTelemetryManager();
 
             InstanceDiscoveryManager = new InstanceDiscoveryManager(
-                HttpManager,
-                shouldClearCaches,
+                HttpManager,                
                 config.CustomInstanceDiscoveryMetadata,
                 config.CustomInstanceDiscoveryMetadataUri);
 
             WsTrustWebRequestManager = new WsTrustWebRequestManager(HttpManager);
             ThrottlingManager = SingletonThrottlingManager.GetInstance();
-            DeviceAuthManager = config.DeviceAuthManagerForTest ?? PlatformProxy.CreateDeviceAuthManager();
-
-            if (shouldClearCaches) // for test
-            {
-                AuthorityManager.ClearValidationCache();
-                PoPCryptoProviderFactory.Reset();
-            }
+            DeviceAuthManager = config.DeviceAuthManagerForTest ?? PlatformProxy.CreateDeviceAuthManager();           
         }
 
         /// <summary>

src/client/Microsoft.Identity.Client/PlatformsCommon/Shared/InMemoryPartitionedAppTokenCacheAccessor.cs

@@ -256,5 +256,11 @@ private ref int GetEntryCountRef()
             return ref _tokenCacheAccessorOptions.UseSharedCache ? ref s_entryCount : ref _entryCount;
         }
 
+        public static void ClearStaticCacheForTest()
+        {
+            s_accessTokenCacheDictionary.Clear();
+            s_appMetadataDictionary.Clear();
+            Interlocked.Exchange(ref s_entryCount, 0);
+        }
     }
 }

src/client/Microsoft.Identity.Client/PlatformsCommon/Shared/InMemoryPartitionedUserTokenCacheAccessor.cs

@@ -5,6 +5,7 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading;
 using Microsoft.Identity.Client.Cache;
 using Microsoft.Identity.Client.Cache.Items;
 using Microsoft.Identity.Client.Cache.Keys;
@@ -168,7 +169,7 @@ public void DeleteAccessToken(MsalAccessTokenCacheItem item)
 
             if (AccessTokenCacheDictionary.TryGetValue(partitionKey, out var partition))
             {
-                bool removed = partition.TryRemove(item.CacheKey, out _);                
+                bool removed = partition.TryRemove(item.CacheKey, out _);
                 if (removed)
                 {
                     System.Threading.Interlocked.Decrement(ref GetEntryCountRef());
@@ -365,5 +366,15 @@ private ref int GetEntryCountRef()
         {
             return ref _tokenCacheAccessorOptions.UseSharedCache ? ref s_entryCount : ref _entryCount;
         }
+
+        public static void ClearStaticCacheForTest()
+        {
+            s_accessTokenCacheDictionary.Clear();
+            s_refreshTokenCacheDictionary.Clear();
+            s_idTokenCacheDictionary.Clear();
+            s_accountCacheDictionary.Clear();
+            s_appMetadataDictionary.Clear();
+            Interlocked.Exchange(ref s_entryCount, 0);
+        }
     }
 }

src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt

@@ -2,3 +2,4 @@ const Microsoft.Identity.Client.MsalError.InvalidCertificate = "invalid_certific
 Microsoft.Identity.Client.ManagedIdentityApplication.GetManagedIdentitySourceAsync() -> System.Threading.Tasks.Task<Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource>
 Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource.ImdsV2 = 8 -> Microsoft.Identity.Client.ManagedIdentity.ManagedIdentitySource
 Microsoft.Identity.Client.ManagedIdentityApplicationBuilder.WithExtraQueryParameters(System.Collections.Generic.IDictionary<string, string> extraQueryParameters) -> Microsoft.Identity.Client.ManagedIdentityApplicationBuilder
+static Microsoft.Identity.Client.ApplicationBase.ResetStateForTest() -> void