internal/regtest/regtest.go

// Copyright 2023 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package regtest provides helpers for end-to-end testing
// involving counter and upload packages.
package regtest

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime/debug"
	"strings"
	"testing"

	"golang.org/x/telemetry/counter/countertest"
)

const (
	telemetryDirEnvVar = "_COUNTERTEST_RUN_TELEMETRY_DIR"
	entryPointEnvVar   = "_COUNTERTEST_ENTRYPOINT"
)

var (
	telemetryDirEnvVarValue = os.Getenv(telemetryDirEnvVar)
	entryPointEnvVarValue   = os.Getenv(entryPointEnvVar)
)

// Program is a value that can be used to identify a program in the test.
type Program string

// NewProgram returns a Program value that can be used to identify a program
// to run by RunProg. The program must be registered with NewProgram before
// the first call to RunProg in the test function.
//
// RunProg runs this binary in a separate process with special environment
// variables that specify the entry point. When this binary runs with the
// environment variables that match the specified name, NewProgram calls
// the given fn and exits with the return value. Note that all the code
// before NewProgram is executed in both the main process and the subprocess.
func NewProgram(t *testing.T, name string, fn func() int) Program {
	if telemetryDirEnvVarValue != "" && entryPointEnvVarValue == name {
		// We are running the separate process that was spawned by RunProg.
		fmt.Fprintf(os.Stderr, "running program %q\n", name)
		countertest.Open(telemetryDirEnvVarValue)
		os.Exit(fn())
	}

	testName, _, _ := strings.Cut(t.Name(), "/")
	registered, ok := registeredPrograms[testName]
	if !ok {
		registered = make(map[string]bool)
	}
	if registered[name] {
		t.Fatalf("program %q was already registered", name)
	}
	registered[name] = true
	return Program(name)
}

// registeredPrograms stores all registered program names to detect duplicate registrations.
var registeredPrograms = make(map[string]map[string]bool) // test name -> program name -> exist

// RunProg runs the program prog in a separate process with the specified
// telemetry directory. RunProg can be called multiple times in the same test,
// but all the programs must be registered with NewProgram before the first
// call to RunProg.
func RunProg(t *testing.T, telemetryDir string, prog Program) ([]byte, error) {
	if telemetryDirEnvVarValue != "" {
		fmt.Fprintf(os.Stderr, "unknown program %q\n %s %s", prog, telemetryDirEnvVarValue, entryPointEnvVarValue)
		os.Exit(2)
	}
	testName, _, _ := strings.Cut(t.Name(), "/")
	testBin, err := os.Executable()
	if err != nil {
		return nil, fmt.Errorf("cannot determine the current process's executable name: %v", err)
	}

	// Spawn a subprocess to run the 'prog' by setting telemetryDirEnvVar.
	cmd := exec.Command(testBin, "-test.run", fmt.Sprintf("^%s$", testName))
	cmd.Env = append(cmd.Env, telemetryDirEnvVar+"="+telemetryDir, entryPointEnvVar+"="+string(prog))
	return cmd.CombinedOutput()
}

// ProgInfo returns the go version, program name and version info the process would record in its counter file.
func ProgInfo(t *testing.T) (goVersion, progVersion, progName string) {
	info, ok := debug.ReadBuildInfo()
	if !ok {
		t.Fatal("cannot read build info - it's likely this setup is unsupported by the counter package")
	}
	goVers := info.GoVersion
	if strings.Contains(goVers, "devel") || strings.Contains(goVers, "-") {
		goVers = "devel"
	}
	progPkgPath := info.Path
	if progPkgPath == "" {
		progPkgPath = strings.TrimSuffix(filepath.Base(os.Args[0]), ".exe")
	}
	progVers := info.Main.Version
	if strings.Contains(progVers, "devel") || strings.Contains(progVers, "-") {
		progVers = "devel"
	}
	return goVers, progVers, progPkgPath
}