gopls/internal/regtest: eliminate DiagnosticsFor

As noted in the associated TODO, DiagnosticsFor is racy. Remove it,
replacing uses with the atomic ReadDiagnostics expectation.

Also replace nearby Awaits with AfterChange, to make them fail eagerly.

Change-Id: Idb5c2420cd0d9b28e701343bc8d2883eea4e144a
Reviewed-on: https://go-review.googlesource.com/c/tools/+/461875
Reviewed-by: Alan Donovan <adonovan@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Run-TryBot: Robert Findley <rfindley@google.com>
gopls-CI: kokoro <noreply+kokoro@google.com>
diff --git a/gopls/internal/lsp/regtest/env.go b/gopls/internal/lsp/regtest/env.go
index 192e8ed..c3f48bf 100644
--- a/gopls/internal/lsp/regtest/env.go
+++ b/gopls/internal/lsp/regtest/env.go
@@ -302,18 +302,6 @@
 	return finalVerdict, summary.String()
 }
 
-// DiagnosticsFor returns the current diagnostics for the file. It is useful
-// after waiting on AnyDiagnosticAtCurrentVersion, when the desired diagnostic
-// is not simply described by DiagnosticAt.
-//
-// TODO(rfindley): this method is inherently racy. Replace usages of this
-// method with the atomic OnceMet(..., ReadDiagnostics) pattern.
-func (a *Awaiter) DiagnosticsFor(name string) *protocol.PublishDiagnosticsParams {
-	a.mu.Lock()
-	defer a.mu.Unlock()
-	return a.state.diagnostics[name]
-}
-
 func (e *Env) Await(expectations ...Expectation) {
 	e.T.Helper()
 	if err := e.Awaiter.Await(e.Ctx, expectations...); err != nil {
diff --git a/gopls/internal/regtest/diagnostics/diagnostics_test.go b/gopls/internal/regtest/diagnostics/diagnostics_test.go
index c7f091c..84bd5d4 100644
--- a/gopls/internal/regtest/diagnostics/diagnostics_test.go
+++ b/gopls/internal/regtest/diagnostics/diagnostics_test.go
@@ -1280,23 +1280,23 @@
 	Run(t, dir, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
 		env.OpenFile("other.go")
-		x := env.Awaiter.DiagnosticsFor("main.go")
-		if x == nil {
-			t.Fatalf("expected 1 diagnostic, got none")
+		var mainDiags, otherDiags protocol.PublishDiagnosticsParams
+		env.AfterChange(
+			ReadDiagnostics("main.go", &mainDiags),
+			ReadDiagnostics("other.go", &otherDiags),
+		)
+		if len(mainDiags.Diagnostics) != 1 {
+			t.Fatalf("main.go, got %d diagnostics, expected 1", len(mainDiags.Diagnostics))
 		}
-		if len(x.Diagnostics) != 1 {
-			t.Fatalf("main.go, got %d diagnostics, expected 1", len(x.Diagnostics))
+		keep := mainDiags.Diagnostics[0]
+		if len(otherDiags.Diagnostics) != 1 {
+			t.Fatalf("other.go: got %d diagnostics, expected 1", len(otherDiags.Diagnostics))
 		}
-		keep := x.Diagnostics[0]
-		y := env.Awaiter.DiagnosticsFor("other.go")
-		if len(y.Diagnostics) != 1 {
-			t.Fatalf("other.go: got %d diagnostics, expected 1", len(y.Diagnostics))
-		}
-		if len(y.Diagnostics[0].RelatedInformation) != 1 {
-			t.Fatalf("got %d RelatedInformations, expected 1", len(y.Diagnostics[0].RelatedInformation))
+		if len(otherDiags.Diagnostics[0].RelatedInformation) != 1 {
+			t.Fatalf("got %d RelatedInformations, expected 1", len(otherDiags.Diagnostics[0].RelatedInformation))
 		}
 		// check that the RelatedInformation matches the error from main.go
-		c := y.Diagnostics[0].RelatedInformation[0]
+		c := otherDiags.Diagnostics[0].RelatedInformation[0]
 		if c.Location.Range != keep.Range {
 			t.Errorf("locations don't match. Got %v expected %v", c.Location.Range, keep.Range)
 		}
diff --git a/gopls/internal/regtest/diagnostics/undeclared_test.go b/gopls/internal/regtest/diagnostics/undeclared_test.go
index c3456fa..bba72c4 100644
--- a/gopls/internal/regtest/diagnostics/undeclared_test.go
+++ b/gopls/internal/regtest/diagnostics/undeclared_test.go
@@ -44,23 +44,30 @@
 
 		// 'x' is undeclared, but still necessary.
 		env.OpenFile("a/a.go")
-		env.Await(env.DiagnosticAtRegexp("a/a.go", "x"))
-		diags := env.Awaiter.DiagnosticsFor("a/a.go")
-		if got := len(diags.Diagnostics); got != 1 {
+		var adiags protocol.PublishDiagnosticsParams
+		env.AfterChange(
+			env.DiagnosticAtRegexp("a/a.go", "x"),
+			ReadDiagnostics("a/a.go", &adiags),
+		)
+		env.Await()
+		if got := len(adiags.Diagnostics); got != 1 {
 			t.Errorf("len(Diagnostics) = %d, want 1", got)
 		}
-		if diag := diags.Diagnostics[0]; isUnnecessary(diag) {
+		if diag := adiags.Diagnostics[0]; isUnnecessary(diag) {
 			t.Errorf("%v tagged unnecessary, want necessary", diag)
 		}
 
 		// 'y = y' is pointless, and should be detected as unnecessary.
 		env.OpenFile("b/b.go")
-		env.Await(env.DiagnosticAtRegexp("b/b.go", "y = y"))
-		diags = env.Awaiter.DiagnosticsFor("b/b.go")
-		if got := len(diags.Diagnostics); got != 1 {
+		var bdiags protocol.PublishDiagnosticsParams
+		env.AfterChange(
+			env.DiagnosticAtRegexp("b/b.go", "y = y"),
+			ReadDiagnostics("b/b.go", &bdiags),
+		)
+		if got := len(bdiags.Diagnostics); got != 1 {
 			t.Errorf("len(Diagnostics) = %d, want 1", got)
 		}
-		if diag := diags.Diagnostics[0]; !isUnnecessary(diag) {
+		if diag := bdiags.Diagnostics[0]; !isUnnecessary(diag) {
 			t.Errorf("%v tagged necessary, want unnecessary", diag)
 		}
 	})
diff --git a/gopls/internal/regtest/modfile/modfile_test.go b/gopls/internal/regtest/modfile/modfile_test.go
index 1d2ade2..e15fa4c 100644
--- a/gopls/internal/regtest/modfile/modfile_test.go
+++ b/gopls/internal/regtest/modfile/modfile_test.go
@@ -554,13 +554,15 @@
 	}.Run(t, files, func(t *testing.T, env *Env) {
 		env.OpenFile("a/main.go")
 		env.OpenFile("a/go.mod")
-		env.Await(
+		var modDiags protocol.PublishDiagnosticsParams
+		env.AfterChange(
 			// We would like for the error to appear in the v2 module, but
 			// as of writing non-workspace packages are not diagnosed.
 			env.DiagnosticAtRegexpWithMessage("a/main.go", `"example.com/blah/v2"`, "cannot find module providing"),
 			env.DiagnosticAtRegexpWithMessage("a/go.mod", `require example.com/blah/v2`, "cannot find module providing"),
+			ReadDiagnostics("a/go.mod", &modDiags),
 		)
-		env.ApplyQuickFixes("a/go.mod", env.Awaiter.DiagnosticsFor("a/go.mod").Diagnostics)
+		env.ApplyQuickFixes("a/go.mod", modDiags.Diagnostics)
 		const want = `module mod.com
 
 go 1.12
@@ -571,7 +573,7 @@
 )
 `
 		env.SaveBuffer("a/go.mod")
-		env.Await(EmptyDiagnostics("a/main.go"))
+		env.AfterChange(EmptyDiagnostics("a/main.go"))
 		if got := env.BufferText("a/go.mod"); got != want {
 			t.Fatalf("suggested fixes failed:\n%s", compare.Text(want, got))
 		}
diff --git a/gopls/internal/regtest/template/template_test.go b/gopls/internal/regtest/template/template_test.go
index 91c704d..afe5797 100644
--- a/gopls/internal/regtest/template/template_test.go
+++ b/gopls/internal/regtest/template/template_test.go
@@ -70,8 +70,15 @@
 		},
 	).Run(t, files, func(t *testing.T, env *Env) {
 		// TODO: can we move this diagnostic onto {{}}?
-		env.Await(env.DiagnosticAtRegexp("hello.tmpl", "()Hello {{}}"))
-		d := env.Awaiter.DiagnosticsFor("hello.tmpl").Diagnostics // issue 50786: check for Source
+		var diags protocol.PublishDiagnosticsParams
+		env.Await(
+			OnceMet(
+				InitialWorkspaceLoad,
+				env.DiagnosticAtRegexp("hello.tmpl", "()Hello {{}}"),
+				ReadDiagnostics("hello.tmpl", &diags),
+			),
+		)
+		d := diags.Diagnostics // issue 50786: check for Source
 		if len(d) != 1 {
 			t.Errorf("expected 1 diagnostic, got %d", len(d))
 			return
@@ -91,7 +98,7 @@
 		}
 
 		env.WriteWorkspaceFile("hello.tmpl", "{{range .Planets}}\nHello {{.}}\n{{end}}")
-		env.Await(EmptyDiagnostics("hello.tmpl"))
+		env.AfterChange(EmptyDiagnostics("hello.tmpl"))
 	})
 }