Add support for xtest packages in Go Packages Driver

go/tools/bazel_testing/bazel_testing.go

@@ -206,15 +206,16 @@ func RunBazel(args ...string) error {
 // If the command starts but exits with a non-zero status, a *StderrExitError
 // will be returned which wraps the original *exec.ExitError.
 func BazelOutput(args ...string) ([]byte, error) {
-	return BazelOutputWithInput(nil, args...)
+	stdout, _, err := BazelOutputWithInput(nil, args...)
+	return stdout, err
 }
 
 // BazelOutputWithInput invokes a bazel command with a list of arguments and
 // an input stream and returns the content of stdout.
 //
 // If the command starts but exits with a non-zero status, a *StderrExitError
 // will be returned which wraps the original *exec.ExitError.
-func BazelOutputWithInput(stdin io.Reader, args ...string) ([]byte, error) {
+func BazelOutputWithInput(stdin io.Reader, args ...string) ([]byte, []byte, error) {
 	cmd := BazelCmd(args...)
 	stdout := &bytes.Buffer{}
 	stderr := &bytes.Buffer{}
@@ -228,7 +229,7 @@ func BazelOutputWithInput(stdin io.Reader, args ...string) ([]byte, error) {
 		eErr.Stderr = stderr.Bytes()
 		err = &StderrExitError{Err: eErr}
 	}
-	return stdout.Bytes(), err
+	return stdout.Bytes(), stderr.Bytes(), err
 }
 
 // StderrExitError wraps *exec.ExitError and prints the complete stderr output

go/tools/gopackagesdriver/flatpackage.go

@@ -110,6 +110,61 @@ func (fp *FlatPackage) FilterFilesForBuildTags() {
 	fp.CompiledGoFiles = filterSourceFilesForTags(fp.CompiledGoFiles)
 }
 
+func (fp *FlatPackage) filterTestSuffix(files []string) (err error, testFiles []string, xTestFiles, nonTestFiles []string) {
+	for _, filename := range files {
+		if strings.HasSuffix(filename, "_test.go") {
+			fset := token.NewFileSet()
+			f, err := parser.ParseFile(fset, filename, nil, parser.PackageClauseOnly)
+			if err != nil {
+				return err, nil, nil, nil
+			}
+			if f.Name.Name == fp.Name {
+				testFiles = append(testFiles, filename)
+			} else {
+				xTestFiles = append(xTestFiles, filename)
+			}
+		} else {
+			nonTestFiles = append(nonTestFiles, filename)
+		}
+	}
+	return
+}
+
+func (fp *FlatPackage) MoveTestFiles() *FlatPackage {
+	err, tgf, xtgf, gf := fp.filterTestSuffix(fp.GoFiles)
+
+	if err != nil {
+		return nil
+	}
+	fp.GoFiles = append(gf, tgf...)
+	fp.CompiledGoFiles = append(gf, tgf...)
+
+	if len(xtgf) == 0 {
+		return nil
+	}
+
+	newImports := make(map[string]string, len(fp.Imports))
+	for k, v := range fp.Imports {
+		newImports[k] = v
+	}
+
+	newImports[fp.PkgPath] = fp.ID
+
+	// Clone package, only xtgf files
+	return &FlatPackage{
+		ID:              fp.ID + "_xtest",
+		Name:            fp.Name + "_test",
+		PkgPath:         fp.PkgPath,
+		Imports:         newImports,
+		Errors:          fp.Errors,
+		GoFiles:         append([]string{}, xtgf...),
+		CompiledGoFiles: append([]string{}, xtgf...),
+		OtherFiles:      fp.OtherFiles,
+		ExportFile:      fp.ExportFile,
+		Standard:        fp.Standard,
+	}
+}
+
 func (fp *FlatPackage) IsStdlib() bool {
 	return fp.Standard
 }

go/tools/gopackagesdriver/packageregistry.go

@@ -71,6 +71,10 @@ func (pr *PackageRegistry) ResolveImports() error {
 		if err := pkg.ResolveImports(resolve); err != nil {
 			return err
 		}
+		testFp := pkg.MoveTestFiles()
+		if testFp != nil {
+			pr.packagesByID[testFp.ID] = testFp
+		}
 	}
 
 	return nil
@@ -108,6 +112,10 @@ func (pr *PackageRegistry) Match(labels []string) ([]string, []*FlatPackage) {
 			}
 		} else {
 			roots[label] = struct{}{}
+			// If an xtest package exists for this package add it to the roots
+			if _, ok := pr.packagesByID[label+"_xtest"]; ok {
+				roots[label+"_xtest"] = struct{}{}
+			}
 		}
 	}
 

tests/integration/gopackagesdriver/gopackagesdriver_test.go

@@ -19,7 +19,7 @@ func TestMain(m *testing.M) {
 	bazel_testing.TestMain(m, bazel_testing.Args{
 		Main: `
 -- BUILD.bazel --
-load("@io_bazel_rules_go//go:def.bzl", "go_library")
+load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "hello",
@@ -28,6 +28,15 @@ go_library(
     visibility = ["//visibility:public"],
 )
 
+go_test(
+	name = "hello_test",
+	srcs = [
+		"hello_test.go",
+		"hello_external_test.go",
+	],
+	embed = [":hello"],
+)
+
 -- hello.go --
 package hello
 
@@ -36,6 +45,20 @@ import "os"
 func main() {
 	fmt.Fprintln(os.Stderr, "Hello World!")
 }
+
+-- hello_test.go --
+package hello
+
+import "testing"
+
+func TestHelloInternal(t *testing.T) {}
+
+-- hello_external_test.go --
+package hello_test
+
+import "testing"
+
+func TestHelloExternal(t *testing.T) {}
 		`,
 	})
 }
@@ -46,7 +69,7 @@ const (
 
 func TestBaseFileLookup(t *testing.T) {
 	reader := strings.NewReader("{}")
-	out, err := bazel_testing.BazelOutputWithInput(reader, "run", "@io_bazel_rules_go//go/tools/gopackagesdriver", "--", "file=hello.go")
+	out, _, err := bazel_testing.BazelOutputWithInput(reader, "run", "@io_bazel_rules_go//go/tools/gopackagesdriver", "--", "file=hello.go")
 	if err != nil {
 		t.Errorf("Unexpected error: %w", err.Error())
 		return
@@ -124,3 +147,50 @@ func TestBaseFileLookup(t *testing.T) {
 		}
 	})
 }
+
+func TestExternalTests(t *testing.T) {
+	reader := strings.NewReader("{}")
+	out, stderr, err := bazel_testing.BazelOutputWithInput(reader, "run", "@io_bazel_rules_go//go/tools/gopackagesdriver", "--", "file=hello_external_test.go")
+	if err != nil {
+		t.Errorf("Unexpected error: %w\n=====\n%s\n=====", err.Error(), stderr)
+	}
+	var resp response
+	err = json.Unmarshal(out, &resp)
+	if err != nil {
+		t.Errorf("Failed to unmarshal packages driver response: %w\n%w", err.Error(), out)
+	}
+
+	if len(resp.Roots) != 2 {
+		t.Errorf("Expected exactly two roots for package: %+v", resp.Roots)
+	}
+
+	var testId, xTestId string
+	for _, id := range resp.Roots {
+		if strings.HasSuffix(id, "_xtest") {
+			xTestId = id
+		} else {
+			testId = id
+		}
+	}
+
+	for _, p := range resp.Packages {
+		if p.ID == xTestId {
+			assertSuffixesInList(t, p.GoFiles, "/hello_external_test.go")
+		} else if p.ID == testId {
+			assertSuffixesInList(t, p.GoFiles, "/hello.go", "/hello_test.go")
+		}
+	}
+}
+
+func assertSuffixesInList(t *testing.T, list []string, expectedSuffixes ...string) {
+	for _, suffix := range expectedSuffixes {
+		itemFound := false
+		for _, listItem := range list {
+			itemFound = itemFound || strings.HasSuffix(listItem, suffix)
+		}
+
+		if !itemFound {
+			t.Errorf("Expected suffix %q in list, but was not found: %+v", suffix, list)
+		}
+	}
+}