gopls/internal/regtest: use AfterChange in more places

Update all the places (I think) where we await diagnostics unbounded,
except for a couple that can't yet be made eager due to lacking
instrumentation (for example, around didChangeConfiguration).

Updates golang/go#39384

Change-Id: Ib00d28ce0350737d39a123fd200fd435cc89f967
Reviewed-on: https://go-review.googlesource.com/c/tools/+/462495
gopls-CI: kokoro <noreply+kokoro@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Alan Donovan <adonovan@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
diff --git a/gopls/internal/regtest/completion/completion_test.go b/gopls/internal/regtest/completion/completion_test.go
index 2f9b550..6d0bfc5 100644
--- a/gopls/internal/regtest/completion/completion_test.go
+++ b/gopls/internal/regtest/completion/completion_test.go
@@ -173,9 +173,7 @@
 			Run(t, files, func(t *testing.T, env *Env) {
 				if tc.content != nil {
 					env.WriteWorkspaceFile(tc.filename, *tc.content)
-					env.Await(
-						env.DoneWithChangeWatchedFiles(),
-					)
+					env.Await(env.DoneWithChangeWatchedFiles())
 				}
 				env.OpenFile(tc.filename)
 				completions := env.Completion(tc.filename, env.RegexpSearch(tc.filename, tc.triggerRegexp))
@@ -318,7 +316,7 @@
 		env.AcceptCompletion("main.go", pos, item)
 
 		// Await the diagnostics to add example.com/blah to the go.mod file.
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("main.go", `"example.com/blah"`)),
 		)
 	})
diff --git a/gopls/internal/regtest/diagnostics/diagnostics_test.go b/gopls/internal/regtest/diagnostics/diagnostics_test.go
index 4ce509c..800274b 100644
--- a/gopls/internal/regtest/diagnostics/diagnostics_test.go
+++ b/gopls/internal/regtest/diagnostics/diagnostics_test.go
@@ -73,9 +73,7 @@
 	log.Println()
 }
 `)
-		env.Await(
-			Diagnostics(env.AtRegexp("main.go", "log")),
-		)
+		env.AfterChange(Diagnostics(env.AtRegexp("main.go", "log")))
 		env.SaveBuffer("main.go")
 		env.AfterChange(NoDiagnostics(ForFile("main.go")))
 	})
@@ -88,7 +86,7 @@
 `
 	Run(t, brokenFile, func(t *testing.T, env *Env) {
 		env.CreateBuffer("broken.go", brokenFile)
-		env.Await(Diagnostics(env.AtRegexp("broken.go", "\"abc")))
+		env.AfterChange(Diagnostics(env.AtRegexp("broken.go", "\"abc")))
 	})
 }
 
@@ -128,13 +126,13 @@
 func TestDiagnosticClearingOnDelete_Issue37049(t *testing.T) {
 	Run(t, badPackage, func(t *testing.T, env *Env) {
 		env.OpenFile("a.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("a.go", "a = 1")),
 			Diagnostics(env.AtRegexp("b.go", "a = 2")),
 		)
 		env.RemoveWorkspaceFile("b.go")
 
-		env.Await(
+		env.AfterChange(
 			NoDiagnostics(ForFile("a.go")),
 			NoDiagnostics(ForFile("b.go")),
 		)
@@ -214,7 +212,7 @@
 // not break the workspace.
 func TestDeleteTestVariant(t *testing.T) {
 	Run(t, test38878, func(t *testing.T, env *Env) {
-		env.Await(Diagnostics(env.AtRegexp("a_test.go", `f\((3)\)`)))
+		env.AfterChange(Diagnostics(env.AtRegexp("a_test.go", `f\((3)\)`)))
 		env.RemoveWorkspaceFile("a_test.go")
 		env.AfterChange(NoDiagnostics(ForFile("a_test.go")))
 
@@ -231,7 +229,7 @@
 func TestDeleteTestVariant_DiskOnly(t *testing.T) {
 	Run(t, test38878, func(t *testing.T, env *Env) {
 		env.OpenFile("a_test.go")
-		env.Await(Diagnostics(AtPosition("a_test.go", 5, 3)))
+		env.AfterChange(Diagnostics(AtPosition("a_test.go", 5, 3)))
 		env.Sandbox.Workdir.RemoveFile(context.Background(), "a_test.go")
 		env.AfterChange(Diagnostics(AtPosition("a_test.go", 5, 3)))
 	})
@@ -281,7 +279,8 @@
 	})
 	t.Run("initialized", func(t *testing.T) {
 		Run(t, noMod, func(t *testing.T, env *Env) {
-			env.Await(
+			env.OnceMet(
+				InitialWorkspaceLoad,
 				Diagnostics(env.AtRegexp("main.go", `"mod.com/bob"`)),
 			)
 			env.RunGoCommand("mod", "init", "mod.com")
@@ -296,7 +295,8 @@
 		WithOptions(
 			Modes(Default),
 		).Run(t, noMod, func(t *testing.T, env *Env) {
-			env.Await(
+			env.OnceMet(
+				InitialWorkspaceLoad,
 				Diagnostics(env.AtRegexp("main.go", `"mod.com/bob"`)),
 			)
 			if err := env.Sandbox.RunGoCommand(env.Ctx, "", "mod", []string{"init", "mod.com"}, true); err != nil {
@@ -345,7 +345,7 @@
 
 	Run(t, testPackage, func(t *testing.T, env *Env) {
 		env.OpenFile("lib_test.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(AtPosition("lib_test.go", 10, 2)),
 			Diagnostics(AtPosition("lib_test.go", 11, 2)),
 		)
@@ -420,8 +420,12 @@
 		env.OpenFile("print.go")
 		// Check that gopackages correctly loaded this dependency. We should get a
 		// diagnostic for the wrong formatting type.
-		// TODO: we should be able to easily also match the diagnostic message.
-		env.Await(Diagnostics(env.AtRegexp("print.go", "fmt.Printf")))
+		env.AfterChange(
+			Diagnostics(
+				env.AtRegexp("print.go", "fmt.Printf"),
+				WithMessage("wrong type int"),
+			),
+		)
 	})
 }
 
@@ -444,7 +448,9 @@
 `
 	Run(t, adHoc, func(t *testing.T, env *Env) {
 		env.OpenFile("b/b.go")
-		env.Await(Diagnostics(env.AtRegexp("b/b.go", "x")))
+		env.AfterChange(
+			Diagnostics(env.AtRegexp("b/b.go", "x")),
+		)
 	})
 }
 
@@ -542,13 +548,13 @@
 `
 	Run(t, noModule, func(t *testing.T, env *Env) {
 		env.OpenFile("a.go")
-		env.Await(
+		env.AfterChange(
 			// Expect the adHocPackagesWarning.
 			OutstandingWork(lsp.WorkspaceLoadFailure, "outside of a module"),
 		)
 		// Deleting the import dismisses the warning.
 		env.RegexpReplace("a.go", `import "mod.com/hello"`, "")
-		env.Await(
+		env.AfterChange(
 			NoOutstandingWork(),
 		)
 	})
@@ -564,7 +570,8 @@
 			WithOptions(
 				EnvVars{"GO111MODULE": go111module},
 			).Run(t, files, func(t *testing.T, env *Env) {
-				env.Await(
+				env.OnceMet(
+					InitialWorkspaceLoad,
 					NoOutstandingWork(),
 				)
 			})
@@ -836,7 +843,7 @@
 	var x int
 }
 `)
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("foo/bar_test.go", "x")),
 		)
 	})
@@ -918,7 +925,9 @@
 
 		// We should still get diagnostics for files that exist.
 		env.RegexpReplace("b/b.go", `a.A`, "a.Nonexistant")
-		env.Await(Diagnostics(env.AtRegexp("b/b.go", `Nonexistant`)))
+		env.AfterChange(
+			Diagnostics(env.AtRegexp("b/b.go", `Nonexistant`)),
+		)
 	})
 }
 
@@ -964,7 +973,9 @@
 		Run(t, mod, func(t *testing.T, env *Env) {
 			writeGoVim(env, "p/p.go", p)
 			writeGoVim(env, "main.go", main)
-			env.Await(Diagnostics(env.AtRegexp("main.go", "5")))
+			env.AfterChange(
+				Diagnostics(env.AtRegexp("main.go", "5")),
+			)
 		})
 	})
 
@@ -993,7 +1004,7 @@
 	p.DoIt(5)
 }
 `)
-			env.Await(
+			env.AfterChange(
 				Diagnostics(env.AtRegexp("main.go", "5")),
 				Diagnostics(env.AtRegexp("p/p_test.go", "5")),
 				Diagnostics(env.AtRegexp("p/x_test.go", "5")),
@@ -1026,7 +1037,7 @@
 		WorkspaceFolders(),
 	).Run(t, mod, func(t *testing.T, env *Env) {
 		env.OpenFile("a/a.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("a/a.go", "x")),
 		)
 	})
@@ -1074,9 +1085,7 @@
 `
 	Run(t, basic, func(t *testing.T, env *Env) {
 		env.Editor.CreateBuffer(env.Ctx, "foo.go", `package main`)
-		env.Await(
-			env.DoneWithOpen(),
-		)
+		env.AfterChange()
 		env.CloseBuffer("foo.go")
 		env.AfterChange(NoLogMatching(protocol.Info, "packages=0"))
 	})
@@ -1118,7 +1127,7 @@
 	var x int
 }
 `))
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("main.go", "x")),
 		)
 	})
@@ -1138,8 +1147,11 @@
 `
 	Run(t, pkgDefault, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
-		env.Await(
-			Diagnostics(env.AtRegexp("main.go", "default"), WithMessage("expected 'IDENT'")),
+		env.AfterChange(
+			Diagnostics(
+				env.AtRegexp("main.go", "default"),
+				WithMessage("expected 'IDENT'"),
+			),
 		)
 	})
 }
@@ -1284,11 +1296,11 @@
 `
 	Run(t, files, func(t *testing.T, env *Env) {
 		env.OpenFile("a/a.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("a/a.go", "x")),
 		)
 		env.OpenFile("a/a_exclude.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("a/a_exclude.go", "package (a)")),
 		)
 	})
@@ -1453,7 +1465,7 @@
 `
 	Run(t, pkg, func(t *testing.T, env *Env) {
 		env.OpenFile("go.mod")
-		env.Await(
+		env.AfterChange(
 			OutstandingWork(lsp.WorkspaceLoadFailure, "unknown directive"),
 		)
 		env.EditBuffer("go.mod", fake.NewEdit(0, 0, 3, 0, `module mod.com
@@ -1463,12 +1475,12 @@
 		// As of golang/go#42529, go.mod changes do not reload the workspace until
 		// they are saved.
 		env.SaveBufferWithoutActions("go.mod")
-		env.Await(
+		env.AfterChange(
 			OutstandingWork(lsp.WorkspaceLoadFailure, "invalid go version"),
 		)
 		env.RegexpReplace("go.mod", "go 1.hello", "go 1.12")
 		env.SaveBufferWithoutActions("go.mod")
-		env.Await(
+		env.AfterChange(
 			NoOutstandingWork(),
 		)
 	})
@@ -1541,7 +1553,8 @@
 import _ "mod.com/triple/a"
 `
 	Run(t, mod, func(t *testing.T, env *Env) {
-		env.Await(
+		env.OnceMet(
+			InitialWorkspaceLoad,
 			Diagnostics(env.AtRegexp("self/self.go", `_ "mod.com/self"`), WithMessage("import cycle not allowed")),
 			Diagnostics(env.AtRegexp("double/a/a.go", `_ "mod.com/double/b"`), WithMessage("import cycle not allowed")),
 			Diagnostics(env.AtRegexp("triple/a/a.go", `_ "mod.com/triple/b"`), WithMessage("import cycle not allowed")),
@@ -1587,7 +1600,7 @@
 		)
 		env.RegexpReplace("b/b.go", `const B = a\.B`, "")
 		env.SaveBuffer("b/b.go")
-		env.Await(
+		env.AfterChange(
 			NoDiagnostics(ForFile("a/a.go")),
 			NoDiagnostics(ForFile("b/b.go")),
 		)
@@ -1609,7 +1622,8 @@
 `
 	t.Run("module", func(t *testing.T) {
 		Run(t, mod, func(t *testing.T, env *Env) {
-			env.Await(
+			env.OnceMet(
+				InitialWorkspaceLoad,
 				Diagnostics(env.AtRegexp("main.go", `"nosuchpkg"`), WithMessage(`could not import nosuchpkg (no required module provides package "nosuchpkg"`)),
 			)
 		})
@@ -1620,7 +1634,8 @@
 			EnvVars{"GO111MODULE": "off"},
 			Modes(Default),
 		).Run(t, mod, func(t *testing.T, env *Env) {
-			env.Await(
+			env.OnceMet(
+				InitialWorkspaceLoad,
 				Diagnostics(env.AtRegexp("main.go", `"nosuchpkg"`), WithMessage(`cannot find package "nosuchpkg" in any of`)),
 			)
 		})
@@ -1749,7 +1764,7 @@
 	Run(t, files, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
 		env.OpenFile("other.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("main.go", "asdf")),
 			Diagnostics(env.AtRegexp("main.go", "fdas")),
 		)
diff --git a/gopls/internal/regtest/diagnostics/undeclared_test.go b/gopls/internal/regtest/diagnostics/undeclared_test.go
index a6289c9..ac5f598 100644
--- a/gopls/internal/regtest/diagnostics/undeclared_test.go
+++ b/gopls/internal/regtest/diagnostics/undeclared_test.go
@@ -49,7 +49,6 @@
 			Diagnostics(env.AtRegexp("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)
 		}
diff --git a/gopls/internal/regtest/misc/failures_test.go b/gopls/internal/regtest/misc/failures_test.go
index aad8360..9ab456d 100644
--- a/gopls/internal/regtest/misc/failures_test.go
+++ b/gopls/internal/regtest/misc/failures_test.go
@@ -69,14 +69,14 @@
 
 	Run(t, badPackageDup, func(t *testing.T, env *Env) {
 		env.OpenFile("b.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("b.go", `a = 2`), WithMessage("a redeclared")),
 			Diagnostics(env.AtRegexp("a.go", `a = 1`), WithMessage("other declaration")),
 		)
 
 		// Fix the error by editing the const name in b.go to `b`.
 		env.RegexpReplace("b.go", "(a) = 2", "b")
-		env.Await(
+		env.AfterChange(
 			NoDiagnostics(ForFile("a.go")),
 			NoDiagnostics(ForFile("b.go")),
 		)
diff --git a/gopls/internal/regtest/misc/generate_test.go b/gopls/internal/regtest/misc/generate_test.go
index 69a28e2..85dd9a7 100644
--- a/gopls/internal/regtest/misc/generate_test.go
+++ b/gopls/internal/regtest/misc/generate_test.go
@@ -59,7 +59,8 @@
 `
 
 	Run(t, generatedWorkspace, func(t *testing.T, env *Env) {
-		env.Await(
+		env.OnceMet(
+			InitialWorkspaceLoad,
 			Diagnostics(env.AtRegexp("main.go", "lib1.(Answer)")),
 		)
 		env.RunGenerate("./lib1")
diff --git a/gopls/internal/regtest/misc/shared_test.go b/gopls/internal/regtest/misc/shared_test.go
index de44167..eaa4c7e 100644
--- a/gopls/internal/regtest/misc/shared_test.go
+++ b/gopls/internal/regtest/misc/shared_test.go
@@ -54,8 +54,8 @@
 		env2.RegexpReplace("main.go", "\\)\n(})", "")
 
 		// Now check that we got different diagnostics in each environment.
-		env1.Await(Diagnostics(env1.AtRegexp("main.go", "Printl")))
-		env2.Await(Diagnostics(env2.AtRegexp("main.go", "$")))
+		env1.AfterChange(Diagnostics(env1.AtRegexp("main.go", "Printl")))
+		env2.AfterChange(Diagnostics(env2.AtRegexp("main.go", "$")))
 
 		// Now close editor #2, and verify that operation in editor #1 is
 		// unaffected.
diff --git a/gopls/internal/regtest/modfile/modfile_test.go b/gopls/internal/regtest/modfile/modfile_test.go
index d3df215..f44eefa 100644
--- a/gopls/internal/regtest/modfile/modfile_test.go
+++ b/gopls/internal/regtest/modfile/modfile_test.go
@@ -101,7 +101,7 @@
 			// Save the buffer, which will format and organize imports.
 			// Confirm that the go.mod file still does not change.
 			env.SaveBuffer("a/main.go")
-			env.Await(
+			env.AfterChange(
 				Diagnostics(env.AtRegexp("a/main.go", "\"example.com/blah\"")),
 			)
 			if got := env.ReadWorkspaceFile("a/go.mod"); got != goModContent {
@@ -124,22 +124,14 @@
 			//
 			// If this proves insufficient, env.RemoveWorkspaceFile can be updated to
 			// retry file lock errors on windows.
-			env.Await(
-				env.DoneWithOpen(),
-				env.DoneWithSave(),
-				env.DoneWithChangeWatchedFiles(),
-			)
+			env.AfterChange()
 			env.RemoveWorkspaceFile("a/main.go")
 
 			// TODO(rfindley): awaiting here shouldn't really be necessary. We should
 			// be consistent eventually.
 			//
 			// Probably this was meant to exercise a race with the change below.
-			env.Await(
-				env.DoneWithOpen(),
-				env.DoneWithSave(),
-				env.DoneWithChangeWatchedFiles(),
-			)
+			env.AfterChange()
 
 			env.WriteWorkspaceFile("a/main.go", mainContent)
 			env.AfterChange(
@@ -595,7 +587,7 @@
 	t.Run("bad", func(t *testing.T) {
 		runner.Run(t, unknown, func(t *testing.T, env *Env) {
 			env.OpenFile("a/go.mod")
-			env.Await(
+			env.AfterChange(
 				Diagnostics(env.AtRegexp("a/go.mod", "example.com v1.2.2")),
 			)
 			env.RegexpReplace("a/go.mod", "v1.2.2", "v1.2.3")
@@ -705,7 +697,7 @@
 		{"nested", WithOptions(ProxyFiles(badProxy))},
 	}.Run(t, module, func(t *testing.T, env *Env) {
 		env.OpenFile("a/go.mod")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("a/go.mod", "require example.com v1.2.3")),
 		)
 	})
@@ -734,7 +726,7 @@
 	).Run(t, mod, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
 		original := env.ReadWorkspaceFile("go.mod")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("main.go", `"example.com/blah"`)),
 		)
 		got := env.ReadWorkspaceFile("go.mod")
@@ -792,8 +784,11 @@
 	WithOptions(
 		ProxyFiles(workspaceProxy),
 	).Run(t, mod, func(t *testing.T, env *Env) {
-		env.Await(
-			Diagnostics(env.AtRegexp("a/go.mod", "example.com v1.2.3"), WithMessage("is not used")),
+		env.AfterChange(
+			Diagnostics(
+				env.AtRegexp("a/go.mod", "example.com v1.2.3"),
+				WithMessage("is not used"),
+			),
 		)
 	})
 }
@@ -819,7 +814,8 @@
 		ProxyFiles(workspaceProxy),
 		Settings{"buildFlags": []string{"-tags", "bob"}},
 	).Run(t, mod, func(t *testing.T, env *Env) {
-		env.Await(
+		env.OnceMet(
+			InitialWorkspaceLoad,
 			Diagnostics(env.AtRegexp("main.go", `"example.com/blah"`)),
 		)
 	})
@@ -839,7 +835,7 @@
 	Run(t, mod, func(t *testing.T, env *Env) {
 		env.OpenFile("go.mod")
 		env.RegexpReplace("go.mod", "module", "modul")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("go.mod", "modul")),
 		)
 	})
@@ -1143,7 +1139,7 @@
 	).Run(t, mod, func(t *testing.T, env *Env) {
 		env.OpenFile("main.go")
 		d := &protocol.PublishDiagnosticsParams{}
-		env.Await(
+		env.AfterChange(
 			Diagnostics(
 				env.AtRegexp("main.go", `"example.com/blah"`),
 				WithMessage(`could not import example.com/blah (no required module provides package "example.com/blah")`),
@@ -1151,7 +1147,7 @@
 			ReadDiagnostics("main.go", d),
 		)
 		env.ApplyQuickFixes("main.go", d.Diagnostics)
-		env.Await(
+		env.AfterChange(
 			NoDiagnostics(ForFile("main.go")),
 			NoDiagnostics(ForFile("go.mod")),
 		)
diff --git a/gopls/internal/regtest/watch/watch_test.go b/gopls/internal/regtest/watch/watch_test.go
index 9086483..edb479a 100644
--- a/gopls/internal/regtest/watch/watch_test.go
+++ b/gopls/internal/regtest/watch/watch_test.go
@@ -89,9 +89,9 @@
 `
 	Run(t, pkg, func(t *testing.T, env *Env) {
 		env.OpenFile("a/a.go")
-		env.Await(env.DoneWithOpen())
+		env.AfterChange()
 		env.WriteWorkspaceFile("b/b.go", `package b; func B() {};`)
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("a/a.go", "b.B")),
 		)
 	})
@@ -167,9 +167,9 @@
 `
 	Run(t, pkg, func(t *testing.T, env *Env) {
 		env.OpenFile("a/a.go")
-		env.Await(env.DoneWithOpen())
+		env.AfterChange()
 		env.RemoveWorkspaceFile("b/b.go")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("a/a.go", "\"mod.com/b\"")),
 		)
 	})
@@ -419,7 +419,7 @@
 			env.RemoveWorkspaceFile("a/a_unneeded.go")
 			env.CloseBuffer("a/a_unneeded.go")
 			env.RegexpReplace("a/a.go", "var _ int", "fmt.Println(\"\")")
-			env.Await(
+			env.AfterChange(
 				Diagnostics(env.AtRegexp("a/a.go", "fmt")),
 			)
 			env.SaveBuffer("a/a.go")
diff --git a/gopls/internal/regtest/workspace/workspace_test.go b/gopls/internal/regtest/workspace/workspace_test.go
index 0573455..3699fa8 100644
--- a/gopls/internal/regtest/workspace/workspace_test.go
+++ b/gopls/internal/regtest/workspace/workspace_test.go
@@ -199,8 +199,7 @@
 replace random.org => %s
 `, env.ReadWorkspaceFile("pkg/go.mod"), dir)
 		env.WriteWorkspaceFile("pkg/go.mod", goModWithReplace)
-		env.Await(
-			env.DoneWithChangeWatchedFiles(),
+		env.AfterChange(
 			UnregistrationMatching("didChangeWatchedFiles"),
 			RegistrationMatching("didChangeWatchedFiles"),
 		)
@@ -323,7 +322,8 @@
 	WithOptions(
 		ProxyFiles(proxy),
 	).Run(t, multiModule, func(t *testing.T, env *Env) {
-		env.Await(
+		env.OnceMet(
+			InitialWorkspaceLoad,
 			Diagnostics(env.AtRegexp("main.go", "x")),
 		)
 	})
@@ -507,7 +507,7 @@
 		)
 		env.RegexpReplace("modb/go.mod", "modul", "module")
 		env.SaveBufferWithoutActions("modb/go.mod")
-		env.Await(
+		env.AfterChange(
 			Diagnostics(env.AtRegexp("modb/b/b.go", "x")),
 		)
 	})