Remove IDocumentMappingService.GetLanguageKind(...) and make it an extension method on RazorCodeDocument (#10851)

For a long while, the `GetLanguageKind(...)` method that determines
whether an index into a document falls within Razor, C# or HTML has been
a bit of a wart on the `IDocumentMappingService`. It really isn't part
of document mapping, and its implementation is completely distinct. In
fact, making the actual change is quite simple, so why hadn't it been
done yet? The answer is mocking.

There are several tests that mock
`IDocumentMappingService.GetLanguageKind(...)` to lie about test inputs.
In my not-so-humble opinion, this represents an abuse of mocking.
Instead of setting up tests to have the necessary inputs that ensure
`GetLanguageKind(...)` would return a real and correct result, the
inputs would often be garbage and an `IDocumentMappingService` mock
would lie about the `GetLanguageKind(...)` result at a particular
location. This makes moving `GetLanguageKind(...)` off of
`IDocumentMappingService` a much larger change than it needs to be. This
is why there are substantial test changes in this PR.

Don't misunderstand me as a mocking hater! Mocking libraries are
definitely useful! In fact, there are new mocks used in this very PR!
However, mocks should be used judiciously and thoughtfully, and in this
case, a mock was used to write lazy tests.

Fixes https://github.com/dotnet/razor/issues/8774

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/CodeActions/CodeActionEndpoint.cs

@@ -177,7 +177,7 @@ internal sealed class CodeActionEndpoint(
 
     private async Task<ImmutableArray<RazorVSInternalCodeAction>> GetDelegatedCodeActionsAsync(DocumentContext documentContext, RazorCodeActionContext context, Guid correlationId, CancellationToken cancellationToken)
     {
-        var languageKind = _documentMappingService.GetLanguageKind(context.CodeDocument, context.Location.AbsoluteIndex, rightAssociative: false);
+        var languageKind = context.CodeDocument.GetLanguageKind(context.Location.AbsoluteIndex, rightAssociative: false);
 
         // No point delegating if we're in a Razor context
         if (languageKind == RazorLanguageKind.Razor)

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Debugging/RazorBreakpointSpanEndpoint.cs

@@ -55,7 +55,7 @@ internal class RazorBreakpointSpanEndpoint(
         }
 
         var projectedIndex = hostDocumentIndex;
-        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: false);
         // If we're in C#, then map to the right position in the generated document
         if (languageKind == RazorLanguageKind.CSharp &&
             !_documentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), hostDocumentIndex, out _, out projectedIndex))

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Debugging/RazorProximityExpressionsEndpoint.cs

@@ -56,7 +56,7 @@ internal class RazorProximityExpressionsEndpoint(
         }
 
         var projectedIndex = hostDocumentIndex;
-        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: false);
         // If we're in C#, then map to the right position in the generated document
         if (languageKind == RazorLanguageKind.CSharp &&
             !_documentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), hostDocumentIndex, out _, out projectedIndex))

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/DocumentPresentation/AbstractTextDocumentPresentationEndpointBase.cs

@@ -3,8 +3,6 @@
 
 using System;
 using System.Collections.Generic;
-using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
@@ -71,7 +69,7 @@ internal abstract class AbstractTextDocumentPresentationEndpointBase<TParams>(
             return null;
         }
 
-        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: false);
         // See if we can handle this directly in Razor. If not, we'll let things flow to the below delegated handling.
         var result = await TryGetRazorWorkspaceEditAsync(languageKind, request, cancellationToken).ConfigureAwait(false);
         if (result is not null)

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Formatting/DocumentOnTypeFormattingEndpoint.cs

@@ -3,16 +3,13 @@
 
 using System;
 using System.Collections.Frozen;
-using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.EndpointContracts;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
-using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Formatting;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/InlineCompletion/InlineCompletionEndPoint.cs

@@ -11,7 +11,6 @@ using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.EndpointContracts;
-using Microsoft.AspNetCore.Razor.LanguageServer.Formatting;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
 using Microsoft.AspNetCore.Razor.PooledObjects;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
@@ -86,7 +85,7 @@ internal sealed class InlineCompletionEndpoint(
         var sourceText = await documentContext.GetSourceTextAsync(cancellationToken).ConfigureAwait(false);
         var hostDocumentIndex = sourceText.GetPosition(request.Position);
 
-        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: false);
 
         // Map to the location in the C# document.
         if (languageKind != RazorLanguageKind.CSharp ||

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/Mapping/RazorLanguageQueryEndpoint.cs

@@ -61,7 +61,7 @@ internal sealed class RazorLanguageQueryEndpoint(IDocumentMappingService documen
 
         var responsePositionIndex = hostDocumentIndex;
 
-        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: false);
         if (languageKind == RazorLanguageKind.CSharp)
         {
             if (_documentMappingService.TryMapToGeneratedDocumentPosition(codeDocument.GetCSharpDocument(), hostDocumentIndex, out Position? projectedPosition, out var projectedIndex))

src/Razor/src/Microsoft.AspNetCore.Razor.LanguageServer/WrapWithTag/WrapWithTagEndpoint.cs

@@ -1,7 +1,6 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
-using System;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
@@ -9,7 +8,6 @@ using Microsoft.AspNetCore.Razor.Language.Syntax;
 using Microsoft.AspNetCore.Razor.LanguageServer.EndpointContracts;
 using Microsoft.AspNetCore.Razor.LanguageServer.Formatting;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
-using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Logging;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Text;
@@ -18,14 +16,9 @@ using Microsoft.VisualStudio.LanguageServer.Protocol;
 namespace Microsoft.AspNetCore.Razor.LanguageServer.WrapWithTag;
 
 [RazorLanguageServerEndpoint(LanguageServerConstants.RazorWrapWithTagEndpoint)]
-internal class WrapWithTagEndpoint(
-    IClientConnection clientConnection,
-    IDocumentMappingService documentMappingService,
-    ILoggerFactory loggerFactory)
-    : IRazorRequestHandler<WrapWithTagParams, WrapWithTagResponse?>
+internal class WrapWithTagEndpoint(IClientConnection clientConnection, ILoggerFactory loggerFactory) : IRazorRequestHandler<WrapWithTagParams, WrapWithTagResponse?>
 {
-    private readonly IClientConnection _clientConnection = clientConnection ?? throw new ArgumentNullException(nameof(clientConnection));
-    private readonly IDocumentMappingService _documentMappingService = documentMappingService ?? throw new ArgumentNullException(nameof(documentMappingService));
+    private readonly IClientConnection _clientConnection = clientConnection;
     private readonly ILogger _logger = loggerFactory.GetOrCreateLogger<WrapWithTagEndpoint>();
 
     public bool MutatesSolutionState => false;
@@ -69,7 +62,7 @@ internal class WrapWithTagEndpoint(
         //
         // Instead of C#, which certainly would be expected to go in an if statement, we'll see HTML, which obviously
         // is the better choice for this operation.
-        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: true);
+        var languageKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: true);
         if (languageKind is not RazorLanguageKind.Html)
         {
             // In general, we don't support C# for obvious reasons, but we can support implicit expressions. ie

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/DocumentMapping/AbstractDocumentMappingService.cs

@@ -8,11 +8,9 @@ using System.Diagnostics;
 using System.Linq;
 using Microsoft.AspNetCore.Razor;
 using Microsoft.AspNetCore.Razor.Language;
-using Microsoft.AspNetCore.Razor.Language.Legacy;
 using Microsoft.AspNetCore.Razor.PooledObjects;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Razor.Logging;
-using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -341,114 +339,6 @@ internal abstract class AbstractDocumentMappingService(IFilePathService filePath
         }
     }
 
-    public RazorLanguageKind GetLanguageKind(RazorCodeDocument codeDocument, int hostDocumentIndex, bool rightAssociative)
-    {
-        var classifiedSpans = GetClassifiedSpans(codeDocument);
-        var tagHelperSpans = GetTagHelperSpans(codeDocument);
-        var documentLength = codeDocument.Source.Text.Length;
-        var languageKind = GetLanguageKindCore(classifiedSpans, tagHelperSpans, hostDocumentIndex, documentLength, rightAssociative);
-
-        return languageKind;
-    }
-
-    // Internal for testing
-    internal static RazorLanguageKind GetLanguageKindCore(
-        ImmutableArray<ClassifiedSpanInternal> classifiedSpans,
-        ImmutableArray<TagHelperSpanInternal> tagHelperSpans,
-        int hostDocumentIndex,
-        int hostDocumentLength,
-        bool rightAssociative)
-    {
-        var length = classifiedSpans.Length;
-        for (var i = 0; i < length; i++)
-        {
-            var classifiedSpan = classifiedSpans[i];
-            var span = classifiedSpan.Span;
-
-            if (span.AbsoluteIndex <= hostDocumentIndex)
-            {
-                var end = span.AbsoluteIndex + span.Length;
-                if (end >= hostDocumentIndex)
-                {
-                    if (end == hostDocumentIndex)
-                    {
-                        // We're at an edge.
-
-                        if (classifiedSpan.SpanKind is SpanKindInternal.MetaCode or SpanKindInternal.Transition)
-                        {
-                            // If we're on an edge of a transition of some kind (MetaCode representing an open or closing piece of syntax such as <|,
-                            // and Transition representing an explicit transition to/from razor syntax, such as @|), prefer to classify to the span
-                            // to the right to better represent where the user clicks
-                            continue;
-                        }
-
-                        // If we're right associative, then we don't want to use the classification that we're at the end
-                        // of, if we're also at the start of the next one
-                        if (rightAssociative)
-                        {
-                            if (i < classifiedSpans.Length - 1 && classifiedSpans[i + 1].Span.AbsoluteIndex == hostDocumentIndex)
-                            {
-                                // If we're at the start of the next span, then use that span
-                                return GetLanguageFromClassifiedSpan(classifiedSpans[i + 1]);
-                            }
-
-                            // Otherwise, we did not find a match using right associativity, so check for tag helpers
-                            break;
-                        }
-                    }
-
-                    return GetLanguageFromClassifiedSpan(classifiedSpan);
-                }
-            }
-        }
-
-        foreach (var tagHelperSpan in tagHelperSpans)
-        {
-            var span = tagHelperSpan.Span;
-
-            if (span.AbsoluteIndex <= hostDocumentIndex)
-            {
-                var end = span.AbsoluteIndex + span.Length;
-                if (end >= hostDocumentIndex)
-                {
-                    if (end == hostDocumentIndex)
-                    {
-                        // We're at an edge. TagHelper spans never own their edge and aren't represented by marker spans
-                        continue;
-                    }
-
-                    // Found intersection
-                    return RazorLanguageKind.Html;
-                }
-            }
-        }
-
-        // Use the language of the last classified span if we're at the end
-        // of the document.
-        if (classifiedSpans.Length != 0 && hostDocumentIndex == hostDocumentLength)
-        {
-            var lastClassifiedSpan = classifiedSpans.Last();
-            return GetLanguageFromClassifiedSpan(lastClassifiedSpan);
-        }
-
-        // Default to Razor
-        return RazorLanguageKind.Razor;
-
-        static RazorLanguageKind GetLanguageFromClassifiedSpan(ClassifiedSpanInternal classifiedSpan)
-        {
-            // Overlaps with request
-            return classifiedSpan.SpanKind switch
-            {
-                SpanKindInternal.Markup => RazorLanguageKind.Html,
-                SpanKindInternal.Code => RazorLanguageKind.CSharp,
-
-                // Content type was non-C# or Html or we couldn't find a classified span overlapping the request position.
-                // All other classified span kinds default back to Razor
-                _ => RazorLanguageKind.Razor,
-            };
-        }
-    }
-
     private bool TryMapToHostDocumentRangeStrict(IRazorGeneratedDocument generatedDocument, LinePositionSpan generatedDocumentRange, out LinePositionSpan hostDocumentRange)
     {
         hostDocumentRange = default;
@@ -666,36 +556,4 @@ internal abstract class AbstractDocumentMappingService(IFilePathService filePath
             return sourceText.TryGetAbsoluteIndex(linePosition, out _);
         }
     }
-
-    private static ImmutableArray<ClassifiedSpanInternal> GetClassifiedSpans(RazorCodeDocument document)
-    {
-        // Since this service is called so often, we get a good performance improvement by caching these values
-        // for this code document. If the document changes, as the user types, then the document instance will be
-        // different, so we don't need to worry about invalidating the cache.
-        if (!document.Items.TryGetValue(typeof(ClassifiedSpanInternal), out ImmutableArray<ClassifiedSpanInternal> classifiedSpans))
-        {
-            var syntaxTree = document.GetSyntaxTree();
-            classifiedSpans = ClassifiedSpanVisitor.VisitRoot(syntaxTree);
-
-            document.Items[typeof(ClassifiedSpanInternal)] = classifiedSpans;
-        }
-
-        return classifiedSpans;
-    }
-
-    private static ImmutableArray<TagHelperSpanInternal> GetTagHelperSpans(RazorCodeDocument document)
-    {
-        // Since this service is called so often, we get a good performance improvement by caching these values
-        // for this code document. If the document changes, as the user types, then the document instance will be
-        // different, so we don't need to worry about invalidating the cache.
-        if (!document.Items.TryGetValue(typeof(TagHelperSpanInternal), out ImmutableArray<TagHelperSpanInternal> tagHelperSpans))
-        {
-            var syntaxTree = document.GetSyntaxTree();
-            tagHelperSpans = TagHelperSpanVisitor.VisitRoot(syntaxTree);
-
-            document.Items[typeof(TagHelperSpanInternal)] = tagHelperSpans;
-        }
-
-        return tagHelperSpans;
-    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/DocumentMapping/IDocumentMappingService.cs

@@ -3,7 +3,6 @@
 
 using System.Collections.Generic;
 using Microsoft.AspNetCore.Razor.Language;
-using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Text;
 
 namespace Microsoft.CodeAnalysis.Razor.DocumentMapping;
@@ -21,6 +20,4 @@ internal interface IDocumentMappingService
     bool TryMapToGeneratedDocumentPosition(IRazorGeneratedDocument generatedDocument, int hostDocumentIndex, out LinePosition generatedPosition, out int generatedIndex);
 
     bool TryMapToGeneratedDocumentOrNextCSharpPosition(IRazorGeneratedDocument generatedDocument, int hostDocumentIndex, out LinePosition generatedPosition, out int generatedIndex);
-
-    RazorLanguageKind GetLanguageKind(RazorCodeDocument codeDocument, int hostDocumentIndex, bool rightAssociative);
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/DocumentMapping/IDocumentMappingServiceExtensions.cs

@@ -48,7 +48,7 @@ internal static class IDocumentMappingServiceExtensions
         var sourceText = codeDocument.Source.Text;
         var position = sourceText.GetPosition(hostDocumentIndex);
 
-        var languageKind = service.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: false);
         if (languageKind is not RazorLanguageKind.Razor)
         {
             var generatedDocument = languageKind is RazorLanguageKind.CSharp

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Extensions/RazorCodeDocumentExtensions.cs

@@ -2,10 +2,12 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Immutable;
 using System.Diagnostics;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using Microsoft.AspNetCore.Razor.Language.Intermediate;
+using Microsoft.AspNetCore.Razor.Language.Legacy;
 using Microsoft.CodeAnalysis.Razor;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Workspaces;
@@ -139,4 +141,142 @@ internal static class RazorCodeDocumentExtensions
 
         return namespaceNode.Content == fullyQualifiedNamespace;
     }
+
+    public static RazorLanguageKind GetLanguageKind(this RazorCodeDocument codeDocument, int hostDocumentIndex, bool rightAssociative)
+    {
+        var classifiedSpans = GetClassifiedSpans(codeDocument);
+        var tagHelperSpans = GetTagHelperSpans(codeDocument);
+        var documentLength = codeDocument.Source.Text.Length;
+
+        return GetLanguageKindCore(classifiedSpans, tagHelperSpans, hostDocumentIndex, documentLength, rightAssociative);
+    }
+
+    private static ImmutableArray<ClassifiedSpanInternal> GetClassifiedSpans(RazorCodeDocument document)
+    {
+        // Since this service is called so often, we get a good performance improvement by caching these values
+        // for this code document. If the document changes, as the user types, then the document instance will be
+        // different, so we don't need to worry about invalidating the cache.
+        if (!document.Items.TryGetValue(typeof(ClassifiedSpanInternal), out ImmutableArray<ClassifiedSpanInternal> classifiedSpans))
+        {
+            var syntaxTree = document.GetSyntaxTree();
+            classifiedSpans = syntaxTree.GetClassifiedSpans();
+
+            document.Items[typeof(ClassifiedSpanInternal)] = classifiedSpans;
+        }
+
+        return classifiedSpans;
+    }
+
+    private static ImmutableArray<TagHelperSpanInternal> GetTagHelperSpans(RazorCodeDocument document)
+    {
+        // Since this service is called so often, we get a good performance improvement by caching these values
+        // for this code document. If the document changes, as the user types, then the document instance will be
+        // different, so we don't need to worry about invalidating the cache.
+        if (!document.Items.TryGetValue(typeof(TagHelperSpanInternal), out ImmutableArray<TagHelperSpanInternal> tagHelperSpans))
+        {
+            var syntaxTree = document.GetSyntaxTree();
+            tagHelperSpans = syntaxTree.GetTagHelperSpans();
+
+            document.Items[typeof(TagHelperSpanInternal)] = tagHelperSpans;
+        }
+
+        return tagHelperSpans;
+    }
+
+    private static RazorLanguageKind GetLanguageKindCore(
+        ImmutableArray<ClassifiedSpanInternal> classifiedSpans,
+        ImmutableArray<TagHelperSpanInternal> tagHelperSpans,
+        int hostDocumentIndex,
+        int hostDocumentLength,
+        bool rightAssociative)
+    {
+        var length = classifiedSpans.Length;
+        for (var i = 0; i < length; i++)
+        {
+            var classifiedSpan = classifiedSpans[i];
+            var span = classifiedSpan.Span;
+
+            if (span.AbsoluteIndex <= hostDocumentIndex)
+            {
+                var end = span.AbsoluteIndex + span.Length;
+                if (end >= hostDocumentIndex)
+                {
+                    if (end == hostDocumentIndex)
+                    {
+                        // We're at an edge.
+
+                        if (classifiedSpan.SpanKind is SpanKindInternal.MetaCode or SpanKindInternal.Transition)
+                        {
+                            // If we're on an edge of a transition of some kind (MetaCode representing an open or closing piece of syntax such as <|,
+                            // and Transition representing an explicit transition to/from razor syntax, such as @|), prefer to classify to the span
+                            // to the right to better represent where the user clicks
+                            continue;
+                        }
+
+                        // If we're right associative, then we don't want to use the classification that we're at the end
+                        // of, if we're also at the start of the next one
+                        if (rightAssociative)
+                        {
+                            if (i < classifiedSpans.Length - 1 && classifiedSpans[i + 1].Span.AbsoluteIndex == hostDocumentIndex)
+                            {
+                                // If we're at the start of the next span, then use that span
+                                return GetLanguageFromClassifiedSpan(classifiedSpans[i + 1]);
+                            }
+
+                            // Otherwise, we did not find a match using right associativity, so check for tag helpers
+                            break;
+                        }
+                    }
+
+                    return GetLanguageFromClassifiedSpan(classifiedSpan);
+                }
+            }
+        }
+
+        foreach (var tagHelperSpan in tagHelperSpans)
+        {
+            var span = tagHelperSpan.Span;
+
+            if (span.AbsoluteIndex <= hostDocumentIndex)
+            {
+                var end = span.AbsoluteIndex + span.Length;
+                if (end >= hostDocumentIndex)
+                {
+                    if (end == hostDocumentIndex)
+                    {
+                        // We're at an edge. TagHelper spans never own their edge and aren't represented by marker spans
+                        continue;
+                    }
+
+                    // Found intersection
+                    return RazorLanguageKind.Html;
+                }
+            }
+        }
+
+        // Use the language of the last classified span if we're at the end
+        // of the document.
+        if (classifiedSpans.Length != 0 && hostDocumentIndex == hostDocumentLength)
+        {
+            var lastClassifiedSpan = classifiedSpans.Last();
+            return GetLanguageFromClassifiedSpan(lastClassifiedSpan);
+        }
+
+        // Default to Razor
+        return RazorLanguageKind.Razor;
+
+        static RazorLanguageKind GetLanguageFromClassifiedSpan(ClassifiedSpanInternal classifiedSpan)
+        {
+            // Overlaps with request
+            return classifiedSpan.SpanKind switch
+            {
+                SpanKindInternal.Markup => RazorLanguageKind.Html,
+                SpanKindInternal.Code => RazorLanguageKind.CSharp,
+
+                // Content type was non-C# or Html or we couldn't find a classified span overlapping the request position.
+                // All other classified span kinds default back to Razor
+                _ => RazorLanguageKind.Razor,
+            };
+        }
+    }
 }

src/Razor/src/Microsoft.CodeAnalysis.Razor.Workspaces/Formatting/RazorFormattingService.cs

@@ -31,7 +31,6 @@ internal class RazorFormattingService : IRazorFormattingService
     private static readonly FrozenSet<string> s_htmlTriggerCharacterSet = FrozenSet.ToFrozenSet(["\n", "{", "}", ";"], StringComparer.Ordinal);
 
     private readonly IFormattingCodeDocumentProvider _codeDocumentProvider;
-    private readonly IDocumentMappingService _documentMappingService;
     private readonly IAdhocWorkspaceFactory _workspaceFactory;
 
     private readonly ImmutableArray<IFormattingPass> _documentFormattingPasses;
@@ -46,7 +45,6 @@ internal class RazorFormattingService : IRazorFormattingService
         ILoggerFactory loggerFactory)
     {
         _codeDocumentProvider = codeDocumentProvider;
-        _documentMappingService = documentMappingService;
         _workspaceFactory = workspaceFactory;
 
         _htmlOnTypeFormattingPass = new HtmlOnTypeFormattingPass(loggerFactory);
@@ -200,7 +198,7 @@ internal class RazorFormattingService : IRazorFormattingService
 
     public bool TryGetOnTypeFormattingTriggerKind(RazorCodeDocument codeDocument, int hostDocumentIndex, string triggerCharacter, out RazorLanguageKind triggerCharacterKind)
     {
-        triggerCharacterKind = _documentMappingService.GetLanguageKind(codeDocument, hostDocumentIndex, rightAssociative: false);
+        triggerCharacterKind = codeDocument.GetLanguageKind(hostDocumentIndex, rightAssociative: false);
 
         return triggerCharacterKind switch
         {

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/DocumentHighlight/RemoteDocumentHighlightService.cs

@@ -53,7 +53,7 @@ internal sealed partial class RemoteDocumentHighlightService(in ServiceArgs args
 
         var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
 
-        var languageKind = _documentMappingService.GetLanguageKind(codeDocument, index, rightAssociative: true);
+        var languageKind = codeDocument.GetLanguageKind(index, rightAssociative: true);
         if (languageKind is RazorLanguageKind.Html)
         {
             return Response.CallHtml;

src/Razor/src/Microsoft.CodeAnalysis.Remote.Razor/UriPresentation/RemoteUriPresentationService.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.CodeAnalysis.ExternalAccess.Razor;
 using Microsoft.CodeAnalysis.Razor.DocumentPresentation;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
@@ -51,7 +52,7 @@ internal sealed partial class RemoteUriPresentationService(in ServiceArgs args)
 
         var codeDocument = await context.GetCodeDocumentAsync(cancellationToken).ConfigureAwait(false);
 
-        var languageKind = DocumentMappingService.GetLanguageKind(codeDocument, index, rightAssociative: true);
+        var languageKind = codeDocument.GetLanguageKind(index, rightAssociative: true);
         if (languageKind is not RazorLanguageKind.Html)
         {
             // Roslyn doesn't currently support Uri presentation, and whilst it might seem counter intuitive,

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/CodeActions/CodeActionEndpointTest.cs

@@ -12,8 +12,8 @@ using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.CodeActions.Models;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
+using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.AspNetCore.Razor.Test.Common.LanguageServer;
-using Microsoft.AspNetCore.Razor.Threading;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Protocol.CodeActions;
@@ -26,49 +26,17 @@ using Xunit.Abstractions;
 
 namespace Microsoft.AspNetCore.Razor.LanguageServer.CodeActions;
 
-public class CodeActionEndpointTest : LanguageServerTestBase
+public class CodeActionEndpointTest(ITestOutputHelper testOutput) : LanguageServerTestBase(testOutput)
 {
-    private readonly IDocumentMappingService _documentMappingService;
-    private readonly LanguageServerFeatureOptions _languageServerFeatureOptions;
-    private readonly IClientConnection _clientConnection;
-
-    public CodeActionEndpointTest(ITestOutputHelper testOutput)
-        : base(testOutput)
-    {
-        _documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.TryMapToGeneratedDocumentRange(
-                It.IsAny<IRazorGeneratedDocument>(),
-                It.IsAny<LinePositionSpan>(),
-                out It.Ref<LinePositionSpan>.IsAny) == false &&
-
-                s.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp,
-
-            MockBehavior.Strict);
-
-        _languageServerFeatureOptions = Mock.Of<LanguageServerFeatureOptions>(
-            l => l.SupportsFileManipulation == true,
-            MockBehavior.Strict);
-
-        _clientConnection = Mock.Of<IClientConnection>(MockBehavior.Strict);
-    }
+    private static readonly LinePositionSpan s_defaultRange = new(new(5, 2), new(5, 2));
 
     [Fact]
     public async Task Handle_NoDocument()
     {
         // Arrange
         var documentPath = new Uri("C:/path/to/Page.razor");
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            Array.Empty<IRazorCodeActionProvider>(),
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint();
+
         var request = new VSCodeActionParams()
         {
             TextDocument = new VSTextDocumentIdentifier { Uri = documentPath },
@@ -79,7 +47,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var requestContext = CreateRazorRequestContext(documentContext: null);
 
         // Act
-        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, default);
+        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, DisposalToken);
 
         // Assert
         Assert.Null(commandOrCodeActionContainer);
@@ -93,28 +61,19 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
         codeDocument.SetUnsupported();
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            Array.Empty<IRazorCodeActionProvider>(),
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint();
+
         var request = new VSCodeActionParams()
         {
             TextDocument = new VSTextDocumentIdentifier { Uri = documentPath },
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
-        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, default);
+        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, DisposalToken);
 
         // Assert
         Assert.Null(commandOrCodeActionContainer);
@@ -127,31 +86,23 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            Array.Empty<IRazorCodeActionProvider>(),
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint();
+
         var request = new VSCodeActionParams()
         {
             TextDocument = new VSTextDocumentIdentifier { Uri = documentPath },
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
-        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, default);
+        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, DisposalToken);
 
         // Assert
-        Assert.Empty(commandOrCodeActionContainer!);
+        Assert.NotNull(commandOrCodeActionContainer);
+        Assert.Empty(commandOrCodeActionContainer);
     }
 
     [Fact]
@@ -161,20 +112,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider()
-            },
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint(razorCodeActionProviders: [CreateRazorCodeActionProvider()]);
 
         var request = new VSCodeActionParams()
         {
@@ -182,10 +120,11 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
-        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, default);
+        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, DisposalToken);
 
         // Assert
         Assert.NotNull(commandOrCodeActionContainer);
@@ -199,22 +138,11 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var documentMappingService = CreateDocumentMappingService();
-        var languageServer = CreateLanguageServer();
-        var codeActionEndpoint = new CodeActionEndpoint(
-            documentMappingService,
-            Array.Empty<IRazorCodeActionProvider>(),
-            new ICSharpCodeActionProvider[] {
-                    new MockCSharpCodeActionProvider()
-            },
-            Array.Empty<IHtmlCodeActionProvider>(),
-            languageServer,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            documentMappingService: CreateDocumentMappingService(s_defaultRange),
+            csharpCodeActionProviders: [CreateCSharpCodeActionProvider()],
+            clientConnection: TestClientConnection.Instance);
 
         var request = new VSCodeActionParams()
         {
@@ -222,10 +150,11 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
-        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, default);
+        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, DisposalToken);
 
         // Assert
         Assert.NotNull(commandOrCodeActionContainer);
@@ -239,20 +168,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockMultipleRazorCodeActionProvider(),
-            },
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint(razorCodeActionProviders: [CreateMultipleRazorCodeActionProvider()]);
 
         var request = new VSCodeActionParams()
         {
@@ -260,10 +176,11 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
-        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, default);
+        var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, DisposalToken);
 
         // Assert
         Assert.NotNull(commandOrCodeActionContainer);
@@ -277,27 +194,17 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var documentMappingService = CreateDocumentMappingService();
-        var languageServer = CreateLanguageServer();
-        var codeActionEndpoint = new CodeActionEndpoint(
-            documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockMultipleRazorCodeActionProvider(),
-                    new MockMultipleRazorCodeActionProvider(),
-                    new MockRazorCodeActionProvider(),
-            },
-            new ICSharpCodeActionProvider[] {
-                    new MockCSharpCodeActionProvider(),
-                    new MockCSharpCodeActionProvider()
-            },
-            Array.Empty<IHtmlCodeActionProvider>(),
-            languageServer,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            documentMappingService: CreateDocumentMappingService(s_defaultRange),
+            razorCodeActionProviders: [
+                CreateMultipleRazorCodeActionProvider(),
+                CreateMultipleRazorCodeActionProvider(),
+                CreateRazorCodeActionProvider()],
+            csharpCodeActionProviders: [
+                CreateCSharpCodeActionProvider(),
+                CreateCSharpCodeActionProvider()],
+            clientConnection: TestClientConnection.Instance);
 
         var request = new VSCodeActionParams()
         {
@@ -322,27 +229,17 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var documentMappingService = CreateDocumentMappingService();
-        var languageServer = CreateLanguageServer();
-        var codeActionEndpoint = new CodeActionEndpoint(
-            documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider(),
-                    new MockRazorCodeActionProvider(),
-                    new MockRazorCodeActionProvider(),
-            },
-            new ICSharpCodeActionProvider[] {
-                    new MockCSharpCodeActionProvider(),
-                    new MockCSharpCodeActionProvider()
-            },
-            Array.Empty<IHtmlCodeActionProvider>(),
-            languageServer,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            documentMappingService: CreateDocumentMappingService(s_defaultRange),
+            razorCodeActionProviders: [
+                CreateRazorCodeActionProvider(),
+                CreateRazorCodeActionProvider(),
+                CreateRazorCodeActionProvider()],
+            csharpCodeActionProviders: [
+                CreateCSharpCodeActionProvider(),
+                CreateCSharpCodeActionProvider()],
+            clientConnection: TestClientConnection.Instance);
 
         var request = new VSCodeActionParams()
         {
@@ -350,6 +247,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -367,20 +265,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockEmptyRazorCodeActionProvider()
-            },
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint(razorCodeActionProviders: [CreateEmptyRazorCodeActionProvider()]);
 
         var request = new VSCodeActionParams()
         {
@@ -388,13 +273,15 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
         var commandOrCodeActionContainer = await codeActionEndpoint.HandleRequestAsync(request, requestContext, default);
 
         // Assert
-        Assert.Empty(commandOrCodeActionContainer!);
+        Assert.NotNull(commandOrCodeActionContainer);
+        Assert.Empty(commandOrCodeActionContainer);
     }
 
     [Fact]
@@ -404,28 +291,18 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var documentMappingService = CreateDocumentMappingService();
-        var languageServer = CreateLanguageServer();
-        var codeActionEndpoint = new CodeActionEndpoint(
-            documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider(),
-                    new MockEmptyRazorCodeActionProvider(),
-                    new MockRazorCodeActionProvider(),
-                    new MockEmptyRazorCodeActionProvider(),
-            },
-            new ICSharpCodeActionProvider[] {
-                    new MockCSharpCodeActionProvider(),
-                    new MockCSharpCodeActionProvider()
-            },
-            Array.Empty<IHtmlCodeActionProvider>(),
-            languageServer,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            documentMappingService: CreateDocumentMappingService(s_defaultRange),
+            razorCodeActionProviders: [
+                CreateRazorCodeActionProvider(),
+                CreateEmptyRazorCodeActionProvider(),
+                CreateRazorCodeActionProvider(),
+                CreateEmptyRazorCodeActionProvider()],
+            csharpCodeActionProviders: [
+                CreateCSharpCodeActionProvider(),
+                CreateCSharpCodeActionProvider()],
+            clientConnection: TestClientConnection.Instance);
 
         var request = new VSCodeActionParams()
         {
@@ -433,6 +310,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -450,22 +328,13 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider(),
-                    new MockRazorCommandProvider(),
-                    new MockEmptyRazorCodeActionProvider()
-            },
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = true
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            razorCodeActionProviders: [
+                CreateRazorCodeActionProvider(),
+                CreateRazorCommandProvider(),
+                CreateEmptyRazorCodeActionProvider()],
+            supportsCodeActionResolve: true);
 
         var request = new VSCodeActionParams()
         {
@@ -473,6 +342,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -500,24 +370,13 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var documentMappingService = CreateDocumentMappingService();
-        var languageServer = CreateLanguageServer();
-        var codeActionEndpoint = new CodeActionEndpoint(
-            documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider(),
-            },
-            new ICSharpCodeActionProvider[] {
-                    new MockCSharpCodeActionProvider()
-            },
-            Array.Empty<IHtmlCodeActionProvider>(),
-            languageServer,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = true
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            documentMappingService: CreateDocumentMappingService(s_defaultRange),
+            razorCodeActionProviders: [CreateRazorCodeActionProvider()],
+            csharpCodeActionProviders: [CreateCSharpCodeActionProvider()],
+            clientConnection: TestClientConnection.Instance,
+            supportsCodeActionResolve: true);
 
         var request = new VSCodeActionParams()
         {
@@ -525,6 +384,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -554,22 +414,12 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider(),
-                    new MockRazorCommandProvider(),
-                    new MockEmptyRazorCodeActionProvider()
-            },
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            razorCodeActionProviders: [
+                CreateRazorCodeActionProvider(),
+                CreateRazorCommandProvider(),
+                CreateEmptyRazorCodeActionProvider()]);
 
         var request = new VSCodeActionParams()
         {
@@ -577,6 +427,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             Range = VsLspFactory.CreateZeroWidthRange(0, 1),
             Context = new VSInternalCodeActionContext()
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -587,14 +438,15 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         Assert.Collection(commandOrCodeActionContainer,
             c =>
             {
-                Assert.True(c.TryGetFirst(out var command1));
-                var command = Assert.IsType<Command>(command1);
-                var codeActionParamsToken = (JsonObject)command.Arguments!.First();
+                Assert.True(c.TryGetFirst(out var first));
+                var command = Assert.IsType<Command>(first);
+                Assert.NotNull(command.Arguments);
+                var codeActionParamsToken = (JsonObject)command.Arguments.First();
                 var codeActionParams = codeActionParamsToken.Deserialize<RazorCodeActionResolutionParams>();
                 Assert.NotNull(codeActionParams);
                 Assert.Equal(LanguageServerConstants.CodeActions.EditBasedCodeActionCommand, codeActionParams.Action);
             },
-            c => Assert.True(c.TryGetFirst(out var _)));
+            c => Assert.True(c.TryGetFirst(out _)));
     }
 
     [Fact]
@@ -604,20 +456,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider()
-            },
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint(razorCodeActionProviders: [CreateRazorCodeActionProvider()]);
 
         var initialRange = VsLspFactory.CreateZeroWidthRange(0, 1);
         var selectionRange = VsLspFactory.CreateZeroWidthRange(0, 5);
@@ -646,20 +485,7 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            _documentMappingService,
-            new IRazorCodeActionProvider[] {
-                    new MockRazorCodeActionProvider()
-            },
-            Array.Empty<ICSharpCodeActionProvider>(),
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+        var codeActionEndpoint = CreateEndpoint(razorCodeActionProviders: [CreateRazorCodeActionProvider()]);
 
         var initialRange = VsLspFactory.CreateZeroWidthRange(0, 1);
         var request = new VSCodeActionParams()
@@ -687,24 +513,8 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var documentPath = new Uri("C:/path/to/Page.razor");
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
-        LinePositionSpan projectedRange = default;
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            d => d.TryMapToGeneratedDocumentRange(It.IsAny<IRazorGeneratedDocument>(), It.IsAny<LinePositionSpan>(), out projectedRange) == false
-        , MockBehavior.Strict);
-        var codeActionEndpoint = new CodeActionEndpoint(
-            documentMappingService,
-            Array.Empty<IRazorCodeActionProvider>(),
-            new ICSharpCodeActionProvider[] {
-                    new MockCSharpCodeActionProvider()
-            },
-            Array.Empty<IHtmlCodeActionProvider>(),
-            _clientConnection,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+
+        var codeActionEndpoint = CreateEndpoint(csharpCodeActionProviders: [CreateCSharpCodeActionProvider()]);
 
         var initialRange = VsLspFactory.CreateZeroWidthRange(0, 1);
         var request = new VSCodeActionParams()
@@ -733,22 +543,11 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         var codeDocument = CreateCodeDocument("@code {}");
         var documentContext = CreateDocumentContext(documentPath, codeDocument);
         var projectedRange = VsLspFactory.CreateZeroWidthRange(15, 2);
-        var documentMappingService = CreateDocumentMappingService(projectedRange.ToLinePositionSpan());
-        var languageServer = CreateLanguageServer();
-        var codeActionEndpoint = new CodeActionEndpoint(
-            documentMappingService,
-            Array.Empty<IRazorCodeActionProvider>(),
-            new ICSharpCodeActionProvider[] {
-                    new MockCSharpCodeActionProvider()
-            },
-            Array.Empty<IHtmlCodeActionProvider>(),
-            languageServer,
-            _languageServerFeatureOptions,
-            LoggerFactory,
-            telemetryReporter: null)
-        {
-            _supportsCodeActionResolve = false
-        };
+
+        var codeActionEndpoint = CreateEndpoint(
+            documentMappingService: CreateDocumentMappingService(projectedRange.ToLinePositionSpan()),
+            csharpCodeActionProviders: [CreateCSharpCodeActionProvider()],
+            clientConnection: TestClientConnection.Instance);
 
         var initialRange = VsLspFactory.CreateZeroWidthRange(0, 1);
         var request = new VSCodeActionParams()
@@ -765,11 +564,12 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         Assert.NotNull(context);
 
         // Act
-        var results = await codeActionEndpoint.GetCodeActionsFromLanguageServerAsync(RazorLanguageKind.CSharp, documentContext, context, Guid.Empty, cancellationToken: default);
+        var results = await codeActionEndpoint.GetCodeActionsFromLanguageServerAsync(RazorLanguageKind.CSharp, documentContext, context, Guid.Empty, cancellationToken: DisposalToken);
 
         // Assert
         var result = Assert.Single(results);
-        var diagnostics = result.Diagnostics!.ToArray();
+        Assert.NotNull(result.Diagnostics);
+        var diagnostics = result.Diagnostics.ToArray();
         Assert.Equal(2, diagnostics.Length);
 
         // Diagnostic ranges contain the projected range for
@@ -782,136 +582,142 @@ public class CodeActionEndpointTest : LanguageServerTestBase
         Assert.Equal(projectedRange, diagnostics[1].Range);
     }
 
-    private static IDocumentMappingService CreateDocumentMappingService(LinePositionSpan projectedRange = default)
+    private CodeActionEndpoint CreateEndpoint(
+        IDocumentMappingService? documentMappingService = null,
+        ImmutableArray<IRazorCodeActionProvider> razorCodeActionProviders = default,
+        ImmutableArray<ICSharpCodeActionProvider> csharpCodeActionProviders = default,
+        ImmutableArray<IHtmlCodeActionProvider> htmlCodeActionProviders = default,
+        IClientConnection? clientConnection = null,
+        LanguageServerFeatureOptions? languageServerFeatureOptions = null,
+        bool supportsCodeActionResolve = false)
     {
-        if (projectedRange == default)
+        return new CodeActionEndpoint(
+            documentMappingService ?? CreateDocumentMappingService(),
+            razorCodeActionProviders.NullToEmpty(),
+            csharpCodeActionProviders.NullToEmpty(),
+            htmlCodeActionProviders.NullToEmpty(),
+            clientConnection ?? StrictMock.Of<IClientConnection>(),
+            languageServerFeatureOptions ?? StrictMock.Of<LanguageServerFeatureOptions>(x => x.SupportsFileManipulation == true),
+            LoggerFactory,
+            telemetryReporter: null)
         {
-            projectedRange = new LinePositionSpan(new(5, 2), new(5, 2));
-        }
-
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            d => d.TryMapToGeneratedDocumentRange(It.IsAny<IRazorGeneratedDocument>(), It.IsAny<LinePositionSpan>(), out projectedRange) == true &&
-                 d.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp
-        , MockBehavior.Strict);
-        return documentMappingService;
+            _supportsCodeActionResolve = supportsCodeActionResolve
+        };
     }
 
-    private static IClientConnection CreateLanguageServer()
+    private static IDocumentMappingService CreateDocumentMappingService(LinePositionSpan? projectedRange = null)
     {
-        return new TestLanguageServer();
+        var mock = new StrictMock<IDocumentMappingService>();
+
+        // If a range was provided, use that and return true; otherwise, return false.
+        var (outRange, result) = projectedRange is LinePositionSpan
+            ? (projectedRange.GetValueOrDefault(), true)
+            : (It.Ref<LinePositionSpan>.IsAny, false);
+
+        mock.Setup(x => x.TryMapToGeneratedDocumentRange(It.IsAny<IRazorGeneratedDocument>(), It.IsAny<LinePositionSpan>(), out outRange))
+            .Returns(result);
+
+        return mock.Object;
     }
 
     private static RazorCodeDocument CreateCodeDocument(string text)
     {
-        var codeDocument = TestRazorCodeDocument.Create(text);
         var sourceDocument = TestRazorSourceDocument.Create(text);
-        var syntaxTree = RazorSyntaxTree.Parse(sourceDocument);
-        codeDocument.SetSyntaxTree(syntaxTree);
-        return codeDocument;
+        var projectEngine = RazorProjectEngine.Create(builder => { });
+
+        return projectEngine.ProcessDesignTime(sourceDocument, "mvc", importSources: [], tagHelpers: []);
     }
 
-    private class MockRazorCodeActionProvider : IRazorCodeActionProvider
+    private static IRazorCodeActionProvider CreateEmptyRazorCodeActionProvider()
+        => CreateRazorCodeActionProvider([]);
+
+    private static IRazorCodeActionProvider CreateRazorCodeActionProvider()
+        => CreateRazorCodeActionProvider(new RazorVSInternalCodeAction());
+
+    private static IRazorCodeActionProvider CreateMultipleRazorCodeActionProvider()
+        => CreateRazorCodeActionProvider(
+            new RazorVSInternalCodeAction(),
+            new RazorVSInternalCodeAction());
+
+    private static IRazorCodeActionProvider CreateRazorCommandProvider()
+        => CreateRazorCodeActionProvider(
+            new RazorVSInternalCodeAction()
+            {
+                Title = "SomeTitle",
+                Data = JsonSerializer.SerializeToElement(new AddUsingsCodeActionParams()
+                {
+                    Namespace = "Test",
+                    Uri = new Uri("C:/path/to/Page.razor")
+                })
+            });
+
+    private static IRazorCodeActionProvider CreateRazorCodeActionProvider(params ImmutableArray<RazorVSInternalCodeAction> codeActions)
     {
-        public Task<ImmutableArray<RazorVSInternalCodeAction>> ProvideAsync(RazorCodeActionContext context, CancellationToken cancellationToken)
+        var mock = new StrictMock<IRazorCodeActionProvider>();
+
+        mock.Setup(x => x.ProvideAsync(It.IsAny<RazorCodeActionContext>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(() => codeActions);
+
+        return mock.Object;
+    }
+
+    private static ICSharpCodeActionProvider CreateCSharpCodeActionProvider()
+        => CreateCSharpCodeActionProvider([new RazorVSInternalCodeAction()]);
+
+    private static ICSharpCodeActionProvider CreateCSharpCodeActionProvider(params ImmutableArray<RazorVSInternalCodeAction> codeActions)
+    {
+        var mock = new StrictMock<ICSharpCodeActionProvider>();
+
+        mock.Setup(x => x.ProvideAsync(It.IsAny<RazorCodeActionContext>(), It.IsAny<ImmutableArray<RazorVSInternalCodeAction>>(), It.IsAny<CancellationToken>()))
+            .ReturnsAsync(() => codeActions);
+
+        return mock.Object;
+    }
+
+    private sealed class TestClientConnection : IClientConnection
+    {
+        public static readonly IClientConnection Instance = new TestClientConnection();
+
+        private static readonly string[] s_customTags = ["CodeActionName"];
+
+        private TestClientConnection()
         {
-            return Task.FromResult<ImmutableArray<RazorVSInternalCodeAction>>([new RazorVSInternalCodeAction()]);
         }
-    }
 
-    private class MockMultipleRazorCodeActionProvider : IRazorCodeActionProvider
-    {
-        public Task<ImmutableArray<RazorVSInternalCodeAction>> ProvideAsync(RazorCodeActionContext context, CancellationToken cancellationToken)
-        {
-            return Task.FromResult<ImmutableArray<RazorVSInternalCodeAction>>(
-                [
-                    new RazorVSInternalCodeAction(),
-                    new RazorVSInternalCodeAction()
-                ]);
-        }
-    }
-
-    private class MockCSharpCodeActionProvider : ICSharpCodeActionProvider
-    {
-        public Task<ImmutableArray<RazorVSInternalCodeAction>> ProvideAsync(RazorCodeActionContext _1, ImmutableArray<RazorVSInternalCodeAction> _2, CancellationToken _3)
-        {
-            return Task.FromResult<ImmutableArray<RazorVSInternalCodeAction>>(
-                [
-                    new RazorVSInternalCodeAction()
-                ]);
-        }
-    }
-
-    private class MockRazorCommandProvider : IRazorCodeActionProvider
-    {
-        public Task<ImmutableArray<RazorVSInternalCodeAction>> ProvideAsync(RazorCodeActionContext _1, CancellationToken _2)
-        {
-            // O# Code Actions don't have `Data`, but `Commands` do
-            return Task.FromResult<ImmutableArray<RazorVSInternalCodeAction>>(
-                [
-                    new RazorVSInternalCodeAction() {
-                        Title = "SomeTitle",
-                        Data = JsonSerializer.SerializeToElement(new AddUsingsCodeActionParams()
-                        {
-                            Namespace="Test",
-                            Uri = new Uri("C:/path/to/Page.razor")
-                        })
-                    }
-                ]);
-        }
-    }
-
-    private class MockEmptyRazorCodeActionProvider : IRazorCodeActionProvider
-    {
-        public Task<ImmutableArray<RazorVSInternalCodeAction>> ProvideAsync(RazorCodeActionContext _1, CancellationToken _2)
-        {
-            return SpecializedTasks.EmptyImmutableArray<RazorVSInternalCodeAction>();
-        }
-    }
-
-    private class TestLanguageServer : IClientConnection
-    {
         public Task SendNotificationAsync<TParams>(string method, TParams @params, CancellationToken cancellationToken)
         {
-            if (method != CustomMessageNames.RazorProvideCodeActionsEndpoint)
-            {
-                throw new InvalidOperationException($"Unexpected method {method}");
-            }
+            Assert.Equal(CustomMessageNames.RazorProvideCodeActionsEndpoint, method);
 
             return Task.CompletedTask;
         }
 
         public Task SendNotificationAsync(string method, CancellationToken cancellationToken)
-        {
-            throw new NotImplementedException();
-        }
+            => throw new NotImplementedException();
 
         public Task<TResponse> SendRequestAsync<TParams, TResponse>(string method, TParams @params, CancellationToken cancellationToken)
         {
-            if (method != CustomMessageNames.RazorProvideCodeActionsEndpoint)
-            {
-                throw new InvalidOperationException($"Unexpected method {method}");
-            }
+            Assert.Equal(CustomMessageNames.RazorProvideCodeActionsEndpoint, method);
 
-            if (@params is not DelegatedCodeActionParams delegatedCodeActionParams ||
-                delegatedCodeActionParams.CodeActionParams is not VSCodeActionParams codeActionParams ||
-                codeActionParams.Context is not VSInternalCodeActionContext codeActionContext)
-            {
-                throw new InvalidOperationException(@params!.GetType().FullName);
-            }
+            Assert.NotNull(@params);
+            var delegatedCodeActionParams = Assert.IsType<DelegatedCodeActionParams>(@params);
+
+            Assert.NotNull(delegatedCodeActionParams.CodeActionParams);
+            Assert.NotNull(delegatedCodeActionParams.CodeActionParams.Context);
 
             var diagnostics = new List<Diagnostic>
             {
-                new Diagnostic()
+                new()
                 {
-                    Range = codeActionParams.Range,
+                    Range = delegatedCodeActionParams.CodeActionParams.Range,
                     Message = "Range"
                 }
             };
-            if (codeActionContext.SelectionRange is not null)
+
+            if (delegatedCodeActionParams.CodeActionParams.Context.SelectionRange is { } selectionRange)
             {
-                diagnostics.Add(new Diagnostic()
+                diagnostics.Add(new()
                 {
-                    Range = codeActionContext.SelectionRange,
+                    Range = selectionRange,
                     Message = "Selection Range"
                 });
             }
@@ -922,12 +728,12 @@ public class CodeActionEndpointTest : LanguageServerTestBase
             // is correct rather than providing specific test hooks in the CodeActionEndpoint
             var result = new[]
             {
-                    new RazorVSInternalCodeAction()
-                    {
-                        Data = JsonSerializer.SerializeToElement(new { CustomTags = new object[] { "CodeActionName" } }),
-                        Diagnostics = diagnostics.ToArray()
-                    }
-                };
+                new RazorVSInternalCodeAction()
+                {
+                    Data = JsonSerializer.SerializeToElement(new { CustomTags = s_customTags }),
+                    Diagnostics = [.. diagnostics]
+                }
+            };
 
             return Task.FromResult((TResponse)(object)result);
         }

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/DocumentPresentation/TextDocumentTextPresentationEndpointTests.cs

@@ -2,15 +2,14 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
+using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.AspNetCore.Razor.Test.Common.LanguageServer;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Protocol.DocumentPresentation;
-using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
 using Moq;
 using Xunit;
@@ -24,116 +23,83 @@ public class TextDocumentTextPresentationEndpointTests(ITestOutputHelper testOut
     public async Task Handle_Html_MakesRequest()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
+        TestCode code = "<[|d|]iv></div>";
+
+        var codeDocument = CreateCodeDocument(code.Text);
 
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorTextPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response)
-            .Verifiable();
-
-        var endpoint = new TextDocumentTextPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null, verifiable: true);
+        var endpoint = CreateEndpoint(clientConnection);
 
         var parameters = new TextPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
-            Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
+            TextDocument = new() { Uri = uri },
+            Range = codeDocument.Source.Text.GetRange(code.Span),
             Text = "Hi there"
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
-        var result = await endpoint.HandleRequestAsync(parameters, requestContext, DisposalToken);
+        _ = await endpoint.HandleRequestAsync(parameters, requestContext, DisposalToken);
 
         // Assert
-        clientConnection.Verify();
+        Mock.Get(clientConnection).Verify();
     }
 
     [Fact]
     public async Task Handle_CSharp_DoesNotMakeRequest()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("@counter");
-        var csharpDocument = codeDocument.GetCSharpDocument();
+        TestCode code = "@[|c|]ounter";
+
+        var codeDocument = CreateCodeDocument(code.Text);
+
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var projectedRange = It.IsAny<LinePositionSpan>();
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp &&
-            s.TryMapToGeneratedDocumentRange(csharpDocument, It.IsAny<LinePositionSpan>(), out projectedRange) == true, MockBehavior.Strict);
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-
-        var endpoint = new TextDocumentTextPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            LoggerFactory);
+        var clientConnection = StrictMock.Of<IClientConnection>();
+        var endpoint = CreateEndpoint(clientConnection);
 
         var parameters = new TextPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
-            Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
+            TextDocument = new() { Uri = uri },
+            Range = codeDocument.Source.Text.GetRange(code.Span),
             Text = "Hi there"
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
-        var result = await endpoint.HandleRequestAsync(parameters, requestContext, DisposalToken);
+        _ = await endpoint.HandleRequestAsync(parameters, requestContext, DisposalToken);
 
         // Assert
-        clientConnection.Verify();
+        Mock.Get(clientConnection)
+            .VerifySendRequest<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorTextPresentationEndpoint, Times.Never);
     }
 
     [Fact]
     public async Task Handle_DocumentNotFound_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
+        TestCode code = "<[|d|]iv></div>";
+
+        var codeDocument = CreateCodeDocument(code.Text);
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorTextPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentTextPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null);
+        var endpoint = CreateEndpoint(clientConnection);
 
         var parameters = new TextPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
-            Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
+            TextDocument = new() { Uri = uri },
+            Range = codeDocument.Source.Text.GetRange(code.Span),
             Text = "Hi there"
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -147,35 +113,24 @@ public class TextDocumentTextPresentationEndpointTests(ITestOutputHelper testOut
     public async Task Handle_UnsupportedCodeDocument_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
+        TestCode code = "<[|d|]iv></div>";
+
+        var codeDocument = CreateCodeDocument(code.Text);
         codeDocument.SetUnsupported();
+
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
 
-        var response = new WorkspaceEdit();
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorTextPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentTextPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: new WorkspaceEdit());
+        var endpoint = CreateEndpoint(clientConnection);
 
         var parameters = new TextPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
-            Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
+            TextDocument = new() { Uri = uri },
+            Range = codeDocument.Source.Text.GetRange(code.Span),
             Text = "Hi there"
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -184,4 +139,13 @@ public class TextDocumentTextPresentationEndpointTests(ITestOutputHelper testOut
         // Assert
         Assert.Null(result);
     }
+
+    private TextDocumentTextPresentationEndpoint CreateEndpoint(IClientConnection clientConnection)
+        => new(StrictMock.Of<IDocumentMappingService>(), clientConnection, FilePathService, LoggerFactory);
+
+    private static IClientConnection CreateClientConnection(WorkspaceEdit? response, bool verifiable = false)
+        => TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorTextPresentationEndpoint, response, verifiable);
+        });
 }

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/DocumentPresentation/TextDocumentUriPresentationEndpointTests.cs

@@ -2,11 +2,11 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
-using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
 using Microsoft.AspNetCore.Razor.ProjectSystem;
+using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.AspNetCore.Razor.Test.Common.LanguageServer;
 using Microsoft.AspNetCore.Razor.Test.Common.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.DocumentMapping;
@@ -15,7 +15,6 @@ using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.CodeAnalysis.Razor.Protocol.DocumentPresentation;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
-using Moq;
 using Xunit;
 using Xunit.Abstractions;
 using static Microsoft.AspNetCore.Razor.Language.CommonMetadata;
@@ -30,46 +29,43 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
         // Arrange
         var projectManager = CreateProjectSnapshotManager();
 
-        var project = await projectManager.UpdateAsync(updater => updater.CreateAndAddProject("c:/path/project.csproj"));
+        var project = await projectManager.UpdateAsync(updater =>
+        {
+            return updater.CreateAndAddProject("c:/path/project.csproj");
+        });
+
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/index.razor");
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/MyTagHelper.razor");
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-
         var droppedUri = new Uri("file:///c:/path/MyTagHelper.razor");
         var builder = TagHelperDescriptorBuilder.Create("MyTagHelper", "MyAssembly");
         builder.SetMetadata(TypeNameIdentifier("MyTagHelper"), TypeNamespace("TestRootNamespace"));
-        var tagHelperDescriptor = builder.Build();
 
-        await projectManager.UpdateAsync(updater => updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([tagHelperDescriptor])));
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([builder.Build()]));
+        });
 
         var razorFilePath = "c:/path/index.razor";
         var uri = new Uri(razorFilePath);
 
-        await projectManager.UpdateAsync(updater => updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>")));
-        var documentSnapshot = projectManager.GetLoadedProject(project.Key).GetDocument(razorFilePath).AssumeNotNull();
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>"));
+        });
+
         var documentContextFactory = new DocumentContextFactory(projectManager, LoggerFactory);
-        Assert.True(documentContextFactory.TryCreate(uri, null, out var documentContext));
+        Assert.True(documentContextFactory.TryCreate(uri, projectContext: null, out var documentContext));
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var endpoint = CreateEndpoint(documentContextFactory);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
             Uris = [droppedUri]
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -77,7 +73,12 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
 
         // Assert
         Assert.NotNull(result);
-        Assert.Equal("<MyTagHelper />", result.DocumentChanges!.Value.First[0].Edits[0].NewText);
+        Assert.NotNull(result.DocumentChanges);
+
+        var documentChanges = result.DocumentChanges.GetValueOrDefault();
+        Assert.True(documentChanges.TryGetFirst(out var documentEdits));
+
+        Assert.Equal("<MyTagHelper />", documentEdits[0].Edits[0].NewText);
     }
 
     [OSSkipConditionFact(["OSX", "Linux"])]
@@ -86,43 +87,39 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
         // Arrange
         var projectManager = CreateProjectSnapshotManager();
 
-        var project = await projectManager.UpdateAsync(updater => updater.CreateAndAddProject("c:/path/project.csproj"));
+        var project = await projectManager.UpdateAsync(updater =>
+        {
+            return updater.CreateAndAddProject("c:/path/project.csproj");
+        });
+
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/index.razor");
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/MyTagHelper.razor");
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-
         var droppedUri = new Uri("file:///c:/path/MyTagHelper.razor");
         var builder = TagHelperDescriptorBuilder.Create("MyTagHelper", "MyAssembly");
         builder.SetMetadata(TypeNameIdentifier("MyTagHelper"), TypeNamespace("TestRootNamespace"));
-        var tagHelperDescriptor = builder.Build();
 
-        await projectManager.UpdateAsync(updater => updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([tagHelperDescriptor])));
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([builder.Build()]));
+        });
 
         var razorFilePath = "c:/path/index.razor";
         var uri = new Uri(razorFilePath);
 
-        await projectManager.UpdateAsync(updater => updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>")));
-        var documentSnapshot = projectManager.GetLoadedProject(project.Key).GetDocument(razorFilePath).AssumeNotNull();
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>"));
+        });
+
         var documentContextFactory = new DocumentContextFactory(projectManager, LoggerFactory);
         Assert.True(documentContextFactory.TryCreate(uri, null, out var documentContext));
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var endpoint = CreateEndpoint(documentContextFactory);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
             Uris =
             [
@@ -131,6 +128,7 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
                 droppedUri,
             ]
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -138,7 +136,12 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
 
         // Assert
         Assert.NotNull(result);
-        Assert.Equal("<MyTagHelper />", result!.DocumentChanges!.Value.First[0].Edits[0].NewText);
+        Assert.NotNull(result.DocumentChanges);
+
+        var documentChanges = result.DocumentChanges.GetValueOrDefault();
+        Assert.True(documentChanges.TryGetFirst(out var documentEdits));
+
+        Assert.Equal("<MyTagHelper />", documentEdits[0].Edits[0].NewText);
     }
 
     [OSSkipConditionFact(["OSX", "Linux"])]
@@ -147,13 +150,14 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
         // Arrange
         var projectManager = CreateProjectSnapshotManager();
 
-        var project = await projectManager.UpdateAsync(updater => updater.CreateAndAddProject("c:/path/project.csproj"));
+        var project = await projectManager.UpdateAsync(updater =>
+        {
+            return updater.CreateAndAddProject("c:/path/project.csproj");
+        });
+
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/index.razor");
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/fetchdata.razor");
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-
         var droppedUri = new Uri("file:///c:/path/fetchdata.razor");
         var builder = TagHelperDescriptorBuilder.Create("FetchData", "MyAssembly");
         builder.SetMetadata(TypeNameIdentifier("FetchData"), TypeNamespace("TestRootNamespace"));
@@ -163,36 +167,32 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
             b.Name = "MyAttribute";
         });
         builder.BindAttribute(b => b.Name = "MyNonRequiredAttribute");
-        var tagHelperDescriptor = builder.Build();
 
-        await projectManager.UpdateAsync(updater => updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([tagHelperDescriptor])));
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([builder.Build()]));
+        });
 
         var razorFilePath = "c:/path/index.razor";
         var uri = new Uri(razorFilePath);
 
-        await projectManager.UpdateAsync(updater => updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>")));
-        var documentSnapshot = projectManager.GetLoadedProject(project.Key).GetDocument(razorFilePath).AssumeNotNull();
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>"));
+        });
+
         var documentContextFactory = new DocumentContextFactory(projectManager, LoggerFactory);
         Assert.True(documentContextFactory.TryCreate(uri, null, out var documentContext));
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var endpoint = CreateEndpoint(documentContextFactory);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
             Uris = [droppedUri]
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -200,51 +200,35 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
 
         // Assert
         Assert.NotNull(result);
-        Assert.Equal("<FetchData MyAttribute=\"\" />", result.DocumentChanges!.Value.First[0].Edits[0].NewText);
+        Assert.NotNull(result.DocumentChanges);
+
+        var documentChanges = result.DocumentChanges.GetValueOrDefault();
+        Assert.True(documentChanges.TryGetFirst(out var documentEdits));
+
+        Assert.Equal("<FetchData MyAttribute=\"\" />", documentEdits[0].Edits[0].NewText);
     }
 
     [Fact]
     public async Task Handle_NoTypeNameIdentifier_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-
-        var componentCodeDocument = TestRazorCodeDocument.Create("<div></div>");
-        var droppedUri = new Uri("file:///c:/path/MyTagHelper.razor");
-        var builder = TagHelperDescriptorBuilder.Create("MyTagHelper", "MyAssembly");
-        var tagHelperDescriptor = builder.Build();
-
-        var documentSnapshot = Mock.Of<IDocumentSnapshot>(s => s.GetGeneratedOutputAsync(It.IsAny<bool>()) == Task.FromResult(componentCodeDocument), MockBehavior.Strict);
-
+        var codeDocument = CreateCodeDocument("<div></div>");
         var uri = new Uri("file://path/test.razor");
+        var droppedUri = new Uri("file:///c:/path/MyTagHelper.razor");
+
         var documentContextFactory = CreateDocumentContextFactory(uri, codeDocument);
         var documentContext = CreateDocumentContext(uri, codeDocument);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null);
+        var endpoint = CreateEndpoint(documentContextFactory, clientConnection);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
             Uris = [droppedUri]
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -258,36 +242,18 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
     public async Task Handle_MultipleUris_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-
-        var documentSnapshot = Mock.Of<IDocumentSnapshot>(s => s.GetGeneratedOutputAsync(false) == Task.FromResult(codeDocument), MockBehavior.Strict);
+        var codeDocument = CreateCodeDocument("<div></div>");
 
         var uri = new Uri("file://path/test.razor");
         var documentContextFactory = CreateDocumentContextFactory(uri, codeDocument);
         var documentContext = CreateDocumentContext(uri, codeDocument);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null);
+        var endpoint = CreateEndpoint(documentContextFactory, clientConnection);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
             Uris =
             [
@@ -296,6 +262,7 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
                 new Uri("file:///c:/path/MyTagHelper.razor"),
             ]
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -309,40 +276,23 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
     public async Task Handle_NotComponent_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-
-        var documentSnapshot = Mock.Of<IDocumentSnapshot>(s => s.GetGeneratedOutputAsync(false) == Task.FromResult(codeDocument), MockBehavior.Strict);
+        var codeDocument = CreateCodeDocument("<div></div>");
 
         var droppedUri = new Uri("file:///c:/path/MyTagHelper.cshtml");
         var uri = new Uri("file://path/test.razor");
         var documentContextFactory = CreateDocumentContextFactory(uri, codeDocument);
         var documentContext = CreateDocumentContext(uri, codeDocument);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null);
+        var endpoint = CreateEndpoint(documentContextFactory, clientConnection);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
             Uris = [droppedUri]
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -358,47 +308,49 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
         // Arrange
         var projectManager = CreateProjectSnapshotManager();
 
-        var project = await projectManager.UpdateAsync(updater => updater.CreateAndAddProject("c:/path/project.csproj"));
+        var project = await projectManager.UpdateAsync(updater =>
+        {
+            return updater.CreateAndAddProject("c:/path/project.csproj");
+        });
+
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/index.razor");
         await projectManager.CreateAndAddDocumentAsync(project, "c:/path/fetchdata.razor");
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-
         var droppedUri1 = new Uri("file:///c:/path/fetchdata.razor.cs");
         var droppedUri2 = new Uri("file:///c:/path/fetchdata.razor");
         var builder = TagHelperDescriptorBuilder.Create("FetchData", "MyAssembly");
         builder.SetMetadata(TypeNameIdentifier("FetchData"), TypeNamespace("TestRootNamespace"));
-        var tagHelperDescriptor = builder.Build();
 
-        await projectManager.UpdateAsync(updater => updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([tagHelperDescriptor])));
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.ProjectWorkspaceStateChanged(project.Key, ProjectWorkspaceState.Create([builder.Build()]));
+        });
 
         var razorFilePath = "c:/path/index.razor";
         var uri = new Uri(razorFilePath);
 
-        await projectManager.UpdateAsync(updater => updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>")));
-        var documentSnapshot = projectManager.GetLoadedProject(project.Key).GetDocument(razorFilePath).AssumeNotNull();
+        await projectManager.UpdateAsync(updater =>
+        {
+            updater.DocumentOpened(project.Key, razorFilePath, SourceText.From("<div></div>"));
+        });
+
+        var documentSnapshot = projectManager
+            .GetLoadedProject(project.Key)
+            .GetDocument(razorFilePath);
+        Assert.NotNull(documentSnapshot);
+
         var documentContextFactory = new DocumentContextFactory(projectManager, LoggerFactory);
-        Assert.True(documentContextFactory.TryCreate(uri, null, out var documentContext));
+        Assert.True(documentContextFactory.TryCreate(uri, projectContext: null, out var documentContext));
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var endpoint = CreateEndpoint(documentContextFactory);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1),
             Uris = [droppedUri1, droppedUri2]
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -406,47 +358,32 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
 
         // Assert
         Assert.NotNull(result);
-        Assert.Equal("<FetchData />", result!.DocumentChanges!.Value.First[0].Edits[0].NewText);
+        Assert.NotNull(result.DocumentChanges);
+
+        var documentChanges = result.DocumentChanges.GetValueOrDefault();
+        Assert.True(documentChanges.TryGetFirst(out var documentEdits));
+
+        Assert.Equal("<FetchData />", documentEdits[0].Edits[0].NewText);
     }
 
     [Fact]
     public async Task Handle_CSharp_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("@counter");
-        var csharpDocument = codeDocument.GetCSharpDocument();
+        var codeDocument = CreateCodeDocument("@counter");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var projectedRange = It.IsAny<LinePositionSpan>();
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp &&
-            s.TryMapToGeneratedDocumentRange(csharpDocument, It.IsAny<LinePositionSpan>(), out projectedRange) == true, MockBehavior.Strict);
-
-        var documentSnapshot = Mock.Of<IDocumentSnapshot>(s => s.GetGeneratedOutputAsync(false) == Task.FromResult(codeDocument), MockBehavior.Strict);
         var documentContextFactory = CreateDocumentContextFactory(uri, codeDocument);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null);
+        var endpoint = CreateEndpoint(documentContextFactory, clientConnection);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1)
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -460,37 +397,21 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
     public async Task Handle_DocumentNotFound_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
+        var codeDocument = CreateCodeDocument("<div></div>");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
 
-        var documentSnapshot = Mock.Of<IDocumentSnapshot>(s => s.GetGeneratedOutputAsync(false) == Task.FromResult(codeDocument), MockBehavior.Strict);
         var documentContextFactory = CreateDocumentContextFactory(uri, codeDocument);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null);
+        var endpoint = CreateEndpoint(documentContextFactory, clientConnection);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1)
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -504,38 +425,22 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
     public async Task Handle_UnsupportedCodeDocument_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
+        var codeDocument = CreateCodeDocument("<div></div>");
         codeDocument.SetUnsupported();
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
 
-        var documentSnapshot = Mock.Of<IDocumentSnapshot>(s => s.GetGeneratedOutputAsync(false) == Task.FromResult(codeDocument), MockBehavior.Strict);
         var documentContextFactory = CreateDocumentContextFactory(uri, codeDocument);
 
-        var response = new WorkspaceEdit();
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: new WorkspaceEdit());
+        var endpoint = CreateEndpoint(documentContextFactory, clientConnection);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1)
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -549,37 +454,21 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
     public async Task Handle_NoUris_ReturnsNull()
     {
         // Arrange
-        var codeDocument = TestRazorCodeDocument.Create("<div></div>");
+        var codeDocument = CreateCodeDocument("<div></div>");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
 
-        var documentSnapshot = Mock.Of<IDocumentSnapshot>(s => s.GetGeneratedOutputAsync(false) == Task.FromResult(codeDocument), MockBehavior.Strict);
         var documentContextFactory = CreateDocumentContextFactory(uri, codeDocument);
 
-        var response = (WorkspaceEdit?)null;
-
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, It.IsAny<IRazorPresentationParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
-
-        var endpoint = new TextDocumentUriPresentationEndpoint(
-            documentMappingService,
-            clientConnection.Object,
-            FilePathService,
-            documentContextFactory,
-            LoggerFactory);
+        var clientConnection = CreateClientConnection(response: null);
+        var endpoint = CreateEndpoint(documentContextFactory, clientConnection);
 
         var parameters = new UriPresentationParams()
         {
-            TextDocument = new TextDocumentIdentifier
-            {
-                Uri = uri
-            },
+            TextDocument = new() { Uri = uri },
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 1, length: 1)
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -588,4 +477,22 @@ public class TextDocumentUriPresentationEndpointTests(ITestOutputHelper testOutp
         // Assert
         Assert.Null(result);
     }
+
+    private TextDocumentUriPresentationEndpoint CreateEndpoint(
+        IDocumentContextFactory documentContextFactory,
+        IClientConnection? clientConnection = null)
+    {
+        return new TextDocumentUriPresentationEndpoint(
+            StrictMock.Of<IDocumentMappingService>(),
+            clientConnection ?? StrictMock.Of<IClientConnection>(),
+            FilePathService,
+            documentContextFactory,
+            LoggerFactory);
+    }
+
+    private static IClientConnection CreateClientConnection(WorkspaceEdit? response)
+        => TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<IRazorPresentationParams, WorkspaceEdit?>(CustomMessageNames.RazorUriPresentationEndpoint, response);
+        });
 }

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/Formatting_NetFx/FormattingTestBase.cs

@@ -139,7 +139,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
         var filePathService = new LSPFilePathService(TestLanguageServerFeatureOptions.Instance);
         var mappingService = new LspDocumentMappingService(
             filePathService, new TestDocumentContextFactory(), LoggerFactory);
-        var languageKind = mappingService.GetLanguageKind(codeDocument, positionAfterTrigger, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(positionAfterTrigger, rightAssociative: false);
 
         var formattingService = await TestRazorFormattingService.CreateWithFullSupportAsync(LoggerFactory, codeDocument, razorLSPOptions);
         var options = new FormattingOptions()
@@ -213,7 +213,7 @@ public class FormattingTestBase : RazorToolingIntegrationTestBase
 
         var filePathService = new LSPFilePathService(TestLanguageServerFeatureOptions.Instance);
         var mappingService = new LspDocumentMappingService(filePathService, new TestDocumentContextFactory(), LoggerFactory);
-        var languageKind = mappingService.GetLanguageKind(codeDocument, positionAfterTrigger, rightAssociative: false);
+        var languageKind = codeDocument.GetLanguageKind(positionAfterTrigger, rightAssociative: false);
         if (languageKind == RazorLanguageKind.Html)
         {
             throw new NotImplementedException("Code action formatting is not yet supported for HTML in Razor.");

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/RazorDocumentMappingServiceTest.cs

@@ -2,10 +2,8 @@
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
 using System;
-using System.Collections.Generic;
 using System.Collections.Immutable;
 using Microsoft.AspNetCore.Razor.Language;
-using Microsoft.AspNetCore.Razor.Language.Legacy;
 using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.AspNetCore.Razor.Test.Common.LanguageServer;
 using Microsoft.AspNetCore.Razor.Test.Common.Workspaces;
@@ -699,339 +697,12 @@ public class RazorDocumentMappingServiceTest(ITestOutputHelper testOutput) : Too
         Assert.False(result);
     }
 
-    [Fact]
-    public void GetLanguageKindCore_TagHelperElementOwnsName()
-    {
-        // Arrange
-        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
-        descriptor.TagMatchingRule(rule => rule.TagName = "test");
-        descriptor.SetMetadata(TypeName("TestTagHelper"));
-        var text = """
-            @addTagHelper *, TestAssembly
-            <test>@Name</test>
-            """;
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text, new[] { descriptor.Build() });
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 32 + Environment.NewLine.Length, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Html, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_TagHelpersDoNotOwnTrailingEdge()
-    {
-        // Arrange
-        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
-        descriptor.TagMatchingRule(rule => rule.TagName = "test");
-        descriptor.SetMetadata(TypeName("TestTagHelper"));
-        var text = """
-            @addTagHelper *, TestAssembly
-            <test></test>@DateTime.Now
-            """;
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text, new[] { descriptor.Build() });
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 42 + Environment.NewLine.Length, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Razor, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_TagHelperNestedCSharpAttribute()
-    {
-        // Arrange
-        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
-        descriptor.TagMatchingRule(rule => rule.TagName = "test");
-        descriptor.BindAttribute(builder =>
-        {
-            builder.Name = "asp-int";
-            builder.TypeName = typeof(int).FullName;
-            builder.SetMetadata(PropertyName("AspInt"));
-        });
-        descriptor.SetMetadata(TypeName("TestTagHelper"));
-        var text = """
-            @addTagHelper *, TestAssembly
-            <test asp-int='123'></test>
-            """;
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text, new[] { descriptor.Build() });
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 46 + Environment.NewLine.Length, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_CSharp()
-    {
-        // Arrange
-        var text = "<p>@Name</p>";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 5, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_Html()
-    {
-        // Arrange
-        var text = "<p>Hello World</p>";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 5, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Html, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_DefaultsToRazorLanguageIfCannotLocateOwner()
-    {
-        // Arrange
-        var text = "<p>Hello World</p>";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, text.Length + 1, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Razor, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_GetsLastClassifiedSpanLanguageIfAtEndOfDocument()
-    {
-        // Arrange
-        var text = """
-            <strong>Something</strong>
-            <App>
-            """;
-        var classifiedSpans = ImmutableArray.Create<ClassifiedSpanInternal>(
-           new(new SourceSpan(0, 0),
-               blockSpan: new SourceSpan(absoluteIndex: 0, lineIndex: 0, characterIndex: 0, length: text.Length),
-               SpanKindInternal.Transition,
-               blockKind: default,
-               acceptedCharacters: default),
-           new(new SourceSpan(0, 26),
-               blockSpan: default,
-               SpanKindInternal.Markup,
-               blockKind: default,
-               acceptedCharacters: default));
-        var tagHelperSpans = ImmutableArray<TagHelperSpanInternal>.Empty;
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, text.Length, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Html, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_HtmlEdgeEnd()
-    {
-        // Arrange
-        var text = "Hello World";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, text.Length, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Html, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_CSharpEdgeEnd()
-    {
-        // Arrange
-        var text = "@Name";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, text.Length, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_RazorEdgeWithCSharp()
-    {
-        // Arrange
-        var text = "@{}";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 2, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_CSharpEdgeWithCSharpMarker()
-    {
-        // Arrange
-        var text = "@{var x = 1;}";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 12, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_ExplicitExpressionStartCSharp()
-    {
-        // Arrange
-        var text = "@()";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 2, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_ExplicitExpressionInProgressCSharp()
-    {
-        // Arrange
-        var text = "@(Da)";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 4, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_ImplicitExpressionStartCSharp()
-    {
-        // Arrange
-        var text = "@";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 1, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_ImplicitExpressionInProgressCSharp()
-    {
-        // Arrange
-        var text = "@Da";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 3, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_RazorEdgeWithHtml()
-    {
-        // Arrange
-        var text = "@{<br />}";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 2, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Html, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_HtmlInCSharpLeftAssociative()
-    {
-        // Arrange
-        var text = "@if (true) { <br /> }";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 13, text.Length, rightAssociative: false);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_HtmlInCSharpRightAssociative()
-    {
-        // Arrange
-        var text = "@if (true) { <br /> }";
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text);
-
-        // Act\
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 13, text.Length, rightAssociative: true);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Html, languageKind);
-    }
-
-    [Fact]
-    public void GetLanguageKindCore_TagHelperInCSharpRightAssociative()
-    {
-        // Arrange
-        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
-        descriptor.TagMatchingRule(rule => rule.TagName = "test");
-        descriptor.SetMetadata(TypeName("TestTagHelper"));
-        var text = """
-                       @addTagHelper *, TestAssembly
-                       @if {
-                         <test>@Name</test>
-                       }
-                       """;
-        var (classifiedSpans, tagHelperSpans) = GetClassifiedSpans(text, new[] { descriptor.Build() });
-
-        // Act\
-        var languageKind = LspDocumentMappingService.GetLanguageKindCore(classifiedSpans, tagHelperSpans, 40, text.Length, rightAssociative: true);
-
-        // Assert
-        Assert.Equal(RazorLanguageKind.Html, languageKind);
-    }
-
-    private static (ImmutableArray<ClassifiedSpanInternal> classifiedSpans, ImmutableArray<TagHelperSpanInternal> tagHelperSpans) GetClassifiedSpans(string text, IReadOnlyList<TagHelperDescriptor>? tagHelpers = null)
-    {
-        var codeDocument = CreateCodeDocument(text, tagHelpers);
-        var syntaxTree = codeDocument.GetSyntaxTree();
-        var classifiedSpans = syntaxTree.GetClassifiedSpans();
-        var tagHelperSpans = syntaxTree.GetTagHelperSpans();
-        return (classifiedSpans, tagHelperSpans);
-    }
-
-    private static RazorCodeDocument CreateCodeDocument(string text, IReadOnlyList<TagHelperDescriptor>? tagHelpers = null)
-    {
-        tagHelpers ??= Array.Empty<TagHelperDescriptor>();
-        var sourceDocument = TestRazorSourceDocument.Create(text);
-        var projectEngine = RazorProjectEngine.Create(builder => { });
-        var codeDocument = projectEngine.ProcessDesignTime(sourceDocument, "mvc", importSources: default, tagHelpers);
-        return codeDocument;
-    }
-
     private static RazorCodeDocument CreateCodeDocumentWithCSharpProjection(string razorSource, string projectedCSharpSource, ImmutableArray<SourceMapping> sourceMappings)
     {
-        var codeDocument = CreateCodeDocument(razorSource, tagHelpers: []);
+        var sourceDocument = TestRazorSourceDocument.Create(razorSource);
+        var projectEngine = RazorProjectEngine.Create(builder => { });
+        var codeDocument = projectEngine.ProcessDesignTime(sourceDocument, "mvc", importSources: default, tagHelpers: []);
+
         var csharpDocument = new RazorCSharpDocument(
             codeDocument,
             projectedCSharpSource,

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/Refactoring/RenameEndpointTest.cs

@@ -110,7 +110,7 @@ public class RenameEndpointTest(ITestOutputHelper testOutput) : LanguageServerTe
     public async Task Handle_Rename_FileManipulationNotSupported_ReturnsNull()
     {
         // Arrange
-        var options = StrictMock.Of<LanguageServerFeatureOptions>(o =>
+        var options = StrictMock.Of<LanguageServerFeatureOptions>(static o =>
             o.SupportsFileManipulation == false &&
             o.ReturnCodeActionAndRenamePathsWithPrefixedSlash == false);
         var (endpoint, documentContextFactory) = await CreateEndpointAndDocumentContextFactoryAsync(options);
@@ -538,39 +538,36 @@ public class RenameEndpointTest(ITestOutputHelper testOutput) : LanguageServerTe
     public async Task Handle_Rename_SingleServer_CallsDelegatedLanguageServer()
     {
         // Arrange
-        var options = StrictMock.Of<LanguageServerFeatureOptions>(o =>
+        var options = StrictMock.Of<LanguageServerFeatureOptions>(static o =>
             o.SupportsFileManipulation == true &&
             o.SingleServerSupport == true &&
             o.ReturnCodeActionAndRenamePathsWithPrefixedSlash == false);
 
         var delegatedEdit = new WorkspaceEdit();
 
-        var clientConnectionMock = new StrictMock<IClientConnection>();
-        clientConnectionMock
-            .Setup(c => c.SendRequestAsync<IDelegatedParams, WorkspaceEdit>(CustomMessageNames.RazorRenameEndpointName, It.IsAny<DelegatedRenameParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(delegatedEdit);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<IDelegatedParams, WorkspaceEdit>(CustomMessageNames.RazorRenameEndpointName, response: delegatedEdit);
+        });
 
         var documentMappingServiceMock = new StrictMock<IDocumentMappingService>();
-        documentMappingServiceMock
-            .Setup(c => c.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()))
-            .Returns(RazorLanguageKind.CSharp);
 
         var projectedPosition = new LinePosition(1, 1);
         var projectedIndex = 1;
         documentMappingServiceMock
-            .Setup(c => c.TryMapToGeneratedDocumentPosition(It.IsAny<IRazorGeneratedDocument>(), It.IsAny<int>(), out projectedPosition, out projectedIndex))
+            .Setup(x => x.TryMapToGeneratedDocumentPosition(It.IsAny<IRazorGeneratedDocument>(), It.IsAny<int>(), out projectedPosition, out projectedIndex))
             .Returns(true);
 
         var editMappingServiceMock = new StrictMock<IEditMappingService>();
         editMappingServiceMock
-            .Setup(c => c.RemapWorkspaceEditAsync(It.IsAny<IDocumentSnapshot>(), It.IsAny<WorkspaceEdit>(), It.IsAny<CancellationToken>()))
+            .Setup(x => x.RemapWorkspaceEditAsync(It.IsAny<IDocumentSnapshot>(), It.IsAny<WorkspaceEdit>(), It.IsAny<CancellationToken>()))
             .ReturnsAsync(delegatedEdit);
 
         var (endpoint, documentContextFactory) = await CreateEndpointAndDocumentContextFactoryAsync(
             options,
             documentMappingServiceMock.Object,
             editMappingServiceMock.Object,
-            clientConnectionMock.Object);
+            clientConnection);
 
         var uri = PathUtilities.GetUri(s_componentWithParamFilePath);
         var request = new RenameParams
@@ -594,29 +591,19 @@ public class RenameEndpointTest(ITestOutputHelper testOutput) : LanguageServerTe
     public async Task Handle_Rename_SingleServer_DoesNotDelegateForRazor()
     {
         // Arrange
-        var options = StrictMock.Of<LanguageServerFeatureOptions>(o =>
+        var options = StrictMock.Of<LanguageServerFeatureOptions>(static o =>
             o.SupportsFileManipulation == true &&
             o.SingleServerSupport == true &&
             o.ReturnCodeActionAndRenamePathsWithPrefixedSlash == false);
 
-        var clientConnection = StrictMock.Of<IClientConnection>();
-        var documentMappingServiceMock = new StrictMock<IDocumentMappingService>();
-        documentMappingServiceMock
-            .Setup(c => c.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()))
-            .Returns(RazorLanguageKind.Razor);
+        var documentMappingService = StrictMock.Of<IDocumentMappingService>();
 
-        var editMappingService = StrictMock.Of<IEditMappingService>();
-
-        var (endpoint, documentContextFactory) = await CreateEndpointAndDocumentContextFactoryAsync(
-            options,
-            documentMappingServiceMock.Object,
-            editMappingService,
-            clientConnection);
+        var (endpoint, documentContextFactory) = await CreateEndpointAndDocumentContextFactoryAsync(options, documentMappingService);
 
         var request = new RenameParams
         {
             TextDocument = new() { Uri = PathUtilities.GetUri(s_componentWithParamFilePath) },
-            Position = VsLspFactory.CreatePosition(1, 0),
+            Position = VsLspFactory.CreatePosition(0, 1), // This is right after the '@' in '@namespace'
             NewName = "Test2"
         };
 
@@ -663,10 +650,11 @@ public class RenameEndpointTest(ITestOutputHelper testOutput) : LanguageServerTe
                 return textLoaderMock.Object;
             });
 
-        var projectService = new TestRazorProjectService(
-            remoteTextLoaderFactoryMock.Object,
-            projectManager,
-            LoggerFactory);
+        var projectService = AddDisposable(
+            new TestRazorProjectService(
+                remoteTextLoaderFactoryMock.Object,
+                projectManager,
+                LoggerFactory));
 
         var projectKey1 = await projectService.AddProjectAsync(
             s_projectFilePath1, s_intermediateOutputPath1, RazorConfiguration.Default, RootNamespace1, displayName: null, DisposalToken);
@@ -707,7 +695,7 @@ public class RenameEndpointTest(ITestOutputHelper testOutput) : LanguageServerTe
         await projectService.UpdateDocumentAsync(s_componentWithParamFilePath, SourceText.From(ComponentWithParamText), DisposalToken);
 
         var searchEngine = new RazorComponentSearchEngine(projectManager, LoggerFactory);
-        options ??= StrictMock.Of<LanguageServerFeatureOptions>(o =>
+        options ??= StrictMock.Of<LanguageServerFeatureOptions>(static o =>
             o.SupportsFileManipulation == true &&
             o.SingleServerSupport == false &&
             o.ReturnCodeActionAndRenamePathsWithPrefixedSlash == false);
@@ -715,9 +703,7 @@ public class RenameEndpointTest(ITestOutputHelper testOutput) : LanguageServerTe
         if (documentMappingService == null)
         {
             var documentMappingServiceMock = new StrictMock<IDocumentMappingService>();
-            documentMappingServiceMock
-                .Setup(c => c.GetLanguageKind(It.IsAny<RazorCodeDocument>(), It.IsAny<int>(), It.IsAny<bool>()))
-                .Returns(RazorLanguageKind.Html);
+
             var projectedPosition = new LinePosition(1, 1);
             var projectedIndex = 1;
             documentMappingServiceMock

src/Razor/test/Microsoft.AspNetCore.Razor.LanguageServer.Test/WrapWithTag/WrapWithTagEndpointTests.cs

@@ -8,8 +8,8 @@ using System.Threading.Tasks;
 using Microsoft.AspNetCore.Razor.Language;
 using Microsoft.AspNetCore.Razor.LanguageServer.Formatting;
 using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
+using Microsoft.AspNetCore.Razor.Test.Common;
 using Microsoft.AspNetCore.Razor.Test.Common.LanguageServer;
-using Microsoft.CodeAnalysis.Razor.DocumentMapping;
 using Microsoft.CodeAnalysis.Razor.ProjectSystem;
 using Microsoft.CodeAnalysis.Razor.Protocol;
 using Microsoft.VisualStudio.LanguageServer.Protocol;
@@ -28,24 +28,19 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
         var codeDocument = CreateCodeDocument("<div></div>");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var response = new WrapWithTagResponse();
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, It.IsAny<WrapWithTagParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, response: new(), verifiable: true);
+        });
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-        var endpoint = new WrapWithTagEndpoint(
-            clientConnection.Object,
-            documentMappingService,
-            LoggerFactory);
+        var endpoint = new WrapWithTagEndpoint(clientConnection, LoggerFactory);
 
-        var wrapWithDivParams = new WrapWithTagParams(new TextDocumentIdentifier { Uri = uri })
+        var wrapWithDivParams = new WrapWithTagParams(new() { Uri = uri })
         {
             Range = VsLspFactory.CreateSingleLineRange(start: (0, 0), length: 2),
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -53,7 +48,7 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
 
         // Assert
         Assert.NotNull(result);
-        clientConnection.Verify();
+        Mock.Get(clientConnection).Verify();
     }
 
     [Fact]
@@ -63,24 +58,19 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
         var codeDocument = CreateCodeDocument("@(counter)");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var response = new WrapWithTagResponse();
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, It.IsAny<WrapWithTagParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, response: new(), verifiable: true);
+        });
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp, MockBehavior.Strict);
-        var endpoint = new WrapWithTagEndpoint(
-            clientConnection.Object,
-            documentMappingService,
-            LoggerFactory);
+        var endpoint = new WrapWithTagEndpoint(clientConnection, LoggerFactory);
 
-        var wrapWithDivParams = new WrapWithTagParams(new TextDocumentIdentifier { Uri = uri })
+        var wrapWithDivParams = new WrapWithTagParams(new() { Uri = uri })
         {
             Range = VsLspFactory.CreateSingleLineRange(start: (0, 0), length: 2),
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -88,7 +78,8 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
 
         // Assert
         Assert.Null(result);
-        clientConnection.Verify();
+        Mock.Get(clientConnection)
+            .VerifySendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, Times.Never);
     }
 
     [Fact]
@@ -98,24 +89,19 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
         var codeDocument = CreateCodeDocument("@counter");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var response = new WrapWithTagResponse();
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, It.IsAny<WrapWithTagParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, response: new(), verifiable: true);
+        });
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp, MockBehavior.Strict);
-        var endpoint = new WrapWithTagEndpoint(
-            clientConnection.Object,
-            documentMappingService,
-            LoggerFactory);
+        var endpoint = new WrapWithTagEndpoint(clientConnection, LoggerFactory);
 
-        var wrapWithDivParams = new WrapWithTagParams(new TextDocumentIdentifier { Uri = uri })
+        var wrapWithDivParams = new WrapWithTagParams(new() { Uri = uri })
         {
             Range = VsLspFactory.CreateSingleLineRange(start: (0, 0), length: 8),
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -123,7 +109,7 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
 
         // Assert
         Assert.NotNull(result);
-        clientConnection.Verify();
+        Mock.Get(clientConnection).Verify();
     }
 
     [Fact]
@@ -133,24 +119,19 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
         var codeDocument = CreateCodeDocument("@counter");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var response = new WrapWithTagResponse();
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, It.IsAny<WrapWithTagParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, response: new(), verifiable: true);
+        });
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp, MockBehavior.Strict);
-        var endpoint = new WrapWithTagEndpoint(
-            clientConnection.Object,
-            documentMappingService,
-            LoggerFactory);
+        var endpoint = new WrapWithTagEndpoint(clientConnection, LoggerFactory);
 
-        var wrapWithDivParams = new WrapWithTagParams(new TextDocumentIdentifier { Uri = uri })
+        var wrapWithDivParams = new WrapWithTagParams(new() { Uri = uri })
         {
             Range = VsLspFactory.CreateSingleLineRange(line: 0, character: 2, length: 2),
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -158,7 +139,8 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
 
         // Assert
         Assert.Null(result);
-        clientConnection.Verify();
+        Mock.Get(clientConnection)
+            .VerifySendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, Times.Never);
     }
 
     [Fact]
@@ -168,24 +150,19 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
         var codeDocument = CreateCodeDocument("@counter");
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
-        var response = new WrapWithTagResponse();
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
-        clientConnection
-            .Setup(l => l.SendRequestAsync<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, It.IsAny<WrapWithTagParams>(), It.IsAny<CancellationToken>()))
-            .ReturnsAsync(response);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, response: new(), verifiable: true);
+        });
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.CSharp, MockBehavior.Strict);
-        var endpoint = new WrapWithTagEndpoint(
-            clientConnection.Object,
-            documentMappingService,
-            LoggerFactory);
+        var endpoint = new WrapWithTagEndpoint(clientConnection, LoggerFactory);
 
-        var wrapWithDivParams = new WrapWithTagParams(new TextDocumentIdentifier { Uri = uri })
+        var wrapWithDivParams = new WrapWithTagParams(new() { Uri = uri })
         {
             Range = VsLspFactory.CreateZeroWidthRange(0, 4),
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -193,27 +170,27 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
 
         // Assert
         Assert.NotNull(result);
-        clientConnection.Verify();
+        Mock.Get(clientConnection).Verify();
     }
 
     [Fact]
     public async Task Handle_DocumentNotFound_ReturnsNull()
     {
         // Arrange
-        var codeDocument = CreateCodeDocument("<div></div>");
-        var realUri = new Uri("file://path/test.razor");
         var missingUri = new Uri("file://path/nottest.razor");
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, response: new(), verifiable: true);
+        });
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-        var endpoint = new WrapWithTagEndpoint(clientConnection.Object, documentMappingService, LoggerFactory);
+        var endpoint = new WrapWithTagEndpoint(clientConnection, LoggerFactory);
 
-        var wrapWithDivParams = new WrapWithTagParams(new TextDocumentIdentifier { Uri = missingUri })
+        var wrapWithDivParams = new WrapWithTagParams(new() { Uri = missingUri })
         {
             Range = VsLspFactory.CreateSingleLineRange(start: (0, 0), length: 2),
         };
+
         var requestContext = CreateRazorRequestContext(documentContext: null);
 
         // Act
@@ -221,6 +198,8 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
 
         // Assert
         Assert.Null(result);
+        Mock.Get(clientConnection)
+          .VerifySendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, Times.Never);
     }
 
     [Fact]
@@ -232,16 +211,18 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
         var uri = new Uri("file://path/test.razor");
         var documentContext = CreateDocumentContext(uri, codeDocument);
 
-        var clientConnection = new Mock<IClientConnection>(MockBehavior.Strict);
+        var clientConnection = TestMocks.CreateClientConnection(builder =>
+        {
+            builder.SetupSendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, response: new(), verifiable: true);
+        });
 
-        var documentMappingService = Mock.Of<IDocumentMappingService>(
-            s => s.GetLanguageKind(codeDocument, It.IsAny<int>(), It.IsAny<bool>()) == RazorLanguageKind.Html, MockBehavior.Strict);
-        var endpoint = new WrapWithTagEndpoint(clientConnection.Object, documentMappingService, LoggerFactory);
+        var endpoint = new WrapWithTagEndpoint(clientConnection, LoggerFactory);
 
-        var wrapWithDivParams = new WrapWithTagParams(new TextDocumentIdentifier { Uri = uri })
+        var wrapWithDivParams = new WrapWithTagParams(new() { Uri = uri })
         {
             Range = VsLspFactory.CreateSingleLineRange(start: (0, 0), length: 2),
         };
+
         var requestContext = CreateRazorRequestContext(documentContext);
 
         // Act
@@ -249,28 +230,31 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
 
         // Assert
         Assert.Null(result);
+        Mock.Get(clientConnection)
+          .VerifySendRequest<WrapWithTagParams, WrapWithTagResponse>(LanguageServerConstants.RazorWrapWithTagEndpoint, Times.Never);
     }
 
     [Fact]
     public async Task CleanUpTextEdits_NoTilde()
     {
         var input = """
+            @if (true)
+            {
+            }
+            """;
+
+        var expected = """
+            <div>
                 @if (true)
                 {
                 }
-                """;
-        var expected = """
-                <div>
-                    @if (true)
-                    {
-                    }
-                </div>
-                """;
+            </div>
+            """;
 
         var uri = new Uri("file://path.razor");
         var factory = CreateDocumentContextFactory(uri, input);
         Assert.True(factory.TryCreate(uri, out var context));
-        var inputSourceText = await context!.GetSourceTextAsync(DisposalToken);
+        var inputSourceText = await context.GetSourceTextAsync(DisposalToken);
 
         var computedEdits = new TextEdit[]
         {
@@ -281,7 +265,7 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
                 newText: "    }" + Environment.NewLine + "</div>"),
         };
 
-        var htmlSourceText = await context!.GetHtmlSourceTextAsync(DisposalToken);
+        var htmlSourceText = await context.GetHtmlSourceTextAsync(DisposalToken);
         var edits = HtmlFormatter.FixHtmlTextEdits(htmlSourceText, computedEdits);
         Assert.Same(computedEdits, edits);
 
@@ -293,23 +277,23 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
     public async Task CleanUpTextEdits_BadEditWithTilde()
     {
         var input = """
+            @if (true)
+            {
+            }
+            """;
+
+        var expected = """
+            <div>
                 @if (true)
                 {
                 }
-                """;
-
-        var expected = """
-                <div>
-                    @if (true)
-                    {
-                    }
-                </div>
-                """;
+            </div>
+            """;
 
         var uri = new Uri("file://path.razor");
         var factory = CreateDocumentContextFactory(uri, input);
         Assert.True(factory.TryCreate(uri, out var context));
-        var inputSourceText = await context!.GetSourceTextAsync(DisposalToken);
+        var inputSourceText = await context.GetSourceTextAsync(DisposalToken);
 
         var computedEdits = new TextEdit[]
         {
@@ -321,7 +305,7 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
                 newText: "    ~" + Environment.NewLine + "</div>")
         };
 
-        var htmlSourceText = await context!.GetHtmlSourceTextAsync(DisposalToken);
+        var htmlSourceText = await context.GetHtmlSourceTextAsync(DisposalToken);
         var edits = HtmlFormatter.FixHtmlTextEdits(htmlSourceText, computedEdits);
         Assert.NotSame(computedEdits, edits);
 
@@ -333,18 +317,18 @@ public class WrapWithTagEndpointTest(ITestOutputHelper testOutput) : LanguageSer
     public async Task CleanUpTextEdits_GoodEditWithTilde()
     {
         var input = """
+            @if (true)
+            {
+            ~
+            """;
+
+        var expected = """
+            <div>
                 @if (true)
                 {
                 ~
-                """;
-
-        var expected = """
-                <div>
-                    @if (true)
-                    {
-                    ~
-                </div>
-                """;
+            </div>
+            """;
 
         var uri = new Uri("file://path.razor");
         var factory = CreateDocumentContextFactory(uri, input);

src/Razor/test/Microsoft.AspNetCore.Razor.Test.Common.Tooling/TestMocks.cs

@@ -1,10 +1,14 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT license. See License.txt in the project root for license information.
 
+using System;
 using System.Threading;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.LanguageServer.Hosting;
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.Text;
 using Moq;
+using Moq.Language.Flow;
 
 namespace Microsoft.AspNetCore.Razor.Test.Common;
 
@@ -26,4 +30,58 @@ internal static class TestMocks
 
         return mock.Object;
     }
+
+    public interface IClientConnectionBuilder
+    {
+        void SetupSendRequest<TParams, TResponse>(string method, TResponse response, bool verifiable = false);
+        void SetupSendRequest<TParams, TResponse>(string method, TParams @params, TResponse response, bool verifiable = false);
+    }
+
+    private sealed class ClientConnectionBuilder : IClientConnectionBuilder
+    {
+        public StrictMock<IClientConnection> Mock { get; } = new();
+
+        public void SetupSendRequest<TParams, TResponse>(string method, TResponse response, bool verifiable = false)
+        {
+            var returnsResult = Mock
+                .Setup(x => x.SendRequestAsync<TParams, TResponse>(method, It.IsAny<TParams>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(response);
+
+            if (verifiable)
+            {
+                returnsResult.Verifiable();
+            }
+        }
+
+        public void SetupSendRequest<TParams, TResponse>(string method, TParams @params, TResponse response, bool verifiable = false)
+        {
+            var returnsResult = Mock
+                .Setup(x => x.SendRequestAsync<TParams, TResponse>(method, @params, It.IsAny<CancellationToken>()))
+                .ReturnsAsync(response);
+
+            if (verifiable)
+            {
+                returnsResult.Verifiable();
+            }
+        }
+    }
+
+    public static IClientConnection CreateClientConnection(Action<IClientConnectionBuilder> configure)
+    {
+        var builder = new ClientConnectionBuilder();
+        configure?.Invoke(builder);
+        return builder.Mock.Object;
+    }
+
+    public static void VerifySendRequest<TParams, TResponse>(this Mock<IClientConnection> mock, string method, Times times)
+        => mock.Verify(x => x.SendRequestAsync<TParams, TResponse>(method, It.IsAny<TParams>(), It.IsAny<CancellationToken>()), times);
+
+    public static void VerifySendRequest<TParams, TResponse>(this Mock<IClientConnection> mock, string method, Func<Times> times)
+        => mock.Verify(x => x.SendRequestAsync<TParams, TResponse>(method, It.IsAny<TParams>(), It.IsAny<CancellationToken>()), times);
+
+    public static void VerifySendRequest<TParams, TResponse>(this Mock<IClientConnection> mock, string method, TParams @params, Times times)
+        => mock.Verify(x => x.SendRequestAsync<TParams, TResponse>(method, @params, It.IsAny<CancellationToken>()), times);
+
+    public static void VerifySendRequest<TParams, TResponse>(this Mock<IClientConnection> mock, string method, TParams @params, Func<Times> times)
+        => mock.Verify(x => x.SendRequestAsync<TParams, TResponse>(method, @params, It.IsAny<CancellationToken>()), times);
 }

src/Razor/test/Microsoft.CodeAnalysis.Razor.Workspaces.Test/Extensions/RazorCodeDocumentExtensionsTest.cs

@@ -0,0 +1,340 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT license. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Razor.Language;
+using Microsoft.AspNetCore.Razor.Test.Common;
+using Microsoft.CodeAnalysis.Razor.Protocol;
+using Xunit;
+using Xunit.Abstractions;
+using static Microsoft.AspNetCore.Razor.Language.CommonMetadata;
+
+namespace Microsoft.CodeAnalysis.Razor.Workspaces.Test.Extensions;
+
+public class RazorCodeDocumentExtensionsTest(ITestOutputHelper testOutput) : ToolingTestBase(testOutput)
+{
+    [Fact]
+    public void GetLanguageKind_TagHelperElementOwnsName()
+    {
+        // Arrange
+        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
+        descriptor.TagMatchingRule(rule => rule.TagName = "test");
+        descriptor.SetMetadata(TypeName("TestTagHelper"));
+
+        TestCode code = """
+            @addTagHelper *, TestAssembly
+            <te$$st>@Name</test>
+            """;
+
+        var codeDocument = CreateCodeDocument(code, descriptor.Build());
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Html, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_TagHelpersDoNotOwnTrailingEdge()
+    {
+        // Arrange
+        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
+        descriptor.TagMatchingRule(rule => rule.TagName = "test");
+        descriptor.SetMetadata(TypeName("TestTagHelper"));
+
+        TestCode code = """
+            @addTagHelper *, TestAssembly
+            <test></test>$$@DateTime.Now
+            """;
+
+        var codeDocument = CreateCodeDocument(code, descriptor.Build());
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Razor, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_TagHelperNestedCSharpAttribute()
+    {
+        // Arrange
+        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
+        descriptor.TagMatchingRule(rule => rule.TagName = "test");
+        descriptor.BindAttribute(builder =>
+        {
+            builder.Name = "asp-int";
+            builder.TypeName = typeof(int).FullName;
+            builder.SetMetadata(PropertyName("AspInt"));
+        });
+        descriptor.SetMetadata(TypeName("TestTagHelper"));
+
+        TestCode code = """
+            @addTagHelper *, TestAssembly
+            <test asp-int='12$$3'></test>
+            """;
+
+        var codeDocument = CreateCodeDocument(code, descriptor.Build());
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_CSharp()
+    {
+        // Arrange
+        TestCode code = "<p>@N$$ame</p>";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_Html()
+    {
+        // Arrange
+        TestCode code = "<p>He$$llo World</p>";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Html, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_DefaultsToRazorLanguageIfCannotLocateOwner()
+    {
+        // Arrange
+        TestCode code = "<p>Hello World</p>$$";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position + 1, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Razor, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_GetsLastClassifiedSpanLanguageIfAtEndOfDocument()
+    {
+        // Arrange
+        TestCode code = """
+            <strong>Something</strong>
+            <App>$$
+            """;
+
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Html, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_HtmlEdgeEnd()
+    {
+        // Arrange
+        TestCode code = "Hello World$$";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Html, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_CSharpEdgeEnd()
+    {
+        // Arrange
+        TestCode code = "@Name$$";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_RazorEdgeWithCSharp()
+    {
+        // Arrange
+        TestCode code = "@{$$}";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_CSharpEdgeWithCSharpMarker()
+    {
+        // Arrange
+        TestCode code = "@{var x = 1;$$}";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_ExplicitExpressionStartCSharp()
+    {
+        // Arrange
+        TestCode code = "@($$)";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_ExplicitExpressionInProgressCSharp()
+    {
+        // Arrange
+        TestCode code = "@(Da$$)";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_ImplicitExpressionStartCSharp()
+    {
+        // Arrange
+        TestCode code = "@$$";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_ImplicitExpressionInProgressCSharp()
+    {
+        // Arrange
+        TestCode code = "@Da$$";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_RazorEdgeWithHtml()
+    {
+        // Arrange
+        TestCode code = "@{$$<br />}";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Html, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_HtmlInCSharpLeftAssociative()
+    {
+        // Arrange
+        TestCode code = "@if (true) { $$<br /> }";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: false);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.CSharp, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_HtmlInCSharpRightAssociative()
+    {
+        // Arrange
+        TestCode code = "@if (true) { $$<br /> }";
+        var codeDocument = CreateCodeDocument(code);
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: true);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Html, languageKind);
+    }
+
+    [Fact]
+    public void GetLanguageKind_TagHelperInCSharpRightAssociative()
+    {
+        // Arrange
+        var descriptor = TagHelperDescriptorBuilder.Create("TestTagHelper", "TestAssembly");
+        descriptor.TagMatchingRule(rule => rule.TagName = "test");
+        descriptor.SetMetadata(TypeName("TestTagHelper"));
+
+        TestCode code = """
+            @addTagHelper *, TestAssembly
+            @if {
+                $$<test>@Name</test>
+            }
+            """;
+
+        var codeDocument = CreateCodeDocument(code, descriptor.Build());
+
+        // Act
+        var languageKind = codeDocument.GetLanguageKind(code.Position, rightAssociative: true);
+
+        // Assert
+        Assert.Equal(RazorLanguageKind.Html, languageKind);
+    }
+
+    private static RazorCodeDocument CreateCodeDocument(TestCode code, params ImmutableArray<TagHelperDescriptor> tagHelpers)
+    {
+        tagHelpers = tagHelpers.NullToEmpty();
+
+        var sourceDocument = TestRazorSourceDocument.Create(code.Text);
+        var projectEngine = RazorProjectEngine.Create(builder => { });
+
+        return projectEngine.ProcessDesignTime(sourceDocument, "mvc", importSources: default, tagHelpers);
+    }
+}