gopackagesdriver: move and simplify test (#3856)

* gopackagesdriver: move and simplify test

Moved //tests/integration/gopackagesdriver:gopackagesdriver_test
to //go/tools/gopackagesdriver:gopackagesdriver_test.

Previously the test invoked gopackagesdriver with 'bazel run'
in a temporary workspace, which forced Bazel to build it, slowing
down the test significantly, especially on Windows. Now, the test
embeds the gopackagesdriver library and invokes it directly.
The driver still runs in a temporary workspace and still needs to
invoke bazel, so the test is still slow, but it's faster than before.

Also, a number of other improvements:

- gopackagesdriver now sets --consistent_labels if the bazel version
  is 6.4.0 or later instead of guessing based on the link-stamped
  rules_go repo name.
- bazel_testing.BazelCmd now writes --output_user_root to a .bazelrc
  file instead of passing it to commands. gopackagesdriver invokes
  bazel directly, so this lets it use the same output root.

* restore output directory path

* fix review comment

10 files changed, 167 insertions, 109 deletions

diff --git a/go/tools/bazel_testing/bazel_testing.go b/go/tools/bazel_testing/bazel_testing.go
index 7dbde4ac..78820c8a 100644
--- a/go/tools/bazel_testing/bazel_testing.go
+++ b/go/tools/bazel_testing/bazel_testing.go
@@ -135,7 +135,7 @@ func TestMain(m *testing.M, args Args) {
 	workspaceDir, cleanup, err := setupWorkspace(args, files)
 	defer func() {
 		if err := cleanup(); err != nil {
-			fmt.Fprintf(os.Stderr, "cleanup error: %v\n", err)
+			fmt.Fprintf(os.Stderr, "cleanup warning: %v\n", err)
 			// Don't fail the test on a cleanup error.
 			// Some operating systems (windows, maybe also darwin) can't reliably
 			// delete executable files after they're run.
@@ -175,13 +175,7 @@ func TestMain(m *testing.M, args Args) {
 // hide that this code is executing inside a bazel test.
 func BazelCmd(args ...string) *exec.Cmd {
 	cmd := exec.Command("bazel")
-	if outputUserRoot != "" {
-		cmd.Args = append(cmd.Args,
-			"--output_user_root="+outputUserRoot,
-			"--nosystem_rc",
-			"--nohome_rc",
-		)
-	}
+	cmd.Args = append(cmd.Args, "--nosystem_rc", "--nohome_rc")
 	cmd.Args = append(cmd.Args, args...)
 	for _, e := range os.Environ() {
 		// Filter environment variables set by the bazel test wrapper script.
@@ -291,7 +285,11 @@ func setupWorkspace(args Args, files []string) (dir string, cleanup func() error
 		tmpDir = filepath.Clean(tmpDir)
 		if i := strings.Index(tmpDir, string(os.PathSeparator)+"execroot"+string(os.PathSeparator)); i >= 0 {
 			outBaseDir = tmpDir[:i]
-			outputUserRoot = filepath.Dir(outBaseDir)
+			if dir, err := filepath.Abs(filepath.Dir(outBaseDir)); err == nil {
+				// Use forward slashes, even on Windows. Bazel's rc file parser
+				// reports an error if there are backslashes.
+				outputUserRoot = strings.ReplaceAll(dir, `\`, `/`)
+			}
 			cacheDir = filepath.Join(outBaseDir, "bazel_testing")
 		} else {
 			cacheDir = filepath.Join(tmpDir, "bazel_testing")
@@ -318,14 +316,18 @@ func setupWorkspace(args Args, files []string) (dir string, cleanup func() error
 		return "", cleanup, err
 	}
 
-	// Create a .bazelrc file if GO_BAZEL_TEST_BAZELFLAGS is set.
+	// Create a .bazelrc file with the contents of GO_BAZEL_TEST_BAZELFLAGS is set.
 	// The test can override this with its own .bazelrc or with flags in commands.
+	bazelrcPath := filepath.Join(mainDir, ".bazelrc")
+	bazelrcBuf := &bytes.Buffer{}
+	if outputUserRoot != "" {
+		fmt.Fprintf(bazelrcBuf, "startup --output_user_root=%s\n", outputUserRoot)
+	}
 	if flags := os.Getenv("GO_BAZEL_TEST_BAZELFLAGS"); flags != "" {
-		bazelrcPath := filepath.Join(mainDir, ".bazelrc")
-		content := "build " + flags
-		if err := ioutil.WriteFile(bazelrcPath, []byte(content), 0666); err != nil {
-			return "", cleanup, err
-		}
+		fmt.Fprintf(bazelrcBuf, "build %s\n", flags)
+	}
+	if err := os.WriteFile(bazelrcPath, bazelrcBuf.Bytes(), 0666); err != nil {
+		return "", cleanup, err
 	}
 
 	// Extract test files for the main workspace.
diff --git a/go/tools/gopackagesdriver/BUILD.bazel b/go/tools/gopackagesdriver/BUILD.bazel
index 86ad484d..2ab967f5 100644
--- a/go/tools/gopackagesdriver/BUILD.bazel
+++ b/go/tools/gopackagesdriver/BUILD.bazel
@@ -1,5 +1,6 @@
 load("//go:def.bzl", "go_binary", "go_library")
 load("//go/private:common.bzl", "RULES_GO_REPO_NAME")
+load("//go/tools/bazel_testing:def.bzl", "go_bazel_test")
 
 go_library(
     name = "gopackagesdriver_lib",
@@ -31,6 +32,19 @@ go_binary(
     },
 )
 
+RULES_GO_REPO_NAME_FOR_TEST = RULES_GO_REPO_NAME if RULES_GO_REPO_NAME != "@" else "@io_bazel_rules_go"
+
+go_bazel_test(
+    name = "gopackagesdriver_test",
+    size = "enormous",
+    srcs = ["gopackagesdriver_test.go"],
+    embed = [":gopackagesdriver_lib"],
+    rule_files = ["//:all_files"],
+    x_defs = {
+        "rulesGoRepositoryName": RULES_GO_REPO_NAME_FOR_TEST,
+    },
+)
+
 filegroup(
     name = "all_files",
     testonly = True,
diff --git a/go/tools/gopackagesdriver/bazel.go b/go/tools/gopackagesdriver/bazel.go
index 979faa6d..ce7ceb03 100644
--- a/go/tools/gopackagesdriver/bazel.go
+++ b/go/tools/gopackagesdriver/bazel.go
@@ -26,6 +26,7 @@ import (
 	"os"
 	"os/exec"
 	"path/filepath"
+	"strconv"
 	"strings"
 )
 
@@ -38,7 +39,7 @@ type Bazel struct {
 	workspaceRoot     string
 	bazelStartupFlags []string
 	info              map[string]string
-	version           string
+	version           bazelVersion
 }
 
 // Minimal BEP structs to access the build outputs
@@ -56,7 +57,7 @@ func NewBazel(ctx context.Context, bazelBin, workspaceRoot string, bazelStartupF
 		bazelBin:          bazelBin,
 		workspaceRoot:     workspaceRoot,
 		bazelStartupFlags: bazelStartupFlags,
-		version:           "6",
+		version:           bazelVersion{6, 0, 0}, // assumed until 'bazel info' output parsed
 	}
 	if err := b.fillInfo(ctx); err != nil {
 		return nil, fmt.Errorf("unable to query bazel info: %w", err)
@@ -77,7 +78,9 @@ func (b *Bazel) fillInfo(ctx context.Context) error {
 	}
 	release := strings.Split(b.info["release"], " ")
 	if len(release) == 2 {
-		b.version = release[1]
+		if version, ok := parseBazelVersion(release[1]); ok {
+			b.version = version
+		}
 	}
 	return nil
 }
@@ -168,3 +171,30 @@ func (b *Bazel) ExecutionRoot() string {
 func (b *Bazel) OutputBase() string {
 	return b.info["output_base"]
 }
+
+type bazelVersion [3]int
+
+func parseBazelVersion(raw string) (bazelVersion, bool) {
+	parts := strings.Split(raw, ".")
+	if len(parts) != 3 {
+		return [3]int{}, false
+	}
+	var version [3]int
+	for i, part := range parts {
+		v, err := strconv.Atoi(part)
+		if err != nil {
+			return [3]int{}, false
+		}
+		version[i] = v
+	}
+	return version, true
+}
+
+func (a bazelVersion) compare(b bazelVersion) int {
+	for i := 0; i < len(a); i++ {
+		if c := a[i] - b[i]; c != 0 {
+			return c
+		}
+	}
+	return 0
+}
diff --git a/go/tools/gopackagesdriver/bazel_json_builder.go b/go/tools/gopackagesdriver/bazel_json_builder.go
index c19f0308..1a5f100c 100644
--- a/go/tools/gopackagesdriver/bazel_json_builder.go
+++ b/go/tools/gopackagesdriver/bazel_json_builder.go
@@ -161,8 +161,7 @@ func (b *BazelJSONBuilder) outputGroupsForMode(mode LoadMode) string {
 
 func (b *BazelJSONBuilder) query(ctx context.Context, query string) ([]string, error) {
 	var bzlmodQueryFlags []string
-	if strings.HasPrefix(rulesGoRepositoryName, "@@") {
-		// Requires Bazel 6.4.0.
+	if b.bazel.version.compare(bazelVersion{6, 4, 0}) >= 0 {
 		bzlmodQueryFlags = []string{"--consistent_labels"}
 	}
 	queryArgs := concatStringsArrays(bazelQueryFlags, bzlmodQueryFlags, []string{
diff --git a/tests/integration/gopackagesdriver/gopackagesdriver_test.go b/go/tools/gopackagesdriver/gopackagesdriver_test.go
index b17a3325..3d61ab96 100644
--- a/tests/integration/gopackagesdriver/gopackagesdriver_test.go
+++ b/go/tools/gopackagesdriver/gopackagesdriver_test.go
@@ -1,18 +1,20 @@
-package gopackagesdriver_test
+package main
 
 import (
+	"bytes"
+	"context"
 	"encoding/json"
+	"os"
 	"path"
 	"strings"
 	"testing"
 
 	"github.com/bazelbuild/rules_go/go/tools/bazel_testing"
-	gpd "github.com/bazelbuild/rules_go/go/tools/gopackagesdriver"
 )
 
 type response struct {
 	Roots    []string `json:",omitempty"`
-	Packages []*gpd.FlatPackage
+	Packages []*FlatPackage
 }
 
 func TestMain(m *testing.M) {
@@ -69,18 +71,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")
-	if err != nil {
-		t.Errorf("Unexpected error: %w", err.Error())
-		return
-	}
-	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)
-		return
-	}
+	resp := runForTest(t, "file=hello.go")
 
 	t.Run("roots", func(t *testing.T) {
 		if len(resp.Roots) != 1 {
@@ -95,7 +86,7 @@ func TestBaseFileLookup(t *testing.T) {
 	})
 
 	t.Run("package", func(t *testing.T) {
-		var pkg *gpd.FlatPackage
+		var pkg *FlatPackage
 		for _, p := range resp.Packages {
 			if p.ID == resp.Roots[0] {
 				pkg = p
@@ -130,7 +121,7 @@ func TestBaseFileLookup(t *testing.T) {
 	})
 
 	t.Run("dependency", func(t *testing.T) {
-		var osPkg *gpd.FlatPackage
+		var osPkg *FlatPackage
 		for _, p := range resp.Packages {
 			if p.ID == osPkgID || p.ID == bzlmodOsPkgID {
 				osPkg = p
@@ -150,17 +141,7 @@ 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)
-	}
-
+	resp := runForTest(t, "file=hello_external_test.go")
 	if len(resp.Roots) != 2 {
 		t.Errorf("Expected exactly two roots for package: %+v", resp.Roots)
 	}
@@ -186,7 +167,74 @@ func TestExternalTests(t *testing.T) {
 	}
 }
 
+func runForTest(t *testing.T, args ...string) driverResponse {
+	t.Helper()
+
+	// Remove most environment variables, other than those on an allowlist.
+	//
+	// Bazel sets TEST_* and RUNFILES_* and a bunch of other variables.
+	// If Bazel is invoked when these variables, it assumes (correctly)
+	// that it's being invoked by a test, and it does different things that
+	// we don't want. For example, it randomizes the output directory, which
+	// is extremely expensive here. Out test framework creates an output
+	// directory shared among go_bazel_tests and points to it using .bazelrc.
+	//
+	// This only works if TEST_TMPDIR is not set when invoking bazel.
+	// bazel_testing.BazelCmd normally unsets that, but since gopackagesdriver
+	// invokes bazel directly, we need to unset it here.
+	allowEnv := map[string]struct{}{
+		"HOME":        {},
+		"PATH":        {},
+		"PWD":         {},
+		"SYSTEMDRIVE": {},
+		"SYSTEMROOT":  {},
+		"TEMP":        {},
+		"TMP":         {},
+		"TZ":          {},
+		"USER":        {},
+	}
+	var oldEnv []string
+	for _, env := range os.Environ() {
+		key, value, cut := strings.Cut(env, "=")
+		if !cut {
+			continue
+		}
+		if _, allowed := allowEnv[key]; !allowed {
+			os.Unsetenv(key)
+			oldEnv = append(oldEnv, key, value)
+		}
+	}
+	defer func() {
+		for i := 0; i < len(oldEnv); i += 2 {
+			os.Setenv(oldEnv[i], oldEnv[i+1])
+		}
+	}()
+
+	// Set workspaceRoot global variable.
+	// It's initialized to the BUILD_WORKSPACE_DIRECTORY environment variable
+	// before this point.
+	wd, err := os.Getwd()
+	if err != nil {
+		t.Fatal(err)
+	}
+	oldWorkspaceRoot := workspaceRoot
+	workspaceRoot = wd
+	defer func() { workspaceRoot = oldWorkspaceRoot }()
+
+	in := strings.NewReader("{}")
+	out := &bytes.Buffer{}
+	if err := run(context.Background(), in, out, args); err != nil {
+		t.Fatalf("running gopackagesdriver: %v", err)
+	}
+	var resp driverResponse
+	if err := json.Unmarshal(out.Bytes(), &resp); err != nil {
+		t.Fatalf("unmarshaling response: %v", err)
+	}
+	return resp
+}
+
 func assertSuffixesInList(t *testing.T, list []string, expectedSuffixes ...string) {
+	t.Helper()
 	for _, suffix := range expectedSuffixes {
 		itemFound := false
 		for _, listItem := range list {
diff --git a/go/tools/gopackagesdriver/json_packages_driver.go b/go/tools/gopackagesdriver/json_packages_driver.go
index 00b55b3c..ec71ed44 100644
--- a/go/tools/gopackagesdriver/json_packages_driver.go
+++ b/go/tools/gopackagesdriver/json_packages_driver.go
@@ -23,7 +23,7 @@ type JSONPackagesDriver struct {
 	registry *PackageRegistry
 }
 
-func NewJSONPackagesDriver(jsonFiles []string, prf PathResolverFunc, bazelVersion string) (*JSONPackagesDriver, error) {
+func NewJSONPackagesDriver(jsonFiles []string, prf PathResolverFunc, bazelVersion bazelVersion) (*JSONPackagesDriver, error) {
 	jpd := &JSONPackagesDriver{
 		registry: NewPackageRegistry(bazelVersion),
 	}
diff --git a/go/tools/gopackagesdriver/main.go b/go/tools/gopackagesdriver/main.go
index 0213f6f8..0613a3be 100644
--- a/go/tools/gopackagesdriver/main.go
+++ b/go/tools/gopackagesdriver/main.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"os"
 	"runtime"
 	"strings"
@@ -72,54 +73,56 @@ var (
 	}
 )
 
-func run() (*driverResponse, error) {
-	ctx, cancel := signalContext(context.Background(), os.Interrupt)
-	defer cancel()
-
-	queries := os.Args[1:]
+func run(ctx context.Context, in io.Reader, out io.Writer, args []string) error {
+	queries := args
 
-	request, err := ReadDriverRequest(os.Stdin)
+	request, err := ReadDriverRequest(in)
 	if err != nil {
-		return emptyResponse, fmt.Errorf("unable to read request: %w", err)
+		return fmt.Errorf("unable to read request: %w", err)
 	}
 
 	bazel, err := NewBazel(ctx, bazelBin, workspaceRoot, bazelStartupFlags)
 	if err != nil {
-		return emptyResponse, fmt.Errorf("unable to create bazel instance: %w", err)
+		return fmt.Errorf("unable to create bazel instance: %w", err)
 	}
 
 	bazelJsonBuilder, err := NewBazelJSONBuilder(bazel, request.Tests)
 	if err != nil {
-		return emptyResponse, fmt.Errorf("unable to build JSON files: %w", err)
+		return fmt.Errorf("unable to build JSON files: %w", err)
 	}
 
 	labels, err := bazelJsonBuilder.Labels(ctx, queries)
 	if err != nil {
-		return emptyResponse, fmt.Errorf("unable to lookup package: %w", err)
+		return fmt.Errorf("unable to lookup package: %w", err)
 	}
 
 	jsonFiles, err := bazelJsonBuilder.Build(ctx, labels, request.Mode)
 	if err != nil {
-		return emptyResponse, fmt.Errorf("unable to build JSON files: %w", err)
+		return fmt.Errorf("unable to build JSON files: %w", err)
 	}
 
 	driver, err := NewJSONPackagesDriver(jsonFiles, bazelJsonBuilder.PathResolver(), bazel.version)
 	if err != nil {
-		return emptyResponse, fmt.Errorf("unable to load JSON files: %w", err)
+		return fmt.Errorf("unable to load JSON files: %w", err)
 	}
 
 	// Note: we are returning all files required to build a specific package.
 	// For file queries (`file=`), this means that the CompiledGoFiles will
 	// include more than the only file being specified.
-	return driver.GetResponse(labels), nil
+	resp := driver.GetResponse(labels)
+	data, err := json.Marshal(resp)
+	if err != nil {
+		return fmt.Errorf("unable to marshal response: %v", err)
+	}
+	_, err = out.Write(data)
+	return err
 }
 
 func main() {
-	response, err := run()
-	if err := json.NewEncoder(os.Stdout).Encode(response); err != nil {
-		fmt.Fprintf(os.Stderr, "unable to encode response: %v", err)
-	}
-	if err != nil {
+	ctx, cancel := signalContext(context.Background(), os.Interrupt)
+	defer cancel()
+
+	if err := run(ctx, os.Stdin, os.Stdout, os.Args[1:]); err != nil {
 		fmt.Fprintf(os.Stderr, "error: %v", err)
 		// gopls will check the packages driver exit code, and if there is an
 		// error, it will fall back to go list. Obviously we don't want that,
diff --git a/go/tools/gopackagesdriver/packageregistry.go b/go/tools/gopackagesdriver/packageregistry.go
index 36197707..2dc85faa 100644
--- a/go/tools/gopackagesdriver/packageregistry.go
+++ b/go/tools/gopackagesdriver/packageregistry.go
@@ -17,21 +17,20 @@ package main
 import (
 	"fmt"
 	"os"
-	"strconv"
 	"strings"
 )
 
 type PackageRegistry struct {
 	packagesByID map[string]*FlatPackage
 	stdlib       map[string]string
-	bazelVersion []int
+	bazelVersion bazelVersion
 }
 
-func NewPackageRegistry(bazelVersion string, pkgs ...*FlatPackage) *PackageRegistry {
+func NewPackageRegistry(bazelVersion bazelVersion, pkgs ...*FlatPackage) *PackageRegistry {
 	pr := &PackageRegistry{
 		packagesByID: map[string]*FlatPackage{},
 		stdlib:       map[string]string{},
-		bazelVersion: parseVersion(bazelVersion),
+		bazelVersion: bazelVersion,
 	}
 	pr.Add(pkgs...)
 	return pr
@@ -102,7 +101,7 @@ func (pr *PackageRegistry) Match(labels []string) ([]string, []*FlatPackage) {
 
 	for _, label := range labels {
 		// When packagesdriver is ran from rules go, rulesGoRepositoryName will just be @
-		if pr.bazelVersion[0] >= 6 &&
+		if pr.bazelVersion.compare(bazelVersion{6, 0, 0}) >= 0 &&
 			!strings.HasPrefix(label, "@") {
 			// Canonical labels is only since Bazel 6.0.0
 			label = fmt.Sprintf("@%s", label)
@@ -139,19 +138,3 @@ func (pr *PackageRegistry) Match(labels []string) ([]string, []*FlatPackage) {
 
 	return retRoots, retPkgs
 }
-
-func parseVersion(v string) []int {
-	parts := strings.Split(v, ".")
-	version := make([]int, len(parts))
-
-	var err error
-	for i, p := range parts {
-		version[i], err = strconv.Atoi(p)
-		if err != nil {
-			// Failsafe default
-			return []int{6, 0, 0}
-		}
-	}
-
-	return version
-}
diff --git a/tests/integration/gopackagesdriver/BUILD.bazel b/tests/integration/gopackagesdriver/BUILD.bazel
deleted file mode 100644
index 088663ed..00000000
--- a/tests/integration/gopackagesdriver/BUILD.bazel
+++ /dev/null
@@ -1,13 +0,0 @@
-load("//go/tools/bazel_testing:def.bzl", "go_bazel_test")
-
-go_bazel_test(
-    name = "gopackagesdriver_test",
-    size = "enormous",
-    srcs = ["gopackagesdriver_test.go"],
-    rule_files = [
-        "//:all_files",
-    ],
-    deps = [
-        "@io_bazel_rules_go//go/tools/gopackagesdriver:gopackagesdriver_lib",
-    ]
-)
diff --git a/tests/integration/gopackagesdriver/README.rst b/tests/integration/gopackagesdriver/README.rst
deleted file mode 100644
index 89d2d9ce..00000000
--- a/tests/integration/gopackagesdriver/README.rst
+++ /dev/null
@@ -1,8 +0,0 @@
-Go Packages Driver
-
-gopackagesdriver_test
---------------------
-Verifies that the output of the go packages driver includes the correct output.
-
-Go x/tools is very sensitive to inaccuracies in the package output, so we should
-validate each added feature against what is expected by x/tools.