1 files changed, 16 insertions, 10 deletions
diff --git a/gopls/internal/regtest/diagnostics/undeclared_test.go b/gopls/internal/regtest/diagnostics/undeclared_test.go
index 79f7d4267..ac5f598cc 100644
--- a/gopls/internal/regtest/diagnostics/undeclared_test.go
+++ b/gopls/internal/regtest/diagnostics/undeclared_test.go
@@ -7,8 +7,8 @@ package diagnostics
 import (
 	"testing"
 
-	"golang.org/x/tools/internal/lsp/protocol"
-	. "golang.org/x/tools/internal/lsp/regtest"
+	"golang.org/x/tools/gopls/internal/lsp/protocol"
+	. "golang.org/x/tools/gopls/internal/lsp/regtest"
 )
 
 func TestUndeclaredDiagnostics(t *testing.T) {
@@ -44,23 +44,29 @@ func _() int {
 
 		// 'x' is undeclared, but still necessary.
 		env.OpenFile("a/a.go")
-		env.Await(env.DiagnosticAtRegexp("a/a.go", "x"))
-		diags := env.DiagnosticsFor("a/a.go")
-		if got := len(diags.Diagnostics); got != 1 {
+		var adiags protocol.PublishDiagnosticsParams
+		env.AfterChange(
+			Diagnostics(env.AtRegexp("a/a.go", "x")),
+			ReadDiagnostics("a/a.go", &adiags),
+		)
+		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.DiagnosticsFor("b/b.go")
-		if got := len(diags.Diagnostics); got != 1 {
+		var bdiags protocol.PublishDiagnosticsParams
+		env.AfterChange(
+			Diagnostics(env.AtRegexp("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)
 		}
 	})