gopls/internal/regtest: make AfterChange do the awaiting

The only expected use of the AfterChange helper is in the pattern
env.Await(env.AfterChange(...)). Simplify this pattern by having
AfterChange do the awaiting.

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

gopls/internal/lsp/regtest/expectation.go

@@ -287,10 +287,16 @@ func (e *Env) DoneDiagnosingChanges() Expectation {
 
 // 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...,
+//
+// It awaits the completion of all anticipated work before checking the given
+// expectations.
+func (e *Env) AfterChange(expectations ...Expectation) {
+	e.T.Helper()
+	e.Await(
+		OnceMet(
+			e.DoneDiagnosingChanges(),
+			expectations...,
+		),
 	)
 }
 

gopls/internal/regtest/modfile/modfile_test.go

@@ -632,12 +632,10 @@ func main() {
 			env.SaveBuffer("a/go.mod") // Save to trigger diagnostics.
 
 			d := protocol.PublishDiagnosticsParams{}
-			env.Await(
-				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),
-				),
+			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),
 			)
 			qfs := env.GetQuickFixes("a/go.mod", d.Diagnostics)
 			if len(qfs) == 0 {
@@ -645,11 +643,9 @@ 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(
-				env.AfterChange(
-					EmptyDiagnostics("a/go.mod"),
-					env.DiagnosticAtRegexp("a/main.go", "x = "),
-				),
+			env.AfterChange(
+				EmptyDiagnostics("a/go.mod"),
+				env.DiagnosticAtRegexp("a/main.go", "x = "),
 			)
 		})
 	})
@@ -678,24 +674,18 @@ func main() {
 	t.Run("good", func(t *testing.T) {
 		runner.Run(t, known, func(t *testing.T, env *Env) {
 			env.OpenFile("a/go.mod")
-			env.Await(
-				env.AfterChange(
-					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.AfterChange(
-					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.AfterChange(
-					env.DiagnosticAtRegexp("a/main.go", "x = "),
-				),
+			env.AfterChange(
+				env.DiagnosticAtRegexp("a/main.go", "x = "),
 			)
 		})
 	})

gopls/internal/regtest/workspace/broken_test.go

@@ -206,12 +206,10 @@ package b
 				env.OpenFile("a/a.go")
 				env.OpenFile("a/empty.go")
 				env.OpenFile("b/go.mod")
-				env.Await(
-					env.AfterChange(
-						env.DiagnosticAtRegexp("a/a.go", "package a"),
-						env.DiagnosticAtRegexp("b/go.mod", "module b.com"),
-						OutstandingWork(lsp.WorkspaceLoadFailure, msg),
-					),
+				env.AfterChange(
+					env.DiagnosticAtRegexp("a/a.go", "package a"),
+					env.DiagnosticAtRegexp("b/go.mod", "module b.com"),
+					OutstandingWork(lsp.WorkspaceLoadFailure, msg),
 				)
 
 				// Changing the workspace folders to the valid modules should resolve