Add unit tests for gladstone extension handling service.

src/Features/Core/Portable/Extensions/ExtensionFolder.cs

@@ -28,7 +28,7 @@ private sealed class ExtensionFolder
             /// <summary>
             /// Lazily computed assembly loader for this particular folder.
             /// </summary>
-            private readonly AsyncLazy<(IAnalyzerAssemblyLoaderInternal? assemblyLoader, Exception? extensionException)> _lazyAssemblyLoader;
+            private readonly AsyncLazy<(IExtensionAssemblyLoader? assemblyLoader, Exception? extensionException)> _lazyAssemblyLoader;
 
             /// <summary>
             /// Mapping from assembly file path to the handlers it contains.  Should only be mutated while the <see
@@ -43,52 +43,16 @@ public ExtensionFolder(
                 _extensionMessageHandlerService = extensionMessageHandlerService;
                 _lazyAssemblyLoader = AsyncLazy.Create(cancellationToken =>
                 {
-#if NET
-                    // These lines should always succeed.  If they don't, they indicate a bug in our code that we want
-                    // to bubble out as it must be fixed.
-                    var analyzerAssemblyLoaderProvider = _extensionMessageHandlerService._solutionServices.GetRequiredService<IAnalyzerAssemblyLoaderProvider>();
-                    var analyzerAssemblyLoader = analyzerAssemblyLoaderProvider.CreateNewShadowCopyLoader();
-
-                    // Catch exceptions here related to working with the file system.  If we can't properly enumerate,
-                    // we want to report that back to the client, while not blocking the entire extension service.
-                    try
-                    {
-                        // Allow this assembly loader to load any dll in assemblyFolderPath.
-                        foreach (var dll in Directory.EnumerateFiles(assemblyFolderPath, "*.dll"))
-                        {
-                            cancellationToken.ThrowIfCancellationRequested();
-                            try
-                            {
-                                // Check if the file is a valid .NET assembly.
-                                AssemblyName.GetAssemblyName(dll);
-                            }
-                            catch
-                            {
-                                // The file is not a valid .NET assembly, skip it.
-                                continue;
-                            }
-
-                            analyzerAssemblyLoader.AddDependencyLocation(dll);
-                        }
-
-                        return ((IAnalyzerAssemblyLoaderInternal?)analyzerAssemblyLoader, (Exception?)null);
-                    }
-                    catch (Exception ex) when (ex is not OperationCanceledException)
-                    {
-                        // Capture any exceptions here to be reported back in CreateAssemblyHandlersAsync.
-                        return (null, ex);
-                    }
-#else
-                    return ((IAnalyzerAssemblyLoaderInternal?)null, (Exception?)null);
-#endif
+                    var analyzerAssemblyLoaderProvider = _extensionMessageHandlerService._solutionServices.GetRequiredService<IExtensionAssemblyLoaderProvider>();
+                    return analyzerAssemblyLoaderProvider.CreateNewShadowCopyLoader(assemblyFolderPath, cancellationToken);
                 });
             }
 
             public void Unload()
             {
                 // Only if we've created the assembly loader do we need to do anything.
                 _lazyAssemblyLoader.TryGetValue(out var tuple);
-                tuple.assemblyLoader?.Dispose();
+                tuple.assemblyLoader?.Unload();
             }
 
             private async Task<AssemblyMessageHandlers> CreateAssemblyHandlersAsync(
@@ -109,8 +73,14 @@ private async Task<AssemblyMessageHandlers> CreateAssemblyHandlersAsync(
                 }
 
                 var assembly = analyzerAssemblyLoader.LoadFromPath(assemblyFilePath);
-                var factory = _extensionMessageHandlerService._customMessageHandlerFactory;
-                Contract.ThrowIfNull(factory);
+                var factory = _extensionMessageHandlerService._solutionServices.GetService<IExtensionMessageHandlerFactory>();
+                if (factory is null)
+                {
+                    return new(
+                        DocumentMessageHandlers: ImmutableDictionary<string, IExtensionMessageHandlerWrapper>.Empty,
+                        WorkspaceMessageHandlers: ImmutableDictionary<string, IExtensionMessageHandlerWrapper>.Empty,
+                        ExtensionException: null);
+                }
 
                 // We're calling into code here to analyze the assembly at the specified file and to create handlers we
                 // find within it.  If this throws, then we will capture that exception and return it to the caller to 

src/Features/Core/Portable/Extensions/ExtensionMessageHandlerService.cs

@@ -23,12 +23,10 @@ private readonly record struct AssemblyMessageHandlers(
         Exception? ExtensionException);
 
     private sealed partial class ExtensionMessageHandlerService(
-        SolutionServices solutionServices,
-        IExtensionMessageHandlerFactory? customMessageHandlerFactory)
+        SolutionServices solutionServices)
         : IExtensionMessageHandlerService
     {
         private readonly SolutionServices _solutionServices = solutionServices;
-        private readonly IExtensionMessageHandlerFactory? _customMessageHandlerFactory = customMessageHandlerFactory;
 
         /// <summary>
         /// Lock for <see cref="_folderPathToExtensionFolder"/>, <see cref="_cachedDocumentHandlers_useOnlyUnderLock"/>, and <see

src/Features/Core/Portable/Extensions/ExtensionMessageHandlerServiceFactory.cs

@@ -12,10 +12,9 @@ namespace Microsoft.CodeAnalysis.Extensions;
 [ExportWorkspaceServiceFactory(typeof(IExtensionMessageHandlerService)), Shared]
 [method: ImportingConstructor]
 [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
-internal sealed partial class ExtensionMessageHandlerServiceFactory(
-    [Import(AllowDefault = true)] IExtensionMessageHandlerFactory? customMessageHandlerFactory)
+internal sealed partial class ExtensionMessageHandlerServiceFactory()
     : IWorkspaceServiceFactory
 {
     public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
-        => new ExtensionMessageHandlerService(workspaceServices.SolutionServices, customMessageHandlerFactory);
+        => new ExtensionMessageHandlerService(workspaceServices.SolutionServices);
 }

src/Features/Core/Portable/Extensions/IExtensionAssemblyLoaderProvider.cs

@@ -0,0 +1,102 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Composition;
+using System.IO;
+using System.Reflection;
+using System.Threading;
+using Microsoft.CodeAnalysis.Host;
+using Microsoft.CodeAnalysis.Host.Mef;
+
+namespace Microsoft.CodeAnalysis.Extensions;
+
+/// <summary>
+/// Abstraction around <see cref="IAnalyzerAssemblyLoaderProvider"/> so that we can mock that behavior in tests.
+/// </summary>
+internal interface IExtensionAssemblyLoaderProvider : IWorkspaceService
+{
+    (IExtensionAssemblyLoader? assemblyLoader, Exception? extensionException) CreateNewShadowCopyLoader(
+        string assemblyFolderPath, CancellationToken cancellationToken);
+}
+
+/// <summary>
+/// Abstraction around <see cref="IAnalyzerAssemblyLoader"/> so that we can mock that behavior in tests.
+/// </summary>
+internal interface IExtensionAssemblyLoader
+{
+    Assembly LoadFromPath(string assemblyFilePath);
+    void Unload();
+}
+
+[ExportWorkspaceServiceFactory(typeof(IExtensionAssemblyLoaderProvider)), Shared]
+[method: ImportingConstructor]
+[method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
+internal sealed class DefaultExtensionAssemblyLoaderProviderFactory() : IWorkspaceServiceFactory
+{
+    public IWorkspaceService CreateService(HostWorkspaceServices workspaceServices)
+        => new DefaultExtensionAssemblyLoaderProvider(workspaceServices);
+
+    private sealed class DefaultExtensionAssemblyLoaderProvider(HostWorkspaceServices workspaceServices)
+        : IExtensionAssemblyLoaderProvider
+    {
+#pragma warning disable IDE0052 // Remove unread private members
+        private readonly HostWorkspaceServices _workspaceServices = workspaceServices;
+#pragma warning restore IDE0052 // Remove unread private members
+
+        public (IExtensionAssemblyLoader? assemblyLoader, Exception? extensionException) CreateNewShadowCopyLoader(
+            string assemblyFolderPath, CancellationToken cancellationToken)
+        {
+            cancellationToken.ThrowIfCancellationRequested();
+
+#if NET
+            // These lines should always succeed.  If they don't, they indicate a bug in our code that we want
+            // to bubble out as it must be fixed.
+            var analyzerAssemblyLoaderProvider = _workspaceServices.GetRequiredService<IAnalyzerAssemblyLoaderProvider>();
+            var analyzerAssemblyLoader = analyzerAssemblyLoaderProvider.CreateNewShadowCopyLoader();
+
+            // Catch exceptions here related to working with the file system.  If we can't properly enumerate,
+            // we want to report that back to the client, while not blocking the entire extension service.
+            try
+            {
+                // Allow this assembly loader to load any dll in assemblyFolderPath.
+                foreach (var dll in Directory.EnumerateFiles(assemblyFolderPath, "*.dll"))
+                {
+                    cancellationToken.ThrowIfCancellationRequested();
+                    try
+                    {
+                        // Check if the file is a valid .NET assembly.
+                        AssemblyName.GetAssemblyName(dll);
+                    }
+                    catch
+                    {
+                        // The file is not a valid .NET assembly, skip it.
+                        continue;
+                    }
+
+                    analyzerAssemblyLoader.AddDependencyLocation(dll);
+                }
+
+                return (new DefaultExtensionAssemblyLoader(analyzerAssemblyLoader), null);
+            }
+            catch (Exception ex) when (ex is not OperationCanceledException)
+            {
+                // Capture any exceptions here to be reported back in CreateAssemblyHandlersAsync.
+                return (null, ex);
+            }
+#else
+            return default;
+#endif
+        }
+    }
+
+    private sealed class DefaultExtensionAssemblyLoader(
+        IAnalyzerAssemblyLoaderInternal assemblyLoader) : IExtensionAssemblyLoader
+    {
+        public void Unload() => assemblyLoader.Dispose();
+
+        public Assembly LoadFromPath(string assemblyFilePath)
+            => assemblyLoader.LoadFromPath(assemblyFilePath);
+    }
+}

src/Features/Core/Portable/Extensions/IExtensionMessageHandlerFactory.cs

@@ -5,13 +5,14 @@
 using System.Collections.Immutable;
 using System.Reflection;
 using System.Threading;
+using Microsoft.CodeAnalysis.Host;
 
 namespace Microsoft.CodeAnalysis.Extensions;
 
 /// <summary>
 /// Factory for creating instances of extension message handlers.
 /// </summary>
-internal interface IExtensionMessageHandlerFactory
+internal interface IExtensionMessageHandlerFactory : IWorkspaceService
 {
     /// <summary>
     /// Creates <see cref="IExtensionMessageHandlerWrapper{Solution}"/> instances for each

src/Tools/ExternalAccess/Extensions/Internal/ExtensionMessageHandlerFactory.cs

@@ -11,7 +11,7 @@
 
 namespace Microsoft.CodeAnalysis.Extensions;
 
-[Export(typeof(IExtensionMessageHandlerFactory)), Shared]
+[ExportWorkspaceService(typeof(IExtensionMessageHandlerFactory)), Shared]
 [method: ImportingConstructor]
 [method: Obsolete(MefConstruction.ImportingConstructorMessage, error: true)]
 internal sealed class ExtensionMessageHandlerFactory() : IExtensionMessageHandlerFactory

src/VisualStudio/Core/Test.Next/Services/ServiceHubServicesTests.cs

@@ -38,10 +38,22 @@ namespace Roslyn.VisualStudio.Next.UnitTests.Remote;
 
 [UseExportProvider]
 [Trait(Traits.Feature, Traits.Features.RemoteHost)]
-public sealed class ServiceHubServicesTests
+public sealed partial class ServiceHubServicesTests
 {
-    private static TestWorkspace CreateWorkspace(Type[] additionalParts = null)
-         => new(composition: FeaturesTestCompositions.Features.WithTestHostParts(TestHost.OutOfProcess).AddParts(additionalParts));
+    private static TestWorkspace CreateWorkspace(
+        Type[] additionalParts = null,
+        Type[] additionalRemoteParts = null)
+    {
+        var workspace = new TestWorkspace(composition: FeaturesTestCompositions.Features.WithTestHostParts(TestHost.OutOfProcess).AddParts(additionalParts));
+
+        if (additionalRemoteParts != null)
+        {
+            var clientProvider = (InProcRemoteHostClientProvider)workspace.Services.GetService<IRemoteHostClientProvider>();
+            clientProvider.AdditionalRemoteParts = additionalRemoteParts;
+        }
+
+        return workspace;
+    }
 
     [Fact]
     public async Task TestRemoteHostSynchronize()
@@ -1721,8 +1733,16 @@ private static Solution Populate(Solution solution)
         return solution;
     }
 
-    private static Solution AddProject(Solution solution, string language, string[] documents, string[] additionalDocuments, ProjectId[] p2pReferences)
+    private static Solution AddProject(
+        Solution solution,
+        string language,
+        string[] documents,
+        string[] additionalDocuments = null,
+        ProjectId[] p2pReferences = null)
     {
+        additionalDocuments ??= [];
+        p2pReferences ??= [];
+
         var projectName = $"Project{solution.ProjectIds.Count}";
         var project = solution.AddProject(projectName, $"{projectName}.dll", language)
                               .AddMetadataReference(MetadataReference.CreateFromFile(typeof(object).Assembly.Location))