internal/imports: load test exports of package under test

x_tests can access exports from the test variant of the package under
test. Change loadExportsFromFiles to understand that mode, and use it
where appropriate.

I didn't want to come up with a cache key for for the test variant, so
for now we bypass the cache in these situations.

Fixes golang/go#29979.

Change-Id: I9959a08da97bbee64c5bcd56e06f548486693156
Reviewed-on: https://go-review.googlesource.com/c/tools/+/213221
Run-TryBot: Heschi Kreinick <heschi@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>

internal/imports/fix.go

@@ -629,6 +629,14 @@ func getCandidatePkgs(ctx context.Context, wrappedCallback *scanCallback, filena
 			return notSelf(pkg) && wrappedCallback.packageNameLoaded(pkg)
 		},
 		exportsLoaded: func(pkg *pkg, exports []string) {
+			// If we're an x_test, load the package under test's test variant.
+			if strings.HasSuffix(filePkg, "_test") && pkg.dir == filepath.Dir(filename) {
+				var err error
+				_, exports, err = loadExportsFromFiles(ctx, env, pkg.dir, true)
+				if err != nil {
+					return
+				}
+			}
 			wrappedCallback.exportsLoaded(pkg, exports)
 		},
 	}
@@ -842,7 +850,7 @@ type Resolver interface {
 	scan(ctx context.Context, callback *scanCallback) error
 	// loadExports returns the set of exported symbols in the package at dir.
 	// loadExports may be called concurrently.
-	loadExports(ctx context.Context, pkg *pkg) (string, []string, error)
+	loadExports(ctx context.Context, pkg *pkg, includeTest bool) (string, []string, error)
 
 	ClearForNewScan()
 }
@@ -1192,7 +1200,7 @@ func (r *gopathResolver) scan(ctx context.Context, callback *scanCallback) error
 		if !callback.packageNameLoaded(p) {
 			return
 		}
-		if _, exports, err := r.loadExports(ctx, p); err == nil {
+		if _, exports, err := r.loadExports(ctx, p, false); err == nil {
 			callback.exportsLoaded(p, exports)
 		}
 	}
@@ -1231,11 +1239,11 @@ func filterRoots(roots []gopathwalk.Root, include func(gopathwalk.Root) bool) []
 	return result
 }
 
-func (r *gopathResolver) loadExports(ctx context.Context, pkg *pkg) (string, []string, error) {
-	if info, ok := r.cache.Load(pkg.dir); ok {
+func (r *gopathResolver) loadExports(ctx context.Context, pkg *pkg, includeTest bool) (string, []string, error) {
+	if info, ok := r.cache.Load(pkg.dir); ok && !includeTest {
 		return r.cache.CacheExports(ctx, r.env, info)
 	}
-	return loadExportsFromFiles(ctx, r.env, pkg.dir)
+	return loadExportsFromFiles(ctx, r.env, pkg.dir, includeTest)
 }
 
 // VendorlessPath returns the devendorized version of the import path ipath.
@@ -1251,7 +1259,7 @@ func VendorlessPath(ipath string) string {
 	return ipath
 }
 
-func loadExportsFromFiles(ctx context.Context, env *ProcessEnv, dir string) (string, []string, error) {
+func loadExportsFromFiles(ctx context.Context, env *ProcessEnv, dir string, includeTest bool) (string, []string, error) {
 	var exports []string
 
 	// Look for non-test, buildable .go files which could provide exports.
@@ -1262,7 +1270,7 @@ func loadExportsFromFiles(ctx context.Context, env *ProcessEnv, dir string) (str
 	var files []os.FileInfo
 	for _, fi := range all {
 		name := fi.Name()
-		if !strings.HasSuffix(name, ".go") || strings.HasSuffix(name, "_test.go") {
+		if !strings.HasSuffix(name, ".go") || (!includeTest && strings.HasSuffix(name, "_test.go")) {
 			continue
 		}
 		match, err := env.buildContext().MatchFile(dir, fi.Name())
@@ -1295,6 +1303,10 @@ func loadExportsFromFiles(ctx context.Context, env *ProcessEnv, dir string) (str
 			// handled by MatchFile above.
 			continue
 		}
+		if includeTest && strings.HasSuffix(f.Name.Name, "_test") {
+			// x_test package. We want internal test files only.
+			continue
+		}
 		pkgName = f.Name.Name
 		for name := range f.Scope.Objects {
 			if ast.IsExported(name) {
@@ -1360,7 +1372,9 @@ func findImport(ctx context.Context, pass *pass, candidates []pkgDistance, pkgNa
 				if pass.env.Debug {
 					pass.env.Logf("loading exports in dir %s (seeking package %s)", c.pkg.dir, pkgName)
 				}
-				exports, err := loadExportsForPackage(ctx, pass.env, pkgName, c.pkg)
+				// If we're an x_test, load the package under test's test variant.
+				includeTest := strings.HasSuffix(pass.f.Name.Name, "_test") && c.pkg.dir == pass.srcDir
+				_, exports, err := pass.env.GetResolver().loadExports(ctx, c.pkg, includeTest)
 				if err != nil {
 					if pass.env.Debug {
 						pass.env.Logf("loading exports in dir %s (seeking package %s): %v", c.pkg.dir, pkgName, err)
@@ -1397,17 +1411,6 @@ func findImport(ctx context.Context, pass *pass, candidates []pkgDistance, pkgNa
 	return nil, nil
 }
 
-func loadExportsForPackage(ctx context.Context, env *ProcessEnv, expectPkg string, pkg *pkg) ([]string, error) {
-	pkgName, exports, err := env.GetResolver().loadExports(ctx, pkg)
-	if err != nil {
-		return nil, err
-	}
-	if expectPkg != pkgName {
-		return nil, fmt.Errorf("dir %v is package %v, wanted %v", pkg.dir, pkgName, expectPkg)
-	}
-	return exports, err
-}
-
 // pkgIsCandidate reports whether pkg is a candidate for satisfying the
 // finding which package pkgIdent in the file named by filename is trying
 // to refer to.

internal/imports/fix_test.go

@@ -2489,6 +2489,40 @@ var _ = bytes.Buffer{}
 	}.processTest(t, "foo.com", "foo.go", nil, nil, want)
 }
 
+// Tests that an external test package will import the package under test if it
+// also uses symbols exported only in test files.
+// https://golang.org/issues/29979
+func TestExternalTest(t *testing.T) {
+	const input = `package a_test
+func TestX() {
+	a.X()
+	a.Y()
+}
+`
+	const want = `package a_test
+
+import "foo.com/a"
+
+func TestX() {
+	a.X()
+	a.Y()
+}
+`
+
+	testConfig{
+		modules: []packagestest.Module{
+			{
+				Name: "foo.com/a",
+				Files: fm{
+					"a.go":           "package a\n func X() {}",
+					"export_test.go": "package a\n func Y() {}",
+					"a_test.go":      input,
+				},
+			},
+		},
+	}.processTest(t, "foo.com/a", "a_test.go", nil, nil, want)
+}
+
 // TestStdLibGetCandidates tests that get packages finds std library packages
 // with correct priorities.
 func TestGetCandidates(t *testing.T) {

internal/imports/mod.go

@@ -413,7 +413,7 @@ func (r *ModuleResolver) scan(ctx context.Context, callback *scanCallback) error
 		if !callback.packageNameLoaded(pkg) {
 			return
 		}
-		_, exports, err := r.loadExports(ctx, pkg)
+		_, exports, err := r.loadExports(ctx, pkg, false)
 		if err != nil {
 			return
 		}
@@ -529,14 +529,14 @@ func (r *ModuleResolver) canonicalize(info directoryPackageInfo) (*pkg, error) {
 	return res, nil
 }
 
-func (r *ModuleResolver) loadExports(ctx context.Context, pkg *pkg) (string, []string, error) {
+func (r *ModuleResolver) loadExports(ctx context.Context, pkg *pkg, includeTest bool) (string, []string, error) {
 	if err := r.init(); err != nil {
 		return "", nil, err
 	}
-	if info, ok := r.cacheLoad(pkg.dir); ok {
+	if info, ok := r.cacheLoad(pkg.dir); ok && !includeTest {
 		return r.cacheExports(ctx, r.env, info)
 	}
-	return loadExportsFromFiles(ctx, r.env, pkg.dir)
+	return loadExportsFromFiles(ctx, r.env, pkg.dir, includeTest)
 }
 
 func (r *ModuleResolver) scanDirForPackage(root gopathwalk.Root, dir string) directoryPackageInfo {

internal/imports/mod_cache.go

@@ -215,7 +215,7 @@ func (d *dirInfoCache) CacheExports(ctx context.Context, env *ProcessEnv, info d
 	if reached, err := info.reachedStatus(nameLoaded); reached && err != nil {
 		return "", nil, err
 	}
-	info.packageName, info.exports, info.err = loadExportsFromFiles(ctx, env, info.dir)
+	info.packageName, info.exports, info.err = loadExportsFromFiles(ctx, env, info.dir, false)
 	if info.err == context.Canceled || info.err == context.DeadlineExceeded {
 		return info.packageName, info.exports, info.err
 	}

internal/lsp/testdata/summary.txt.golden

@@ -1,7 +1,7 @@
 -- summary --
 CompletionsCount = 225
-CompletionSnippetCount = 62
-UnimportedCompletionsCount = 4
+CompletionSnippetCount = 63
+UnimportedCompletionsCount = 9
 DeepCompletionsCount = 5
 FuzzyCompletionsCount = 8
 RankedCompletionsCount = 56

internal/lsp/testdata/unimported/export_test.go

@@ -0,0 +1,3 @@
+package unimported
+
+var TestExport int //@item(testexport, "TestExport", "(from \"golang.org/x/tools/internal/lsp/unimported\")", "var")

internal/lsp/testdata/unimported/x_test.go

@@ -0,0 +1,9 @@
+package unimported_test
+
+import (
+	"testing"
+)
+
+func TestSomething(t *testing.T) {
+	_ = unimported.TestExport //@unimported("TestExport", testexport)
+}