Always initialize yatest lib for go

@snermolaev предложил в тикете вынести инициализацию окружения в отдельный модуль и для каждого модуля импортировать его, чтобы не делать лишней работы. Но кажется, что при таком подходе, потом нужно будет либо повторно считывать и парсить контекстный файл, либо сохранять его контент внутри нового модуля и тоже заново парсить.

Учитывая, что для инициализации переменных окружения нам все равно придется считать и распарсить контекстный файл, предлагаю оставить инициализацию как есть и вместе с переменными окружния, практически за бесплатно инициализировать весь контекст.

6 files changed, 364 insertions, 0 deletions

diff --git a/build/conf/go.conf b/build/conf/go.conf
index aaafab1074..9e338c3d4f 100644
--- a/build/conf/go.conf
+++ b/build/conf/go.conf
@@ -935,6 +935,8 @@ module GO_TEST: GO_PROGRAM {
     PEERDIR(${GOSTD}/testing/internal/testdeps)
     PEERDIR(${GOSTD}/testing)
 
+    PEERDIR(library/go/test/yatest)
+
     ADD_YTEST($MODULE_PREFIX$REALPRJNAME go.test)
     ADD_YTEST($MODULE_PREFIX$REALPRJNAME go.bench)
 
diff --git a/build/scripts/go_tool.py b/build/scripts/go_tool.py
index 1560eb33cb..b6b0f4e994 100644
--- a/build/scripts/go_tool.py
+++ b/build/scripts/go_tool.py
@@ -666,6 +666,7 @@ def gen_test_main(args, test_lib_args, xtest_lib_args):
     if test_main_package is None:
         lines.append('    "os"')
     lines.extend(['    "testing"', '    "testing/internal/testdeps"'])
+    lines.extend(['    _ "a.yandex-team.ru/library/go/test/yatest"'])
 
     if len(tests) > 0:
         lines.append('    _test "{}"'.format(test_module_path))
diff --git a/library/go/test/yatest/arcadia.go b/library/go/test/yatest/arcadia.go
new file mode 100644
index 0000000000..f22c6d2f04
--- /dev/null
+++ b/library/go/test/yatest/arcadia.go
@@ -0,0 +1,20 @@
+//go:build arcadia
+// +build arcadia
+
+package yatest
+
+import (
+	"os"
+)
+
+func doInit() {
+	initTestContext()
+}
+
+func init() {
+	if val := os.Getenv("YA_TEST_CONTEXT_FILE"); val != "" {
+		if _, err := os.Stat(val); err == nil {
+			lazyInit()
+		}
+	}
+}
diff --git a/library/go/test/yatest/env.go b/library/go/test/yatest/env.go
new file mode 100644
index 0000000000..d655e4c6da
--- /dev/null
+++ b/library/go/test/yatest/env.go
@@ -0,0 +1,307 @@
+// Package yatest provides access to testing context, when running under ya make -t.
+package yatest
+
+import (
+	"bufio"
+	"encoding/json"
+	"fmt"
+	"os"
+	"path"
+	"path/filepath"
+	"runtime"
+	"sync"
+)
+
+type TestContext struct {
+	Build struct {
+		BuildType string            `json:"build_type"`
+		Flags     map[string]string `json:"flags"`
+		Sanitizer string            `json:"sanitizer"`
+	} `json:"build"`
+
+	Runtime struct {
+		BuildRoot              string            `json:"build_root"`
+		OutputPath             string            `json:"output_path"`
+		ProjectPath            string            `json:"project_path"`
+		PythonBin              string            `json:"python_bin"`
+		PythonLibPath          string            `json:"python_lib_path"`
+		RAMDrivePath           string            `json:"ram_drive_path"`
+		YtHDDPath              string            `json:"yt_hdd_path"`
+		TestOutputRAMDrivePath string            `json:"test_output_ram_drive_path"`
+		SourceRoot             string            `json:"source_root"`
+		WorkPath               string            `json:"work_path"`
+		TestToolPath           string            `json:"test_tool_path"`
+		TestParameters         map[string]string `json:"test_params"`
+	} `json:"runtime"`
+
+	Resources struct {
+		Global map[string]string `json:"global"`
+	} `json:"resources"`
+
+	Internal struct {
+		EnvFile string `json:"env_file"`
+	} `json:"internal"`
+
+	Initialized bool
+}
+
+var (
+	context              TestContext
+	isRunningUnderGoTest bool
+	initOnce             sync.Once
+)
+
+func lazyInit() {
+	initOnce.Do(doInit)
+}
+
+func verifyContext() {
+	if !context.Initialized {
+		panic("test context isn't initialized")
+	}
+}
+
+func initTestContext() {
+	data, err := os.ReadFile(getenv("YA_TEST_CONTEXT_FILE"))
+	if err != nil {
+		panic(err)
+	}
+
+	err = json.Unmarshal([]byte(data), &context)
+	if err != nil {
+		panic(err)
+	}
+	setupEnv()
+	context.Initialized = true
+}
+
+func setupEnv() {
+	file, err := os.Open(context.Internal.EnvFile)
+	if err != nil {
+		return
+	}
+	defer file.Close()
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.Text()
+		var objmap map[string]json.RawMessage
+		var val string
+		err := json.Unmarshal([]byte(line), &objmap)
+		if err != nil {
+			panic(err)
+		}
+		for k, v := range objmap {
+			err := json.Unmarshal(v, &val)
+			if err != nil {
+				panic(err)
+			}
+			err = os.Setenv(k, val)
+			if err != nil {
+				panic(err)
+			}
+		}
+	}
+}
+
+func HasYaTestContext() bool {
+	lazyInit()
+	return context.Initialized
+}
+
+// GlobalResourcePath returns absolute path to a directory
+// containing global build resource.
+func GlobalResourcePath(name string) string {
+	resource, ok := RelaxedGlobalResourcePath(name)
+	if !ok {
+		panic(fmt.Sprintf("global resource %s is not defined", name))
+	}
+	return resource
+}
+
+func RelaxedGlobalResourcePath(name string) (string, bool) {
+	lazyInit()
+	verifyContext()
+	resource, ok := context.Resources.Global[name]
+	return resource, ok
+}
+
+func getenv(name string) string {
+	value := os.Getenv(name)
+	if value == "" {
+		panic(fmt.Sprintf("environment variable %s is not set", name))
+	}
+	return value
+}
+
+func CCompilerPath() string {
+	lazyInit()
+	return getenv("YA_CC")
+}
+
+func CxxCompilerPath() string {
+	lazyInit()
+	return getenv("YA_CXX")
+}
+
+// PythonBinPath returns absolute path to the python
+//
+// Warn: if you are using build with system python (-DUSE_SYSTEM_PYTHON=X) beware that some python bundles
+// are built in a stripped-down form that is needed for building, not running tests.
+// See comments in the file below to find out which version of python is compatible with tests.
+// https://a.yandex-team.ru/arc/trunk/arcadia/build/platform/python/resources.inc
+func PythonBinPath() string {
+	lazyInit()
+	verifyContext()
+	return context.Runtime.PythonBin
+}
+
+func PythonLibPath() string {
+	lazyInit()
+	verifyContext()
+	return context.Runtime.PythonLibPath
+}
+
+// SourcePath returns absolute path to source directory.
+//
+// arcadiaPath must be declared using DATA macro inside ya.make.
+func SourcePath(arcadiaPath string) string {
+	lazyInit()
+	if path.IsAbs(arcadiaPath) {
+		panic(fmt.Sprintf("relative path expected, but got %q", arcadiaPath))
+	}
+
+	// Don't verify context for SourcePath - it can be mined without context
+	return filepath.Join(context.Runtime.SourceRoot, arcadiaPath)
+}
+
+// BuildPath returns absolute path to the build directory.
+func BuildPath(dataPath string) string {
+	lazyInit()
+	if path.IsAbs(dataPath) {
+		panic(fmt.Sprintf("relative path expected, but got %q", dataPath))
+	}
+
+	verifyContext()
+	return filepath.Join(context.Runtime.BuildRoot, dataPath)
+}
+
+// WorkPath returns absolute path to the work directory (initial test cwd).
+func WorkPath(dataPath string) string {
+	lazyInit()
+	if path.IsAbs(dataPath) {
+		panic(fmt.Sprintf("relative path expected, but got %q", dataPath))
+	}
+
+	verifyContext()
+	return filepath.Join(context.Runtime.WorkPath, dataPath)
+}
+
+// OutputPath returns absolute path to the output directory (testing_out_stuff).
+func OutputPath(dataPath string) string {
+	lazyInit()
+	verifyContext()
+	return filepath.Join(context.Runtime.OutputPath, dataPath)
+}
+
+// RAMDrivePath returns absolute path to the ramdrive directory
+func RAMDrivePath(dataPath string) string {
+	lazyInit()
+	if path.IsAbs(dataPath) {
+		panic(fmt.Sprintf("relative path expected, but got %q", dataPath))
+	}
+
+	verifyContext()
+	return filepath.Join(context.Runtime.RAMDrivePath, dataPath)
+}
+
+// YtHDDPath returns absolute path to the directory mounted to ext4 fs in YT
+func YtHDDPath(dataPath string) string {
+	lazyInit()
+	if path.IsAbs(dataPath) {
+		panic(fmt.Sprintf("relative path expected, but got %q", dataPath))
+	}
+
+	verifyContext()
+	return filepath.Join(context.Runtime.YtHDDPath, dataPath)
+}
+
+// OutputRAMDrivePath returns absolute path to the ramdrive output directory
+func OutputRAMDrivePath(dataPath string) string {
+	lazyInit()
+	if path.IsAbs(dataPath) {
+		panic(fmt.Sprintf("relative path expected, but got %q", dataPath))
+	}
+
+	verifyContext()
+	return filepath.Join(context.Runtime.TestOutputRAMDrivePath, dataPath)
+}
+
+// HasRAMDrive returns true if ramdrive is enabled in tests
+func HasRAMDrive() bool {
+	lazyInit()
+	verifyContext()
+	return context.Runtime.RAMDrivePath != ""
+}
+
+func BinaryPath(dataPath string) (string, error) {
+	if runtime.GOOS == "windows" {
+		dataPath += ".exe"
+	}
+	buildPath := BuildPath("")
+	binaryPath := filepath.Join(buildPath, dataPath)
+	if _, err := os.Stat(binaryPath); os.IsNotExist(err) {
+		return "", fmt.Errorf("cannot find binary %s: make sure it was added in the DEPENDS section", dataPath)
+	} else {
+		return binaryPath, nil
+	}
+}
+
+// ProjectPath returns arcadia-relative path to the test project
+func ProjectPath() string {
+	lazyInit()
+	verifyContext()
+	return context.Runtime.ProjectPath
+}
+
+// BuildType returns build type that was used to compile the test
+func BuildType() string {
+	lazyInit()
+	verifyContext()
+	return context.Build.BuildType
+}
+
+// BuildFlag returns the value of the requested build flag
+func BuildFlag(name string) (string, bool) {
+	lazyInit()
+	verifyContext()
+	val, ok := context.Build.Flags[name]
+	return val, ok
+}
+
+// TestParam returns the value of the requested test parameter
+func TestParam(name string) (string, bool) {
+	lazyInit()
+	verifyContext()
+	val, ok := context.Runtime.TestParameters[name]
+	return val, ok
+}
+
+// Sanitizer returns sanitizer name that was used to compile the test
+func Sanitizer() string {
+	lazyInit()
+	verifyContext()
+	return context.Build.Sanitizer
+}
+
+func EnvFile() string {
+	lazyInit()
+	verifyContext()
+	return context.Internal.EnvFile
+}
+
+func TestToolPath() string {
+	lazyInit()
+	verifyContext()
+	return context.Runtime.TestToolPath
+}
diff --git a/library/go/test/yatest/go.go b/library/go/test/yatest/go.go
new file mode 100644
index 0000000000..cc03951257
--- /dev/null
+++ b/library/go/test/yatest/go.go
@@ -0,0 +1,21 @@
+package yatest
+
+import (
+	"os"
+)
+
+func PrepareGOPATH() error {
+	return preparePath("GOPATH")
+}
+
+func PrepareGOCACHE() error {
+	return preparePath("GOCACHE")
+}
+
+func preparePath(name string) error {
+	p, err := os.MkdirTemp(WorkPath(""), "name")
+	if err != nil {
+		return err
+	}
+	return os.Setenv(name, p)
+}
diff --git a/library/go/test/yatest/ya.make b/library/go/test/yatest/ya.make
new file mode 100644
index 0000000000..8f2442686f
--- /dev/null
+++ b/library/go/test/yatest/ya.make
@@ -0,0 +1,13 @@
+GO_LIBRARY()
+
+SRCS(
+    arcadia.go
+    env.go
+    go.go
+)
+
+GO_TEST_SRCS(env_test.go)
+
+END()
+
+RECURSE_FOR_TESTS(gotest)