Merge pull request #43030 from dibarbet/cleanup_razor

Add test accessor + more verification on tests and fix flakiness

src/VisualStudio/Core/Def/Implementation/LanguageClient/InProcLanguageServer.cs

@@ -9,6 +9,7 @@
 using System.Collections.Immutable;
 using System.IO;
 using System.Linq;
+using System.Runtime.Remoting.Messaging;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis;
@@ -240,20 +241,21 @@ private async void DiagnosticService_DiagnosticsUpdated(object sender, Diagnosti
         /// This dictionary stores the previously computed diagnostics for the published file so that we can
         /// union the currently computed diagnostics (e.g. for dA) with previously computed diagnostics (e.g. from dB).
         /// </summary>
-        internal readonly Dictionary<Uri, Dictionary<DocumentId, ImmutableArray<LanguageServer.Protocol.Diagnostic>>> _publishedFileToDiagnostics =
+        private readonly Dictionary<Uri, Dictionary<DocumentId, ImmutableArray<LanguageServer.Protocol.Diagnostic>>> _publishedFileToDiagnostics =
             new Dictionary<Uri, Dictionary<DocumentId, ImmutableArray<LanguageServer.Protocol.Diagnostic>>>();
 
         /// <summary>
         /// Stores the mapping of a document to the uri(s) of diagnostics previously produced for this document.
         /// When we get empty diagnostics for the document we need to find the uris we previously published for this document.
         /// Then we can publish the updated diagnostics set for those uris (either empty or the diagnostic contributions from other documents).
+        /// We use a sorted set to ensure consistency in the order in which we report URIs.
+        /// While it's not necessary to publish a document's mapped file diagnostics in a particular order,
+        /// it does make it much easier to write tests and debug issues if we have a consistent ordering.
         /// </summary>
-        internal readonly Dictionary<DocumentId, ImmutableSortedSet<Uri>> _documentsToPublishedUris = new Dictionary<DocumentId, ImmutableSortedSet<Uri>>();
+        private readonly Dictionary<DocumentId, ImmutableSortedSet<Uri>> _documentsToPublishedUris = new Dictionary<DocumentId, ImmutableSortedSet<Uri>>();
 
         /// <summary>
-        /// Basic comparer for Uris used by <see cref="_documentsToPublishedUris"/> when publishing notifications for mapped files.
-        /// While it's not necessary to publish a document's mapped file diagnostics in a particular order,
-        /// it does make it much easier to write tests and debug issues if we have a consistent ordering.
+        /// Basic comparer for Uris used by <see cref="_documentsToPublishedUris"/> when publishing notifications.
         /// </summary>
         private static readonly Comparer<Uri> s_uriComparer = Comparer<Uri>.Create((uri1, uri2)
             => Uri.Compare(uri1, uri2, UriComponents.AbsoluteUri, UriFormat.SafeUnescaped, StringComparison.OrdinalIgnoreCase));
@@ -267,6 +269,7 @@ internal async Task PublishDiagnosticsAsync(Document document)
             // We need to join the uris from current diagnostics with those previously published
             // so that we clear out any diagnostics in mapped files that are no longer a part
             // of the current diagnostics set (because the diagnostics were fixed).
+            // Use sorted set to have consistent publish ordering for tests and debugging.
             var urisForCurrentDocument = _documentsToPublishedUris.GetOrValue(document.Id, ImmutableSortedSet.Create<Uri>(s_uriComparer)).Union(fileUriToDiagnostics.Keys);
 
             // Update the mapping for this document to be the uris we're about to publish diagnostics for.
@@ -390,5 +393,36 @@ private bool IncludeDiagnostic(DiagnosticData diagnostic)
             var linePositionSpan = DiagnosticData.GetLinePositionSpan(diagnosticDataLocation, text, useMapped: true);
             return ProtocolConversions.LinePositionToRange(linePositionSpan);
         }
+
+        internal TestAccessor GetTestAccessor() => new TestAccessor(this);
+
+        internal class TestAccessor
+        {
+            private readonly InProcLanguageServer _server;
+
+            internal TestAccessor(InProcLanguageServer server)
+            {
+                _server = server;
+            }
+
+            internal ImmutableArray<Uri> GetFileUrisInPublishDiagnostics()
+                => _server._publishedFileToDiagnostics.Keys.ToImmutableArray();
+
+            internal ImmutableArray<DocumentId> GetDocumentIdsInPublishedUris()
+                => _server._documentsToPublishedUris.Keys.ToImmutableArray();
+
+            internal IImmutableSet<Uri> GetFileUrisForDocument(DocumentId documentId)
+                => _server._documentsToPublishedUris.GetOrValue(documentId, ImmutableSortedSet<Uri>.Empty);
+
+            internal ImmutableArray<LanguageServer.Protocol.Diagnostic> GetDiagnosticsForUriAndDocument(DocumentId documentId, Uri uri)
+            {
+                if (_server._publishedFileToDiagnostics.TryGetValue(uri, out var dict) && dict.TryGetValue(documentId, out var diagnostics))
+                {
+                    return diagnostics;
+                }
+
+                return ImmutableArray<LanguageServer.Protocol.Diagnostic>.Empty;
+            }
+        }
     }
 }

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

@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+#nullable enable
+
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
@@ -13,9 +15,11 @@
 using Microsoft.CodeAnalysis.Diagnostics;
 using Microsoft.CodeAnalysis.Editor.UnitTests.Workspaces;
 using Microsoft.CodeAnalysis.LanguageServer;
+using Microsoft.CodeAnalysis.Shared.Extensions;
 using Microsoft.CodeAnalysis.Test.Utilities;
 using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.LanguageServices.Implementation.LanguageService;
+using Microsoft.VisualStudio.Threading;
 using Moq;
 using Nerdbank.Streams;
 using Newtonsoft.Json.Linq;
@@ -45,7 +49,7 @@ public async Task AddDiagnosticTestAsync()
             //
             // We expect one publish diagnostic notification ->
             // 1.  from doc1 with id.
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 1, document).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 1, document).ConfigureAwait(false);
 
             var result = Assert.Single(results);
             Assert.Equal(new Uri(document.FilePath), result.Uri);
@@ -69,7 +73,7 @@ public async Task AddDiagnosticWithMappedFilesTestAsync()
             // We expect two publish diagnostic notifications ->
             // 1.  from m1 with id1 (from 1 above).
             // 2.  from m2 with id2 (from 1 above).
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, expectedNumberOfCallbacks: 2, document).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, expectedNumberOfCallbacks: 2, document).ConfigureAwait(false);
 
             Assert.Equal(2, results.Count);
             Assert.Equal(new Uri(document.FilePath + "m1"), results[0].Uri);
@@ -102,7 +106,7 @@ public async Task AddDiagnosticWithMappedFileToManyDocumentsTestAsync()
             // We expect two publish diagnostic notifications ->
             // 1.  from m1 with doc1Diagnostic (from 1 above).
             // 2.  from m1 with doc1Diagnostic and doc2Diagnostic (from 2 above adding doc2Diagnostic to m1).
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 2, documents[0], documents[1]).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 2, documents[0], documents[1]).ConfigureAwait(false);
 
             Assert.Equal(2, results.Count);
             var expectedUri = new Uri(mappedFilePath);
@@ -134,7 +138,7 @@ public async Task RemoveDiagnosticTestAsync()
             // We expect two publish diagnostic notifications ->
             // 1.  from doc1 with id.
             // 2.  from doc1 with empty (from 2 above clearing out diagnostics from doc1).
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 2, document, document).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 2, document, document).ConfigureAwait(false);
 
             Assert.Equal(2, results.Count);
             Assert.Equal(new Uri(document.FilePath), results[0].Uri);
@@ -143,8 +147,8 @@ public async Task RemoveDiagnosticTestAsync()
             Assert.Equal(new Uri(document.FilePath), results[1].Uri);
             Assert.True(results[1].Diagnostics.IsEmpty());
 
-            Assert.False(languageServer._documentsToPublishedUris.Keys.Any());
-            Assert.False(languageServer._publishedFileToDiagnostics.Keys.Any());
+            Assert.Empty(testAccessor.GetDocumentIdsInPublishedUris());
+            Assert.Empty(testAccessor.GetFileUrisInPublishDiagnostics());
         }
 
         [Fact]
@@ -172,7 +176,7 @@ public async Task RemoveDiagnosticForMappedFilesTestAsync()
             // 2.  from m2 with id2 (from 1 above).
             // 3.  from m1 with empty (from 2 above clearing out diagnostics for m1).
             // 4.  from m2 with id2 (from 2 above clearing out diagnostics for m1).
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 4, document, document).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 4, document, document).ConfigureAwait(false);
 
             var mappedFileURIM1 = new Uri(mappedFilePathM1);
             var mappedFileURIM2 = new Uri(mappedFilePathM2);
@@ -192,6 +196,10 @@ public async Task RemoveDiagnosticForMappedFilesTestAsync()
 
             Assert.Equal(mappedFileURIM2, results[3].Uri);
             Assert.Equal("id2", results[3].Diagnostics.Single().Code);
+
+            Assert.Single(testAccessor.GetFileUrisForDocument(document.Id), mappedFileURIM2);
+            Assert.Equal("id2", testAccessor.GetDiagnosticsForUriAndDocument(document.Id, mappedFileURIM2).Single().Code);
+            Assert.Empty(testAccessor.GetDiagnosticsForUriAndDocument(document.Id, mappedFileURIM1));
         }
 
         [Fact]
@@ -221,7 +229,7 @@ public async Task RemoveDiagnosticForMappedFileToManyDocumentsTestAsync()
             // 1.  from m1 with with doc1Diagnostic (triggered by 1 above to add doc1Diagnostic).
             // 2.  from m1 with doc1Diagnostic and doc2Diagnostic (triggered by 2 above to add doc2Diagnostic).
             // 3.  from m1 with just doc2Diagnostic (triggered by 3 above to remove doc1Diagnostic).
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 3, documents[0], documents[1], documents[0]).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 3, documents[0], documents[1], documents[0]).ConfigureAwait(false);
 
             Assert.Equal(3, results.Count);
             var expectedUri = new Uri(mappedFilePath);
@@ -236,6 +244,10 @@ public async Task RemoveDiagnosticForMappedFileToManyDocumentsTestAsync()
             Assert.Equal(expectedUri, results[2].Uri);
             Assert.Equal(1, results[2].Diagnostics.Length);
             Assert.Contains(results[2].Diagnostics, d => d.Code == "doc2Diagnostic");
+
+            Assert.Single(testAccessor.GetFileUrisForDocument(documents[1].Id), expectedUri);
+            Assert.Equal("doc2Diagnostic", testAccessor.GetDiagnosticsForUriAndDocument(documents[1].Id, expectedUri).Single().Code);
+            Assert.Empty(testAccessor.GetDiagnosticsForUriAndDocument(documents[0].Id, expectedUri));
         }
 
         [Fact(Skip = "https://github.com/dotnet/roslyn/issues/43046")]
@@ -259,7 +271,7 @@ public async Task ClearAllDiagnosticsForMappedFilesTestAsync()
             //
             // We expect four publish diagnostic notifications - the first two are the two mapped files from 1.
             // The second two are the two mapped files being cleared by 2.
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 4, document, document).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 4, document, document).ConfigureAwait(false);
 
             var mappedFileURIM1 = new Uri(document.FilePath + "m1");
             var mappedFileURIM2 = new Uri(document.FilePath + "m2");
@@ -280,8 +292,8 @@ public async Task ClearAllDiagnosticsForMappedFilesTestAsync()
             Assert.Equal(mappedFileURIM2, results[3].Uri);
             Assert.True(results[3].Diagnostics.IsEmpty());
 
-            Assert.False(languageServer._documentsToPublishedUris.Keys.Any());
-            Assert.False(languageServer._publishedFileToDiagnostics.Keys.Any());
+            Assert.Empty(testAccessor.GetDocumentIdsInPublishedUris());
+            Assert.Empty(testAccessor.GetFileUrisInPublishDiagnostics());
         }
 
         [Fact]
@@ -312,7 +324,7 @@ public async Task ClearAllDiagnosticsForMappedFileToManyDocumentsTestAsync()
             // 2.  from URI m1 with doc1Diagnostic and doc2Diagnostic (triggered by 2 above to add doc2Diagnostic).
             // 3.  from URI m1 with just doc2Diagnostic (triggered by 3 above to clear doc1 diagnostic).
             // 4.  from URI m1 with empty (triggered by 4 above to also clear doc2 diagnostic).
-            var (languageServer, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 4, documents[0], documents[1], documents[0], documents[1]).ConfigureAwait(false);
+            var (testAccessor, results) = await RunPublishDiagnosticsAsync(workspace, diagnosticsMock.Object, 4, documents[0], documents[1], documents[0], documents[1]).ConfigureAwait(false);
 
             Assert.Equal(4, results.Count);
             var expectedUri = new Uri(mappedFilePath);
@@ -331,29 +343,37 @@ public async Task ClearAllDiagnosticsForMappedFileToManyDocumentsTestAsync()
             Assert.Equal(expectedUri, results[3].Uri);
             Assert.True(results[3].Diagnostics.IsEmpty());
 
-            Assert.False(languageServer._documentsToPublishedUris.Keys.Any());
-            Assert.False(languageServer._publishedFileToDiagnostics.Keys.Any());
+            Assert.Empty(testAccessor.GetDocumentIdsInPublishedUris());
+            Assert.Empty(testAccessor.GetFileUrisInPublishDiagnostics());
         }
 
-        private async Task<(InProcLanguageServer, List<LSP.PublishDiagnosticParams>)> RunPublishDiagnosticsAsync(Workspace workspace, IDiagnosticService diagnosticService,
+        private async Task<(InProcLanguageServer.TestAccessor, List<LSP.PublishDiagnosticParams>)> RunPublishDiagnosticsAsync(Workspace workspace, IDiagnosticService diagnosticService,
             int expectedNumberOfCallbacks, params Document[] documentsToPublish)
         {
             var (clientStream, serverStream) = FullDuplexStream.CreatePair();
             var languageServer = CreateLanguageServer(serverStream, serverStream, workspace, diagnosticService);
 
+            // Notification target for tests to recieve the notification details
             var callback = new Callback(expectedNumberOfCallbacks);
-            using (var jsonRpc = JsonRpc.Attach(clientStream, callback))
-            {
-                foreach (var document in documentsToPublish)
-                {
-                    await languageServer.PublishDiagnosticsAsync(document).ConfigureAwait(false);
-                }
+            using var jsonRpc = new JsonRpc(clientStream, clientStream, callback);
 
-                await callback.CallbackCompletedTask.Task.ConfigureAwait(false);
+            // The json rpc messages won't necessarily come back in order by default.
+            // So use a synchronization context to preserve the original ordering.
+            // https://github.com/microsoft/vs-streamjsonrpc/blob/bc970c61b90db5db135a1b3d1c72ef355c2112af/doc/resiliency.md#when-message-order-is-important
+            jsonRpc.SynchronizationContext = new RpcOrderPreservingSynchronizationContext();
+            jsonRpc.StartListening();
 
-                return (languageServer, callback.Results);
+            // Triggers language server to send notifications.
+            foreach (var document in documentsToPublish)
+            {
+                await languageServer.PublishDiagnosticsAsync(document).ConfigureAwait(false);
             }
 
+            // Waits for all notifications to be recieved.
+            await callback.CallbackCompletedTask.Task.ConfigureAwait(false);
+
+            return (languageServer.GetTestAccessor(), callback.Results);
+
             static InProcLanguageServer CreateLanguageServer(Stream inputStream, Stream outputStream, Workspace workspace, IDiagnosticService mockDiagnosticService)
             {
                 var protocol = ((TestWorkspace)workspace).ExportProvider.GetExportedValue<LanguageServerProtocol>();
@@ -382,13 +402,13 @@ private void SetupMockWithDiagnostics(Mock<IDiagnosticService> diagnosticService
         private async Task<IEnumerable<DiagnosticData>> CreateMockDiagnosticDataAsync(Document document, string id)
         {
             var descriptor = new DiagnosticDescriptor(id, "", "", "", DiagnosticSeverity.Error, true);
-            var location = Location.Create(await document.GetSyntaxTreeAsync().ConfigureAwait(false), new TextSpan());
+            var location = Location.Create(await document.GetRequiredSyntaxTreeAsync(CancellationToken.None).ConfigureAwait(false), new TextSpan());
             return new DiagnosticData[] { DiagnosticData.Create(Diagnostic.Create(descriptor, location), document) };
         }
 
         private async Task<ImmutableArray<DiagnosticData>> CreateMockDiagnosticDatasWithMappedLocationAsync(Document document, params (string diagnosticId, string mappedFilePath)[] diagnostics)
         {
-            var tree = await document.GetSyntaxTreeAsync().ConfigureAwait(false);
+            var tree = await document.GetRequiredSyntaxTreeAsync(CancellationToken.None).ConfigureAwait(false);
 
             return diagnostics.Select(d => CreateMockDiagnosticDataWithMappedLocation(document, tree, d.diagnosticId, d.mappedFilePath)).ToImmutableArray();
 
@@ -422,6 +442,50 @@ static DiagnosticDataLocation GetDataLocation(Document document, string mappedFi
                 => new DiagnosticDataLocation(document.Id, originalFilePath: document.FilePath, mappedFilePath: mappedFilePath);
         }
 
+        /// <summary>
+        /// Synchronization context to preserve ordering of the RPC messages
+        /// Adapted from https://dev.azure.com/devdiv/DevDiv/VS%20Cloud%20Kernel/_git/DevCore?path=%2Fsrc%2Fclr%2FMicrosoft.ServiceHub.Framework%2FServiceRpcDescriptor%2BRpcOrderPreservingSynchronizationContext.cs
+        /// https://github.com/microsoft/vs-streamjsonrpc/issues/440 tracks exposing functionality so we don't need to copy this.
+        /// </summary>
+        private class RpcOrderPreservingSynchronizationContext : SynchronizationContext, IDisposable
+        {
+            /// <summary>
+            /// The queue of work to execute.
+            /// </summary>
+            private readonly AsyncQueue<(SendOrPostCallback, object?)> _queue = new AsyncQueue<(SendOrPostCallback, object?)>();
+
+            public RpcOrderPreservingSynchronizationContext()
+            {
+                // Process the work in the background.
+                this.ProcessQueueAsync().Forget();
+            }
+
+            public override void Post(SendOrPostCallback d, object? state) => this._queue.Enqueue((d, state));
+
+            public override void Send(SendOrPostCallback d, object? state) => throw new NotSupportedException();
+
+            public override SynchronizationContext CreateCopy() => throw new NotSupportedException();
+
+            /// <summary>
+            /// Causes this <see cref="SynchronizationContext"/> to reject all future posted work and
+            /// releases the queue processor when it is empty.
+            /// </summary>
+            public void Dispose() => this._queue.Complete();
+
+            /// <summary>
+            /// Executes queued work on the threadpool, one at a time.
+            /// Don't catch exceptions - let them bubble up to fail the test.
+            /// </summary>
+            private async Task ProcessQueueAsync()
+            {
+                while (!this._queue.IsCompleted)
+                {
+                    var work = await this._queue.DequeueAsync().ConfigureAwait(false);
+                    work.Item1(work.Item2);
+                }
+            }
+        }
+
         private class Callback
         {
             /// <summary>

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/IDictionaryExtensions.cs

@@ -152,13 +152,19 @@ internal static class IDictionaryExtensions
             return dictionary;
         }
 
-        public static void MultiRemove<TKey, TValue>(this IDictionary<TKey, ImmutableHashSet<TValue>> dictionary, TKey key, TValue value)
+        /// <summary>
+        /// Private implementation we can delegate to for sets.
+        /// This must be a different name as overloads are not resolved based on constraints
+        /// and would conflict with <see cref="MultiRemove{TKey, TValue, TCollection}(IDictionary{TKey, TCollection}, TKey, TValue)"/>
+        /// </summary>
+        private static void MultiRemoveSet<TKey, TValue, TSet>(this IDictionary<TKey, TSet> dictionary, TKey key, TValue value)
             where TKey : notnull
+            where TSet : IImmutableSet<TValue>
         {
             if (dictionary.TryGetValue(key, out var collection))
             {
-                collection = collection.Remove(value);
-                if (collection.IsEmpty)
+                collection = (TSet)collection.Remove(value);
+                if (collection.IsEmpty())
                 {
                     dictionary.Remove(key);
                 }
@@ -169,21 +175,16 @@ internal static class IDictionaryExtensions
             }
         }
 
+        public static void MultiRemove<TKey, TValue>(this IDictionary<TKey, ImmutableHashSet<TValue>> dictionary, TKey key, TValue value)
+            where TKey : notnull
+        {
+            MultiRemoveSet(dictionary, key, value);
+        }
+
         public static void MultiRemove<TKey, TValue>(this IDictionary<TKey, ImmutableSortedSet<TValue>> dictionary, TKey key, TValue value)
             where TKey : notnull
         {
-            if (dictionary.TryGetValue(key, out var collection))
-            {
-                collection = collection.Remove(value);
-                if (collection.IsEmpty)
-                {
-                    dictionary.Remove(key);
-                }
-                else
-                {
-                    dictionary[key] = collection;
-                }
-            }
+            MultiRemoveSet(dictionary, key, value);
         }
 
         public static void MultiRemove<TKey, TValue>(this IDictionary<TKey, ImmutableArray<TValue>> dictionary, TKey key, TValue value)