Add tests and fix bugs

src/Features/LanguageServer/Protocol/Microsoft.CodeAnalysis.LanguageServer.Protocol.csproj

@@ -34,6 +34,7 @@
     <InternalsVisibleTo Include="Microsoft.VisualStudio.LanguageServices.LiveShare.UnitTests" />
     <InternalsVisibleTo Include="Microsoft.CodeAnalysis.Remote.ServiceHub" />
     <InternalsVisibleTo Include="Roslyn.Services.Test.Utilities" />
+    <InternalsVisibleTo Include="Roslyn.VisualStudio.Next.UnitTests" />
   </ItemGroup>
 
   <ItemGroup>

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

@@ -249,7 +249,7 @@ private async void DiagnosticService_DiagnosticsUpdated(object sender, Diagnosti
         /// </summary>
         private readonly Dictionary<DocumentId, ImmutableHashSet<Uri>> _documentsToPublishedUris = new Dictionary<DocumentId, ImmutableHashSet<Uri>>();
 
-        private async Task PublishDiagnosticsAsync(Document document)
+        internal async Task PublishDiagnosticsAsync(Document document)
         {
             // Retrieve all diagnostics for the current document grouped by their actual file uri.
             var fileUriToDiagnostics = await GetDiagnosticsAsync(document, CancellationToken.None).ConfigureAwait(false);
@@ -257,7 +257,7 @@ private async Task PublishDiagnosticsAsync(Document document)
             // Get the list of file uris with diagnostics (for the 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 was fixed).
+            // of the current diagnostics set (because the diagnostics were fixed).
             var urisForCurrentDocument = GetOrValue(_documentsToPublishedUris, document.Id, ImmutableHashSet<Uri>.Empty).Union(fileUriToDiagnostics.Keys);
 
             // Update the mapping for this document to be the uris we're about to publish diagnostics for.
@@ -289,9 +289,18 @@ private async Task PublishDiagnosticsAsync(Document document)
                     _documentsToPublishedUris[document.Id] = _documentsToPublishedUris[document.Id].Remove(fileUri);
                 }
 
-                // Update the published diagnostics map to contain the new diagnostics for this document and mapped uri.
-                _publishedFileToDiagnostics.GetOrAdd(fileUri,
-                    (_) => { return new Dictionary<DocumentId, ImmutableArray<LanguageServer.Protocol.Diagnostic>>(); })[document.Id] = diagnostics;
+                // Update the published diagnostics map to contain the new diagnostics contributed by this document and fileUri.
+                var documentsToPublishedDiagnostics = _publishedFileToDiagnostics.GetOrAdd(fileUri, (_) =>
+                    new Dictionary<DocumentId, ImmutableArray<LanguageServer.Protocol.Diagnostic>>());
+                if (fileUriToDiagnostics.ContainsKey(fileUri))
+                {
+                    documentsToPublishedDiagnostics[document.Id] = fileUriToDiagnostics[fileUri];
+                }
+                else
+                {
+                    // There are no new diagnostics for this document and file uri, if we're tracking it we can stop.
+                    documentsToPublishedDiagnostics.Remove(document.Id);
+                }
             }
         }