internal/lsp, go/packages: fix Go 1.15-related overlay bug

Something about the ordering of `go list` results must have changed in
Go 1.15, so overlays with test variants were failing to get the correct
imports. Technically, the correct solution would have been to support
overlays for a package *and* its test variant, but for the purposes of
gopls, it's actually more important that the overlays apply to the
package itself (rather than its test variant).

Fixes #40685

Change-Id: I3282557502f7f30bf484e1e6c17b90db824bc7d0
Reviewed-on: https://go-review.googlesource.com/c/tools/+/252681
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
Reviewed-by: Michael Matloob <matloob@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>

go/packages/golist_overlay.go

@@ -89,9 +89,19 @@ func (state *golistState) processGolistOverlay(response *responseDeduper) (modif
 					// because the file is generated in another directory.
 					testVariantOf = p
 					continue nextPackage
+				} else if !isTestFile && hasTestFiles(p) {
+					// We're examining a test variant, but the overlaid file is
+					// a non-test file. Because the overlay implementation
+					// (currently) only adds a file to one package, skip this
+					// package, so that we can add the file to the production
+					// variant of the package. (https://golang.org/issue/36857
+					// tracks handling overlays on both the production and test
+					// variant of a package).
+					continue nextPackage
 				}
-				// We must have already seen the package of which this is a test variant.
 				if pkg != nil && p != pkg && pkg.PkgPath == p.PkgPath {
+					// We have already seen the production version of the
+					// for which p is a test variant.
 					if hasTestFiles(p) {
 						testVariantOf = pkg
 					}

go/packages/overlay_test.go

@@ -863,3 +863,72 @@ func testInvalidXTestInGOPATH(t *testing.T, exporter packagestest.Exporter) {
 		t.Fatalf("expected at least 2 CompiledGoFiles for %s, got %v", pkg.PkgPath, len(pkg.CompiledGoFiles))
 	}
 }
+
+// Reproduces golang/go#40685.
+func TestAddImportInOverlay(t *testing.T) {
+	packagestest.TestAll(t, testAddImportInOverlay)
+}
+func testAddImportInOverlay(t *testing.T, exporter packagestest.Exporter) {
+	exported := packagestest.Export(t, exporter, []packagestest.Module{
+		{
+			Name: "golang.org/fake",
+			Files: map[string]interface{}{
+				"a/a.go": `package a
+
+import (
+	"fmt"
+)
+
+func _() {
+	fmt.Println("")
+	os.Stat("")
+}`,
+				"a/a_test.go": `package a
+
+import (
+	"os"
+	"testing"
+)
+
+func TestA(t *testing.T) {
+	os.Stat("")
+}`,
+			},
+		},
+	})
+	defer exported.Cleanup()
+
+	exported.Config.Mode = everythingMode
+	exported.Config.Tests = true
+
+	dir := filepath.Dir(exported.File("golang.org/fake", "a/a.go"))
+	exported.Config.Overlay = map[string][]byte{
+		filepath.Join(dir, "a.go"): []byte(`package a
+
+import (
+	"fmt"
+	"os"
+)
+
+func _() {
+	fmt.Println("")
+	os.Stat("")
+}
+`),
+	}
+	initial, err := packages.Load(exported.Config, "golang.org/fake/a")
+	if err != nil {
+		t.Fatal(err)
+	}
+	pkg := initial[0]
+	var foundOs bool
+	for _, imp := range pkg.Imports {
+		if imp.PkgPath == "os" {
+			foundOs = true
+			break
+		}
+	}
+	if !foundOs {
+		t.Fatalf(`expected import "os", found none: %v`, pkg.Imports)
+	}
+}

gopls/internal/regtest/imports_test.go

@@ -7,10 +7,13 @@ import (
 	"strings"
 	"testing"
 
+	"golang.org/x/tools/internal/lsp"
 	"golang.org/x/tools/internal/lsp/fake"
+	"golang.org/x/tools/internal/lsp/protocol"
 	"golang.org/x/tools/internal/testenv"
 )
 
+// Tests golang/go#38815.
 func TestIssue38815(t *testing.T) {
 	const needs = `
 -- go.mod --
@@ -156,3 +159,52 @@ var _, _ = x.X, y.Y
 		}
 	})
 }
+
+// Tests golang/go#40685.
+func TestAcceptImportsQuickFixTestVariant(t *testing.T) {
+	const pkg = `
+-- go.mod --
+module mod.com
+
+go 1.12
+-- a/a.go --
+package a
+
+import (
+	"fmt"
+)
+
+func _() {
+	fmt.Println("")
+	os.Stat("")
+}
+-- a/a_test.go --
+package a
+
+import (
+	"os"
+	"testing"
+)
+
+func TestA(t *testing.T) {
+	os.Stat("")
+}
+`
+	run(t, pkg, func(t *testing.T, env *Env) {
+		env.Await(
+			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromInitialWorkspaceLoad), 1),
+		)
+		env.OpenFile("a/a.go")
+		metBy := env.Await(
+			env.DiagnosticAtRegexp("a/a.go", "os.Stat"),
+		)
+		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"),
+		)
+	})
+}