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.

//go:build go1.21

// Package regtest provides helpers for end-to-end testing
// involving counter and upload packages. This package requires go1.21 or newer.
package regtest

import (
	"flag"
	"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"
)

// Main is a test main function for use in TestMain, which runs one of the
// given programs when invoked as a separate process via RunProg.
//
// The return value of each program is the exit code of the process.
func Main(m *testing.M, programs map[string]func() int) {
	if d := os.Getenv(telemetryDirEnvVar); d != "" {
		countertest.Open(d)
	}
	if e, ok := os.LookupEnv(entryPointEnvVar); ok {
		if prog, ok := programs[e]; ok {
			os.Exit(prog())
		}
		fmt.Fprintf(os.Stderr, "unknown program %q", e)
		os.Exit(2)
	}
	flag.Parse()
	os.Exit(m.Run())
}

// RunProg runs the named program in a separate process with the specified
// telemetry directory, where prog is one of the programs passed to Main (which
// must be invoked by TestMain).
func RunProg(telemetryDir string, prog string) ([]byte, error) {
	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 subprocessKeyEnvVar and telemetryDirEnvVar.
	cmd := exec.Command(testBin)
	cmd.Env = append(cmd.Env, telemetryDirEnvVar+"="+telemetryDir, entryPointEnvVar+"="+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
}