go_test: Add env attribute

go/core.rst

@@ -492,6 +492,13 @@ Attributes
 | List of Go libraries this test imports directly.                                                 |
 | These may be go_library rules or compatible rules with the GoLibrary_ provider.                  |
 +----------------------------+-----------------------------+---------------------------------------+
+| :param:`env`               | :type:`string_dict`         | :value:`{}`                           |
++----------------------------+-----------------------------+---------------------------------------+
+| Environment variables to set for the test execution.                                             |
+| The values (but not keys) are subject to                                                         |
+| [location expansion](https://docs.bazel.build/versions/main/skylark/macros.html) but not full    |
+| [make variable expansion](https://docs.bazel.build/versions/main/be/make-variables.html).        |
++----------------------------+-----------------------------+---------------------------------------+
 | :param:`embed`             | :type:`label_list`          | :value:`[]`                           |
 +----------------------------+-----------------------------+---------------------------------------+
 | List of Go libraries whose sources should be compiled together with this                         |

go/private/rules/test.bzl

@@ -154,6 +154,10 @@ def _go_test_impl(ctx):
         info_file = ctx.info_file,
     )
 
+    env = {}
+    for k, v in ctx.attr.env.items():
+        env[k] = ctx.expand_location(v, ctx.attr.data)
+
     # Bazel only looks for coverage data if the test target has an
     # InstrumentedFilesProvider. If the provider is found and at least one
     # source file is present, Bazel will set the COVERAGE_OUTPUT_FILE
@@ -175,6 +179,7 @@ def _go_test_impl(ctx):
             dependency_attributes = ["deps", "embed"],
             extensions = ["go"],
         ),
+        testing.TestEnvironment(env),
     ]
 
 _go_test_kwargs = {
@@ -185,6 +190,7 @@ _go_test_kwargs = {
         "deps": attr.label_list(providers = [GoLibrary]),
         "embed": attr.label_list(providers = [GoLibrary]),
         "embedsrcs": attr.label_list(allow_files = True),
+        "env": attr.string_dict(),
         "importpath": attr.string(),
         "gc_goopts": attr.string_list(),
         "gc_linkopts": attr.string_list(),

go/tools/bazel/bazel_test.go

@@ -255,7 +255,7 @@ func TestPythonManifest(t *testing.T) {
 		t.Errorf("failed to init runfiles: %v", runfiles.err)
 	}
 
-	entry, ok := runfiles.index["important.txt"]
+	entry, ok := runfiles.index.GetIgnoringWorkspace("important.txt")
 	if !ok {
 		t.Errorf("failed to locate runfile %s in index", "important.txt")
 	}

go/tools/bazel/runfiles.go

@@ -51,10 +51,21 @@ func Runfile(path string) (string, error) {
 	}
 
 	// Search manifest if we have one.
-	if entry, ok := runfiles.index[path]; ok {
+	if entry, ok := runfiles.index.GetIgnoringWorkspace(path); ok {
 		return entry.Path, nil
 	}
 
+	if strings.HasPrefix(path, "../") || strings.HasPrefix(path, "external/") {
+		pathParts := strings.Split(path, "/")
+		if len(pathParts) >= 3 {
+			workspace := pathParts[1]
+			pathInsideWorkspace := strings.Join(pathParts[2:], "/")
+			if path := runfiles.index.Get(workspace, pathInsideWorkspace); path != "" {
+				return path, nil
+			}
+		}
+    }
+
 	// Search the main workspace.
 	if runfiles.workspace != "" {
 		mainPath := filepath.Join(runfiles.dir, runfiles.workspace, path)
@@ -279,7 +290,7 @@ var runfiles = struct {
 	list []RunfileEntry
 
 	// index maps runfile short paths to absolute paths.
-	index map[string]RunfileEntry
+	index index
 
 	// dir is a path to the runfile directory. Typically this is a directory
 	// named <target>.runfiles, with a subdirectory for each workspace.
@@ -296,6 +307,47 @@ var runfiles = struct {
 	err error
 }{}
 
+type index struct {
+	indexWithWorkspace map[indexKey]*RunfileEntry
+	indexIgnoringWorksapce map[string]*RunfileEntry
+}
+
+func newIndex() index {
+	return index {
+		indexWithWorkspace: make(map[indexKey]*RunfileEntry),
+		indexIgnoringWorksapce: make(map[string]*RunfileEntry),
+	}
+}
+
+func (i *index) Put(entry *RunfileEntry) {
+	i.indexWithWorkspace[indexKey{
+		workspace: entry.Workspace,
+		shortPath: entry.ShortPath,
+	}] = entry
+	i.indexIgnoringWorksapce[entry.ShortPath] = entry
+}
+
+func (i *index) Get(workspace string, shortPath string) string {
+	entry := i.indexWithWorkspace[indexKey{
+		workspace: workspace,
+		shortPath: shortPath,
+	}]
+	if entry == nil {
+		return ""
+	}
+	return entry.Path
+}
+
+func (i *index) GetIgnoringWorkspace(shortPath string) (*RunfileEntry, bool) {
+	entry, ok := i.indexIgnoringWorksapce[shortPath]
+	return entry, ok
+}
+
+type indexKey struct {
+	workspace string
+	shortPath string
+}
+
 func ensureRunfiles() error {
 	runfiles.once.Do(initRunfiles)
 	return runfiles.err
@@ -307,7 +359,7 @@ func initRunfiles() {
 		// On Windows, Bazel doesn't create a symlink tree of runfiles because
 		// Windows doesn't support symbolic links by default. Instead, runfile
 		// locations are written to a manifest file.
-		runfiles.index = make(map[string]RunfileEntry)
+		runfiles.index = newIndex()
 		data, err := ioutil.ReadFile(manifest)
 		if err != nil {
 			runfiles.err = err
@@ -361,7 +413,7 @@ func initRunfiles() {
 			}
 
 			runfiles.list = append(runfiles.list, entry)
-			runfiles.index[entry.ShortPath] = entry
+			runfiles.index.Put(&entry)
 		}
 	}
 

tests/core/go_test/BUILD.bazel

@@ -213,3 +213,15 @@ go_test(
     name = "fuzz_test",
     srcs = ["fuzz_test.go"],
 )
+
+go_test(
+    name = "env_test",
+    srcs = ["env_test.go"],
+    data = ["@go_sdk//:lib/time/zoneinfo.zip"],
+    env = {
+        "ZONEINFO": "$(execpath @go_sdk//:lib/time/zoneinfo.zip)",
+    },
+    deps = [
+        "@io_bazel_rules_go//go/tools/bazel",
+    ],
+)

tests/core/go_test/env_test.go

@@ -0,0 +1,39 @@
+// Copyright 2021 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package env_test
+
+import (
+    "os"
+    "testing"
+
+    "github.com/bazelbuild/rules_go/go/tools/bazel"
+)
+
+func TestEnv(t *testing.T) {
+    v := os.Getenv("ZONEINFO")
+    if v == "" {
+        t.Fatalf("ZONEINFO env var was empty")
+    }
+
+    path, err := bazel.Runfile(v)
+    if err != nil {
+        t.Fatalf("Could not find runfile %v: %v", v, err)
+    }
+
+    if _, err := os.Stat(path); err != nil {
+        t.Fatalf("Could not find file at env var $ZONEINFO (value: %v) at path %v: %v", v, path, err)
+    }
+}
+