Ramsess/test isolation bug fix

src/Authentication/Authentication.Loader/GraphLoadContextInitializer.cs

@@ -22,43 +22,51 @@ public static class GraphLoadContextInitializer
 
         private static readonly object s_lock = new object();
         private static GraphAssemblyLoadContext s_context;
+        private static bool s_resolvingHooked;
 
         /// <summary>
         /// The isolated context; null until <see cref="Initialize"/> has been called.
         /// </summary>
         public static GraphAssemblyLoadContext Context => s_context;
 
         /// <summary>
-        /// Creates the isolated context and registers the default-context redirect. Safe to call multiple times.
+        /// Creates the isolated context (once) and ensures the default-context redirect is registered. Safe to call
+        /// multiple times, including after a <see cref="Shutdown"/> triggered by a module remove: the isolated context is
+        /// reused while the <see cref="AssemblyLoadContext.Resolving"/> hook is re-established so a remove/re-import cycle
+        /// keeps resolving Microsoft.Graph.Authentication.Core.
         /// </summary>
         /// <param name="dependencyFolder">Path to the module's <c>Dependencies</c> folder.</param>
         /// <param name="psEditionDependencyFolder">Path to the module's <c>Dependencies/Core</c> folder.</param>
         public static void Initialize(string dependencyFolder, string psEditionDependencyFolder)
         {
-            if (s_context != null)
-                return;
-
             lock (s_lock)
             {
-                if (s_context != null)
-                    return;
+                if (s_context == null)
+                {
+                    s_context = new GraphAssemblyLoadContext(dependencyFolder, psEditionDependencyFolder);
+                }
 
-                s_context = new GraphAssemblyLoadContext(dependencyFolder, psEditionDependencyFolder);
-                AssemblyLoadContext.Default.Resolving += OnDefaultResolving;
+                if (!s_resolvingHooked)
+                {
+                    AssemblyLoadContext.Default.Resolving += OnDefaultResolving;
+                    s_resolvingHooked = true;
+                }
             }
         }
 
         /// <summary>
         /// Unregisters the default-context redirect. The isolated context itself is not collectible and remains alive
-        /// for the lifetime of the process, which matches PowerShell's own behaviour for binary modules.
+        /// for the lifetime of the process, which matches PowerShell's own behaviour for binary modules; a later
+        /// <see cref="Initialize"/> call re-registers the redirect against the existing context.
         /// </summary>
         public static void Shutdown()
         {
             lock (s_lock)
             {
-                if (s_context == null)
+                if (!s_resolvingHooked)
                     return;
                 AssemblyLoadContext.Default.Resolving -= OnDefaultResolving;
+                s_resolvingHooked = false;
             }
         }
 

src/Authentication/Authentication/ModuleInitializer.cs

@@ -65,7 +65,16 @@ static ModuleInitializer()
         public void OnImport()
         {
             if (s_isPsCore)
+            {
+                // The static constructor only runs on the first load of this assembly (it stays in the default
+                // AssemblyLoadContext for the process lifetime). On a remove/re-import cycle OnRemove -> Shutdown
+                // unhooks the default-context redirect, so re-run Initialize here to re-establish it against the
+                // already-created isolated context.
+                s_loaderInitializer?
+                    .GetMethod("Initialize", BindingFlags.Public | BindingFlags.Static, null, new[] { typeof(string), typeof(string) }, null)?
+                    .Invoke(null, new object[] { s_dependencyFolder, s_psEditionDependencyFolder });
                 return;
+            }
 
             AppDomain.CurrentDomain.AssemblyResolve += ResolvingHandler;
         }

src/Authentication/Authentication/Properties/launchSettings.json

@@ -2,7 +2,7 @@
   "profiles": {
     "Graph.Authentication": {
       "commandName": "Executable",
-      "executablePath": "C:\\Program Files\\PowerShell\\7\\pwsh.exe",
+      "executablePath": "pwsh.exe",
       "commandLineArgs": "-NoExit"
     }
   }

src/Authentication/Authentication/test/Microsoft.Graph.Authentication.Tests.ps1

@@ -72,4 +72,36 @@ Describe "Microsoft.Graph.Authentication module" {
             $PSModuleInfo.Guid.Guid | Should -Be "883916f2-9184-46ee-b1f8-b6a2fb784cee"
         }
     }
+
+    Context "On module re-import" {
+        # Regression guard for the switch to a private AssemblyLoadContext on PowerShell 7+.
+        # Importing, removing and re-importing the module must succeed: the second import must not fail
+        # because the isolated load context (and its already-loaded dependencies) is being initialized again.
+        It 'Should import, remove and re-import without error' {
+            {
+                Import-Module $ModulePath -Force
+                Remove-Module $ModuleName -Force
+                Import-Module $ModulePath -Force
+            } | Should -Not -Throw
+        }
+
+        It 'Should expose its cmdlets after a re-import' {
+            Import-Module $ModulePath -Force
+            Remove-Module $ModuleName -Force
+            $ReimportedInfo = Import-Module $ModulePath -Force -PassThru
+
+            $ReimportedInfo.ExportedCommands.Keys | Should -Contain "Connect-MgGraph"
+            $ReimportedInfo.ExportedCommands.Keys | Should -Contain "Invoke-MgGraphRequest"
+        }
+
+        It 'Should keep the same isolated assemblies loaded after a re-import' {
+            Import-Module $ModulePath -Force
+            Remove-Module $ModuleName -Force
+            Import-Module $ModulePath -Force
+
+            # Invoking a cmdlet forces the isolated dependencies (Azure.Identity / MSAL / Kiota) to bind.
+            # If the re-import left the load context in a broken state this throws instead of returning $null.
+            { Get-MgContext } | Should -Not -Throw
+        }
+    }
 }