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

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.

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,

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
+}

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{

tests/integration/gopackagesdriver/gopackagesdriver_test.go → 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 {

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),
 	}