gopls/internal/regtest: simplify awaiting diagnostics from a change

When awaiting diagnostics, we should almost always wrap expectations in
a OnceMet(precondition, ...), so that tests do not hang indefinitely if
the diagnostic pass completes and the expectations are still not met.

Before this change, the user must be careful to pass in the correct
precondition (or combination of preconditions), else they may be
susceptible to races.

This change adds an AllOf combinator, and uses it to define a new
DoneDiagnosingChanges expectation that waits for all anticipated
diagnostic passes to complete. This should fix the race in
TestUnknownRevision.

We should apply a similar transformation throughout the regression test
suites. To make this easier, add a shorter AfterChange helper that
implements the common pattern.

Fixes golang/go#55070

Change-Id: Ie0e3c4701fba7b1d10de6b43d776562d198ffac9
Reviewed-on: https://go-review.googlesource.com/c/tools/+/448115
Reviewed-by: Alan Donovan <adonovan@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>

gopls/internal/lsp/regtest/expectation.go

@@ -7,6 +7,7 @@ package regtest
 import (
 	"fmt"
 	"regexp"
+	"sort"
 	"strings"
 
 	"golang.org/x/tools/gopls/internal/lsp"
@@ -130,6 +131,33 @@ func AnyOf(anyOf ...Expectation) *SimpleExpectation {
 	}
 }
 
+// AllOf expects that all given expectations are met.
+//
+// TODO(rfindley): the problem with these types of combinators (OnceMet, AnyOf
+// and AllOf) is that we lose the information of *why* they failed: the Awaiter
+// is not smart enough to look inside.
+//
+// Refactor the API such that the Check function is responsible for explaining
+// why an expectation failed. This should allow us to significantly improve
+// test output: we won't need to summarize state at all, as the verdict
+// explanation itself should describe clearly why the expectation not met.
+func AllOf(allOf ...Expectation) *SimpleExpectation {
+	check := func(s State) Verdict {
+		verdict := Met
+		for _, e := range allOf {
+			if v := e.Check(s); v > verdict {
+				verdict = v
+			}
+		}
+		return verdict
+	}
+	description := describeExpectations(allOf...)
+	return &SimpleExpectation{
+		check:       check,
+		description: fmt.Sprintf("All of:\n%s", description),
+	}
+}
+
 // ReadDiagnostics is an 'expectation' that is used to read diagnostics
 // atomically. It is intended to be used with 'OnceMet'.
 func ReadDiagnostics(fileName string, into *protocol.PublishDiagnosticsParams) *SimpleExpectation {
@@ -218,6 +246,54 @@ func ShowMessageRequest(title string) SimpleExpectation {
 	}
 }
 
+// DoneDiagnosingChanges expects that diagnostics are complete from common
+// change notifications: didOpen, didChange, didSave, didChangeWatchedFiles,
+// and didClose.
+//
+// This can be used when multiple notifications may have been sent, such as
+// when a didChange is immediately followed by a didSave. It is insufficient to
+// simply await NoOutstandingWork, because the LSP client has no control over
+// when the server starts processing a notification. Therefore, we must keep
+// track of
+func (e *Env) DoneDiagnosingChanges() Expectation {
+	stats := e.Editor.Stats()
+	statsBySource := map[lsp.ModificationSource]uint64{
+		lsp.FromDidOpen:               stats.DidOpen,
+		lsp.FromDidChange:             stats.DidChange,
+		lsp.FromDidSave:               stats.DidSave,
+		lsp.FromDidChangeWatchedFiles: stats.DidChangeWatchedFiles,
+		lsp.FromDidClose:              stats.DidClose,
+	}
+
+	var expected []lsp.ModificationSource
+	for k, v := range statsBySource {
+		if v > 0 {
+			expected = append(expected, k)
+		}
+	}
+
+	// Sort for stability.
+	sort.Slice(expected, func(i, j int) bool {
+		return expected[i] < expected[j]
+	})
+
+	var all []Expectation
+	for _, source := range expected {
+		all = append(all, CompletedWork(lsp.DiagnosticWorkTitle(source), statsBySource[source], true))
+	}
+
+	return AllOf(all...)
+}
+
+// AfterChange expects that the given expectations will be met after all
+// state-changing notifications have been processed by the server.
+func (e *Env) AfterChange(expectations ...Expectation) Expectation {
+	return OnceMet(
+		e.DoneDiagnosingChanges(),
+		expectations...,
+	)
+}
+
 // DoneWithOpen expects all didOpen notifications currently sent by the editor
 // to be completely processed.
 func (e *Env) DoneWithOpen() Expectation {

gopls/internal/lsp/text_synchronization.go

@@ -40,6 +40,9 @@ const (
 	// FromDidClose is a file modification caused by closing a file.
 	FromDidClose
 
+	// TODO: add FromDidChangeConfiguration, once configuration changes cause a
+	// new snapshot to be created.
+
 	// FromRegenerateCgo refers to file modifications caused by regenerating
 	// the cgo sources for the workspace.
 	FromRegenerateCgo

gopls/internal/regtest/modfile/modfile_test.go

@@ -633,7 +633,7 @@ func main() {
 
 			d := protocol.PublishDiagnosticsParams{}
 			env.Await(
-				OnceMet(
+				env.AfterChange(
 					// Make sure the diagnostic mentions the new version -- the old diagnostic is in the same place.
 					env.DiagnosticAtRegexpWithMessage("a/go.mod", "example.com v1.2.3", "example.com@v1.2.3"),
 					ReadDiagnostics("a/go.mod", &d),
@@ -646,8 +646,10 @@ func main() {
 			env.ApplyCodeAction(qfs[0]) // Arbitrarily pick a single fix to apply. Applying all of them seems to cause trouble in this particular test.
 			env.SaveBuffer("a/go.mod")  // Save to trigger diagnostics.
 			env.Await(
-				EmptyDiagnostics("a/go.mod"),
-				env.DiagnosticAtRegexp("a/main.go", "x = "),
+				env.AfterChange(
+					EmptyDiagnostics("a/go.mod"),
+					env.DiagnosticAtRegexp("a/main.go", "x = "),
+				),
 			)
 		})
 	})
@@ -677,17 +679,23 @@ func main() {
 		runner.Run(t, known, func(t *testing.T, env *Env) {
 			env.OpenFile("a/go.mod")
 			env.Await(
-				env.DiagnosticAtRegexp("a/main.go", "x = "),
+				env.AfterChange(
+					env.DiagnosticAtRegexp("a/main.go", "x = "),
+				),
 			)
 			env.RegexpReplace("a/go.mod", "v1.2.3", "v1.2.2")
 			env.Editor.SaveBuffer(env.Ctx, "a/go.mod") // go.mod changes must be on disk
 			env.Await(
-				env.DiagnosticAtRegexp("a/go.mod", "example.com v1.2.2"),
+				env.AfterChange(
+					env.DiagnosticAtRegexp("a/go.mod", "example.com v1.2.2"),
+				),
 			)
 			env.RegexpReplace("a/go.mod", "v1.2.2", "v1.2.3")
 			env.Editor.SaveBuffer(env.Ctx, "a/go.mod") // go.mod changes must be on disk
 			env.Await(
-				env.DiagnosticAtRegexp("a/main.go", "x = "),
+				env.AfterChange(
+					env.DiagnosticAtRegexp("a/main.go", "x = "),
+				),
 			)
 		})
 	})