gopls/internal/regtest: simplify expectation return values

Returning 'metBy' from expectations was not an ideal API, as extracting
the result value was cumbersome and error prone. It also forced quite a
bit of plumbing.

Using OnceMet, we can improve this by instead using a 'ReadDiagnostics'
expectation that reads diagnostics into a variable.

For golang/go#39384

Change-Id: Ia73e5b496089240df3200626e5696305cb507c9a
Reviewed-on: https://go-review.googlesource.com/c/tools/+/255126
Trust: Robert Findley <rfindley@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/gopls/internal/regtest/diagnostics_test.go b/gopls/internal/regtest/diagnostics_test.go
index d86ff44..790c0b1 100644
--- a/gopls/internal/regtest/diagnostics_test.go
+++ b/gopls/internal/regtest/diagnostics_test.go
@@ -259,13 +259,13 @@
 			env.Await(
 				EmptyDiagnostics("main.go"),
 			)
-			metBy := env.Await(
-				env.DiagnosticAtRegexp("bob/bob.go", "x"),
+			var d protocol.PublishDiagnosticsParams
+			env.Await(
+				OnceMet(
+					env.DiagnosticAtRegexp("bob/bob.go", "x"),
+					ReadDiagnostics("bob/bob.go", &d),
+				),
 			)
-			d, ok := metBy[0].(*protocol.PublishDiagnosticsParams)
-			if !ok {
-				t.Fatalf("unexpected met by result %v (%T)", metBy, metBy)
-			}
 			if len(d.Diagnostics) != 1 {
 				t.Fatalf("expected 1 diagnostic, got %v", len(d.Diagnostics))
 			}
@@ -490,13 +490,13 @@
 	runner.Run(t, generated, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
 		original := env.ReadWorkspaceFile("main.go")
-		metBy := env.Await(
-			DiagnosticAt("main.go", 5, 8),
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				DiagnosticAt("main.go", 5, 8),
+				ReadDiagnostics("main.go", &d),
+			),
 		)
-		d, ok := metBy[0].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("unexpected met by result %v (%T)", metBy, metBy)
-		}
 		// Apply fixes and save the buffer.
 		env.ApplyQuickFixes("main.go", d.Diagnostics)
 		env.SaveBuffer("main.go")
@@ -649,13 +649,13 @@
 		// "github.com/ardanlabs/conf" to the go.mod file.
 		env.OpenFile("go.mod")
 		env.OpenFile("main.go")
-		metBy := env.Await(
-			env.DiagnosticAtRegexp("main.go", `"github.com/ardanlabs/conf"`),
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				env.DiagnosticAtRegexp("main.go", `"github.com/ardanlabs/conf"`),
+				ReadDiagnostics("main.go", &d),
+			),
 		)
-		d, ok := metBy[0].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("unexpected type for metBy (%T)", metBy)
-		}
 		env.ApplyQuickFixes("main.go", d.Diagnostics)
 		env.SaveBuffer("go.mod")
 		env.Await(
@@ -669,14 +669,13 @@
 		)
 		env.SaveBuffer("main.go")
 		// Expect a diagnostic and fix to remove the dependency in the go.mod.
-		metBy = env.Await(
-			EmptyDiagnostics("main.go"),
-			env.DiagnosticAtRegexp("go.mod", "require github.com/ardanlabs/conf"),
+		env.Await(EmptyDiagnostics("main.go"))
+		env.Await(
+			OnceMet(
+				env.DiagnosticAtRegexp("go.mod", "require github.com/ardanlabs/conf"),
+				ReadDiagnostics("go.mod", &d),
+			),
 		)
-		d, ok = metBy[1].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("unexpected type for metBy (%T)", metBy)
-		}
 		env.ApplyQuickFixes("go.mod", d.Diagnostics)
 		env.SaveBuffer("go.mod")
 		env.Await(
@@ -715,13 +714,13 @@
 }
 `))
 		env.SaveBuffer("main.go")
-		metBy := env.Await(
-			env.DiagnosticAtRegexp("main.go", `"github.com/ardanlabs/conf"`),
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				env.DiagnosticAtRegexp("main.go", `"github.com/ardanlabs/conf"`),
+				ReadDiagnostics("main.go", &d),
+			),
 		)
-		d, ok := metBy[0].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("unexpected type for diagnostics (%T)", d)
-		}
 		env.ApplyQuickFixes("main.go", d.Diagnostics)
 		env.Await(
 			EmptyDiagnostics("main.go"),
diff --git a/gopls/internal/regtest/env.go b/gopls/internal/regtest/env.go
index d53abc0..f1b958d 100644
--- a/gopls/internal/regtest/env.go
+++ b/gopls/internal/regtest/env.go
@@ -233,7 +233,7 @@
 
 func (e *Env) checkConditionsLocked() {
 	for id, condition := range e.waiters {
-		if v, _, _ := checkExpectations(e.state, condition.expectations); v != Unmet {
+		if v, _ := checkExpectations(e.state, condition.expectations); v != Unmet {
 			delete(e.waiters, id)
 			condition.verdict <- v
 		}
@@ -241,21 +241,17 @@
 }
 
 // checkExpectations reports whether s meets all expectations.
-func checkExpectations(s State, expectations []Expectation) (Verdict, string, []interface{}) {
+func checkExpectations(s State, expectations []Expectation) (Verdict, string) {
 	finalVerdict := Met
-	var metBy []interface{}
 	var summary strings.Builder
 	for _, e := range expectations {
-		v, mb := e.Check(s)
-		if v == Met {
-			metBy = append(metBy, mb)
-		}
+		v := e.Check(s)
 		if v > finalVerdict {
 			finalVerdict = v
 		}
 		summary.WriteString(fmt.Sprintf("\t%v: %s\n", v, e.Description()))
 	}
-	return finalVerdict, summary.String(), metBy
+	return finalVerdict, summary.String()
 }
 
 // DiagnosticsFor returns the current diagnostics for the file. It is useful
@@ -269,16 +265,16 @@
 
 // Await waits for all expectations to simultaneously be met. It should only be
 // called from the main test goroutine.
-func (e *Env) Await(expectations ...Expectation) []interface{} {
+func (e *Env) Await(expectations ...Expectation) {
 	e.T.Helper()
 	e.mu.Lock()
 	// Before adding the waiter, we check if the condition is currently met or
 	// failed to avoid a race where the condition was realized before Await was
 	// called.
-	switch verdict, summary, metBy := checkExpectations(e.state, expectations); verdict {
+	switch verdict, summary := checkExpectations(e.state, expectations); verdict {
 	case Met:
 		e.mu.Unlock()
-		return metBy
+		return
 	case Unmeetable:
 		e.mu.Unlock()
 		e.T.Fatalf("unmeetable expectations:\n%s\nstate:\n%v", summary, e.state)
@@ -302,12 +298,11 @@
 	}
 	e.mu.Lock()
 	defer e.mu.Unlock()
-	_, summary, metBy := checkExpectations(e.state, expectations)
+	_, summary := checkExpectations(e.state, expectations)
 
 	// Debugging an unmet expectation can be tricky, so we put some effort into
 	// nicely formatting the failure.
 	if err != nil {
 		e.T.Fatalf("waiting on:\n%s\nerr:%v\n\nstate:\n%v", summary, err, e.state)
 	}
-	return metBy
 }
diff --git a/gopls/internal/regtest/expectation.go b/gopls/internal/regtest/expectation.go
index e8697f6..697363f 100644
--- a/gopls/internal/regtest/expectation.go
+++ b/gopls/internal/regtest/expectation.go
@@ -19,7 +19,7 @@
 type Expectation interface {
 	// Check determines whether the state of the editor satisfies the
 	// expectation, returning the results that met the condition.
-	Check(State) (Verdict, interface{})
+	Check(State) Verdict
 	// Description is a human-readable description of the expectation.
 	Description() string
 }
@@ -61,12 +61,12 @@
 
 // SimpleExpectation holds an arbitrary check func, and implements the Expectation interface.
 type SimpleExpectation struct {
-	check       func(State) (Verdict, interface{})
+	check       func(State) Verdict
 	description string
 }
 
 // Check invokes e.check.
-func (e SimpleExpectation) Check(s State) (Verdict, interface{}) {
+func (e SimpleExpectation) Check(s State) Verdict {
 	return e.check(s)
 }
 
@@ -78,18 +78,18 @@
 // OnceMet returns an Expectation that, once the precondition is met, asserts
 // that mustMeet is met.
 func OnceMet(precondition Expectation, mustMeet Expectation) *SimpleExpectation {
-	check := func(s State) (Verdict, interface{}) {
-		switch pre, _ := precondition.Check(s); pre {
+	check := func(s State) Verdict {
+		switch pre := precondition.Check(s); pre {
 		case Unmeetable:
-			return Unmeetable, nil
+			return Unmeetable
 		case Met:
-			verdict, metBy := mustMeet.Check(s)
+			verdict := mustMeet.Check(s)
 			if verdict != Met {
-				return Unmeetable, metBy
+				return Unmeetable
 			}
-			return Met, metBy
+			return Met
 		default:
-			return Unmet, nil
+			return Unmet
 		}
 	}
 	return &SimpleExpectation{
@@ -98,14 +98,31 @@
 	}
 }
 
+// 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 {
+	check := func(s State) Verdict {
+		diags, ok := s.diagnostics[fileName]
+		if !ok {
+			return Unmeetable
+		}
+		*into = *diags
+		return Met
+	}
+	return &SimpleExpectation{
+		check:       check,
+		description: fmt.Sprintf("read diagnostics for %q", fileName),
+	}
+}
+
 // NoOutstandingWork asserts that there is no work initiated using the LSP
 // $/progress API that has not completed.
 func NoOutstandingWork() SimpleExpectation {
-	check := func(s State) (Verdict, interface{}) {
+	check := func(s State) Verdict {
 		if len(s.outstandingWork) == 0 {
-			return Met, nil
+			return Met
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return SimpleExpectation{
 		check:       check,
@@ -115,11 +132,11 @@
 
 // NoShowMessage asserts that the editor has not received a ShowMessage.
 func NoShowMessage() SimpleExpectation {
-	check := func(s State) (Verdict, interface{}) {
+	check := func(s State) Verdict {
 		if len(s.showMessage) == 0 {
-			return Met, "no ShowMessage"
+			return Met
 		}
-		return Unmeetable, nil
+		return Unmeetable
 	}
 	return SimpleExpectation{
 		check:       check,
@@ -130,13 +147,13 @@
 // ShownMessage asserts that the editor has received a ShownMessage with the
 // given title.
 func ShownMessage(title string) SimpleExpectation {
-	check := func(s State) (Verdict, interface{}) {
+	check := func(s State) Verdict {
 		for _, m := range s.showMessage {
 			if strings.Contains(m.Message, title) {
-				return Met, m
+				return Met
 			}
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return SimpleExpectation{
 		check:       check,
@@ -147,19 +164,19 @@
 // ShowMessageRequest asserts that the editor has received a ShowMessageRequest
 // with an action item that has the given title.
 func ShowMessageRequest(title string) SimpleExpectation {
-	check := func(s State) (Verdict, interface{}) {
+	check := func(s State) Verdict {
 		if len(s.showMessageRequest) == 0 {
-			return Unmet, nil
+			return Unmet
 		}
 		// Only check the most recent one.
 		m := s.showMessageRequest[len(s.showMessageRequest)-1]
 		if len(m.Actions) == 0 || len(m.Actions) > 1 {
-			return Unmet, nil
+			return Unmet
 		}
 		if m.Actions[0].Title == title {
-			return Met, m.Actions[0]
+			return Met
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return SimpleExpectation{
 		check:       check,
@@ -172,11 +189,11 @@
 // Since the Progress API doesn't include any hidden metadata, we must use the
 // progress notification title to identify the work we expect to be completed.
 func CompletedWork(title string, atLeast int) SimpleExpectation {
-	check := func(s State) (Verdict, interface{}) {
+	check := func(s State) Verdict {
 		if s.completedWork[title] >= atLeast {
-			return Met, title
+			return Met
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return SimpleExpectation{
 		check:       check,
@@ -187,12 +204,12 @@
 // LogExpectation is an expectation on the log messages received by the editor
 // from gopls.
 type LogExpectation struct {
-	check       func([]*protocol.LogMessageParams) (Verdict, interface{})
+	check       func([]*protocol.LogMessageParams) Verdict
 	description string
 }
 
 // Check implements the Expectation interface.
-func (e LogExpectation) Check(s State) (Verdict, interface{}) {
+func (e LogExpectation) Check(s State) Verdict {
 	return e.check(s.logs)
 }
 
@@ -214,7 +231,7 @@
 	if err != nil {
 		panic(err)
 	}
-	check := func(msgs []*protocol.LogMessageParams) (Verdict, interface{}) {
+	check := func(msgs []*protocol.LogMessageParams) Verdict {
 		var found int
 		for _, msg := range msgs {
 			if msg.Type == typ && rec.Match([]byte(msg.Message)) {
@@ -222,9 +239,9 @@
 			}
 		}
 		if found == count {
-			return Met, nil
+			return Met
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return LogExpectation{
 		check:       check,
@@ -244,16 +261,16 @@
 			panic(err)
 		}
 	}
-	check := func(msgs []*protocol.LogMessageParams) (Verdict, interface{}) {
+	check := func(msgs []*protocol.LogMessageParams) Verdict {
 		for _, msg := range msgs {
 			if msg.Type != typ {
 				continue
 			}
 			if r == nil || r.Match([]byte(msg.Message)) {
-				return Unmeetable, nil
+				return Unmeetable
 			}
 		}
-		return Met, nil
+		return Met
 	}
 	return LogExpectation{
 		check:       check,
@@ -264,12 +281,12 @@
 // RegistrationExpectation is an expectation on the capability registrations
 // received by the editor from gopls.
 type RegistrationExpectation struct {
-	check       func([]*protocol.RegistrationParams) (Verdict, interface{})
+	check       func([]*protocol.RegistrationParams) Verdict
 	description string
 }
 
 // Check implements the Expectation interface.
-func (e RegistrationExpectation) Check(s State) (Verdict, interface{}) {
+func (e RegistrationExpectation) Check(s State) Verdict {
 	return e.check(s.registrations)
 }
 
@@ -285,15 +302,15 @@
 	if err != nil {
 		panic(err)
 	}
-	check := func(params []*protocol.RegistrationParams) (Verdict, interface{}) {
+	check := func(params []*protocol.RegistrationParams) Verdict {
 		for _, p := range params {
 			for _, r := range p.Registrations {
 				if rec.Match([]byte(r.Method)) {
-					return Met, r
+					return Met
 				}
 			}
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return RegistrationExpectation{
 		check:       check,
@@ -304,12 +321,12 @@
 // UnregistrationExpectation is an expectation on the capability
 // unregistrations received by the editor from gopls.
 type UnregistrationExpectation struct {
-	check       func([]*protocol.UnregistrationParams) (Verdict, interface{})
+	check       func([]*protocol.UnregistrationParams) Verdict
 	description string
 }
 
 // Check implements the Expectation interface.
-func (e UnregistrationExpectation) Check(s State) (Verdict, interface{}) {
+func (e UnregistrationExpectation) Check(s State) Verdict {
 	return e.check(s.unregistrations)
 }
 
@@ -325,15 +342,15 @@
 	if err != nil {
 		panic(err)
 	}
-	check := func(params []*protocol.UnregistrationParams) (Verdict, interface{}) {
+	check := func(params []*protocol.UnregistrationParams) Verdict {
 		for _, p := range params {
 			for _, r := range p.Unregisterations {
 				if rec.Match([]byte(r.Method)) {
-					return Met, r
+					return Met
 				}
 			}
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return UnregistrationExpectation{
 		check:       check,
@@ -354,11 +371,11 @@
 }
 
 // Check implements the Expectation interface.
-func (e DiagnosticExpectation) Check(s State) (Verdict, interface{}) {
+func (e DiagnosticExpectation) Check(s State) Verdict {
 	if diags, ok := s.diagnostics[e.path]; ok && e.isMet(diags) {
-		return Met, diags
+		return Met
 	}
-	return Unmet, nil
+	return Unmet
 }
 
 // Description implements the Expectation interface.
@@ -369,11 +386,11 @@
 // EmptyDiagnostics asserts that empty diagnostics are sent for the
 // workspace-relative path name.
 func EmptyDiagnostics(name string) Expectation {
-	check := func(s State) (Verdict, interface{}) {
+	check := func(s State) Verdict {
 		if diags := s.diagnostics[name]; diags != nil && len(diags.Diagnostics) == 0 {
-			return Met, nil
+			return Met
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return SimpleExpectation{
 		check:       check,
@@ -386,11 +403,11 @@
 // with a OnceMet, as it has to check that all outstanding diagnostics have
 // already been delivered.
 func NoDiagnostics(name string) Expectation {
-	check := func(s State) (Verdict, interface{}) {
+	check := func(s State) Verdict {
 		if _, ok := s.diagnostics[name]; !ok {
-			return Met, nil
+			return Met
 		}
-		return Unmet, nil
+		return Unmet
 	}
 	return SimpleExpectation{
 		check:       check,
diff --git a/gopls/internal/regtest/imports_test.go b/gopls/internal/regtest/imports_test.go
index f570fa1..6b95b17 100644
--- a/gopls/internal/regtest/imports_test.go
+++ b/gopls/internal/regtest/imports_test.go
@@ -191,13 +191,13 @@
 	run(t, pkg, func(t *testing.T, env *Env) {
 		env.Await(InitialWorkspaceLoad)
 		env.OpenFile("a/a.go")
-		metBy := env.Await(
-			env.DiagnosticAtRegexp("a/a.go", "os.Stat"),
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				env.DiagnosticAtRegexp("a/a.go", "os.Stat"),
+				ReadDiagnostics("a/a.go", &d),
+			),
 		)
-		d, ok := metBy[0].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("expected *protocol.PublishDiagnosticsParams, got %v (%T)", d, d)
-		}
 		env.ApplyQuickFixes("a/a.go", d.Diagnostics)
 		env.Await(
 			EmptyDiagnostics("a/a.go"),
diff --git a/gopls/internal/regtest/modfile_test.go b/gopls/internal/regtest/modfile_test.go
index bfae9dc..a7f33ab 100644
--- a/gopls/internal/regtest/modfile_test.go
+++ b/gopls/internal/regtest/modfile_test.go
@@ -126,17 +126,14 @@
 `
 	runner.Run(t, mod, func(t *testing.T, env *Env) {
 		env.OpenFile("go.mod")
-		d := env.Await(
-			env.DiagnosticAtRegexp("go.mod", "// indirect"),
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				env.DiagnosticAtRegexp("go.mod", "// indirect"),
+				ReadDiagnostics("go.mod", &d),
+			),
 		)
-		if len(d) == 0 {
-			t.Fatalf("no diagnostics")
-		}
-		params, ok := d[0].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("expected diagnostic of type PublishDiagnosticParams, got %T", d[0])
-		}
-		env.ApplyQuickFixes("go.mod", params.Diagnostics)
+		env.ApplyQuickFixes("go.mod", d.Diagnostics)
 		if got := env.Editor.BufferText("go.mod"); got != want {
 			t.Fatalf("unexpected go.mod content:\n%s", tests.Diff(want, got))
 		}
@@ -184,17 +181,14 @@
 	runner.Run(t, repro, func(t *testing.T, env *Env) {
 		env.OpenFile("go.mod")
 		env.OpenFile("main.go")
-		d := env.Await(
-			env.DiagnosticAtRegexp("main.go", `"github.com/esimov/caire"`),
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				env.DiagnosticAtRegexp("main.go", `"github.com/esimov/caire"`),
+				ReadDiagnostics("main.go", &d),
+			),
 		)
-		if len(d) == 0 {
-			t.Fatalf("no diagnostics")
-		}
-		params, ok := d[0].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("expected diagnostic of type PublishDiagnosticParams, got %T", d[0])
-		}
-		env.ApplyQuickFixes("main.go", params.Diagnostics)
+		env.ApplyQuickFixes("main.go", d.Diagnostics)
 		want := `module mod.com
 
 go 1.14
@@ -308,14 +302,13 @@
 	runner.Run(t, pkg, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
 		env.OpenFile("go.mod")
-		metBy := env.Await(
-			DiagnosticAt("go.mod", 0, 0),
-			NoDiagnostics("main.go"),
+		var d protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				DiagnosticAt("go.mod", 0, 0),
+				ReadDiagnostics("go.mod", &d),
+			),
 		)
-		d, ok := metBy[0].(*protocol.PublishDiagnosticsParams)
-		if !ok {
-			t.Fatalf("unexpected type for metBy (%T)", metBy)
-		}
 		env.ApplyQuickFixes("main.go", d.Diagnostics)
 		const want = `module mod.com