main.go

// bktec fetches and runs test plans generated by Buildkite
// Test Engine.
package main

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strconv"
	"syscall"
	"time"

	"github.com/buildkite/test-engine-client/internal/api"
	"github.com/buildkite/test-engine-client/internal/config"
	"github.com/buildkite/test-engine-client/internal/debug"
	"github.com/buildkite/test-engine-client/internal/plan"
	"github.com/buildkite/test-engine-client/internal/runner"
	"github.com/olekukonko/tablewriter"
	"golang.org/x/sys/unix"
)

var Version = ""

const Logo = `
______ ______ _____
___  /____  /___  /____________
__  __ \_  //_/  __/  _ \  ___/
_  /_/ /  ,<  / /_ /  __/ /__
/_.___//_/|_| \__/ \___/\___/
`

func printStartUpMessage() {
	const green = "\033[32m"
	const reset = "\033[0m"
	fmt.Println("+++ Buildkite Test Engine Client: bktec " + Version + "\n")
	fmt.Println(green + Logo + reset)
}

type TestRunner interface {
	Run(result *runner.RunResult, testCases []plan.TestCase, retry bool) error
	GetExamples(files []string) ([]plan.TestCase, error)
	GetFiles() ([]string, error)
	Name() string
}

func main() {
	debug.SetDebug(os.Getenv("BUILDKITE_TEST_ENGINE_DEBUG_ENABLED") == "true")

	versionFlag := flag.Bool("version", false, "print version information")

	flag.Parse()

	if *versionFlag {
		fmt.Println(Version)
		os.Exit(0)
	}

	printStartUpMessage()

	// get config
	cfg, err := config.New()
	if err != nil {
		logErrorAndExit(16, "Invalid configuration...\n%v", err)
	}

	testRunner, err := runner.DetectRunner(cfg)
	if err != nil {
		logErrorAndExit(16, "Unsupported value for BUILDKITE_TEST_ENGINE_TEST_RUNNER %q: %v", cfg.TestRunner, err)
	}

	files, err := testRunner.GetFiles()
	if err != nil {
		logErrorAndExit(16, "Couldn't get files: %v", err)
	}

	// get plan
	ctx := context.Background()
	apiClient := api.NewClient(api.ClientConfig{
		ServerBaseUrl:    cfg.ServerBaseUrl,
		AccessToken:      cfg.AccessToken,
		OrganizationSlug: cfg.OrganizationSlug,
		Version:          Version,
	})

	testPlan, err := fetchOrCreateTestPlan(ctx, apiClient, cfg, files, testRunner)
	if err != nil {
		logErrorAndExit(16, "Couldn't fetch or create test plan: %v", err)
	}

	debug.Printf("My favourite ice cream is %s", testPlan.Experiment)

	// get plan for this node
	thisNodeTask := testPlan.Tasks[strconv.Itoa(cfg.NodeIndex)]

	// execute tests
	var timeline []api.Timeline
	runResult, err := runTestsWithRetry(testRunner, &thisNodeTask.Tests, cfg.MaxRetries, testPlan.MutedTests, &timeline)

	// Handle errors that prevent the runner from finishing.
	// By finishing, it means that the runner has completed with a readable result.
	if err != nil {
		// runner terminated by signal: exit with 128 + signal number
		if ProcessSignaledError := new(runner.ProcessSignaledError); errors.As(err, &ProcessSignaledError) {
			logSignalAndExit(testRunner.Name(), ProcessSignaledError.Signal)
		}

		// runner exited with error: exit with the exit code
		if exitError := new(exec.ExitError); errors.As(err, &exitError) {
			logErrorAndExit(exitError.ExitCode(), "%s exited with error: %v", testRunner.Name(), err)
		}

		// other errors: exit with 16
		logErrorAndExit(16, "Couldn't run tests: %v", err)
	}

	// At this point, the runner is expected to have completed

	if !testPlan.Fallback {
		sendMetadata(ctx, apiClient, cfg, timeline, runResult.Statistics())
	}

	printReport(runResult, testPlan.SkippedTests, testRunner.Name())

	if runResult.Status() == runner.RunStatusFailed || runResult.Status() == runner.RunStatusError {
		os.Exit(1)
	}
}

func printReport(runResult runner.RunResult, testsSkippedByTestEngine []plan.TestCase, runnerName string) {
	fmt.Println("+++ ========== Buildkite Test Engine Report  ==========")

	switch runResult.Status() {
	case runner.RunStatusPassed:
		fmt.Println("✅ All tests passed.")
	case runner.RunStatusFailed:
		fmt.Println("❌ Some tests failed.")
	case runner.RunStatusError:
		fmt.Printf("🚨 %s\n", runResult.Error())
	}
	fmt.Println("")

	// Print statistics
	statistics := runResult.Statistics()
	data := [][]string{
		{"Passed", "first run", strconv.Itoa(statistics.PassedOnFirstRun)},
		{"Passed", "on retry", strconv.Itoa(statistics.PassedOnRetry)},
		{"Muted", "passed", strconv.Itoa(statistics.MutedPassed)},
		{"Muted", "failed", strconv.Itoa(statistics.MutedFailed)},
		{"Failed", "", strconv.Itoa(statistics.Failed)},
		{"Skipped", "", strconv.Itoa(statistics.Skipped)},
	}
	table := tablewriter.NewWriter(os.Stdout)
	table.AppendBulk(data)
	table.SetFooter([]string{"", "Total", strconv.Itoa(statistics.Total)})
	table.SetFooterAlignment(tablewriter.ALIGN_RIGHT)
	table.SetAutoMergeCellsByColumnIndex([]int{0, 1})
	table.SetRowLine(true)
	table.Render()

	// Print muted and failed tests
	mutedTests := runResult.MutedTests()
	if len(mutedTests) > 0 {
		fmt.Println("")
		fmt.Println("+++ Muted Tests:")
		for _, mutedTest := range runResult.MutedTests() {
			fmt.Printf("- %s %s (%s)\n", mutedTest.Scope, mutedTest.Name, mutedTest.Status)
		}
	}

	failedTests := runResult.FailedTests()
	if len(failedTests) > 0 {
		fmt.Println("")
		fmt.Println("+++ Failed Tests:")
		for _, failedTests := range runResult.FailedTests() {
			fmt.Printf("- %s %s\n", failedTests.Scope, failedTests.Name)
		}
	}

	testsSkippedByRunner := runResult.SkippedTests()
	if len(testsSkippedByRunner) > 0 {
		fmt.Println("")
		fmt.Printf("+++ Skipped by %s:\n", runnerName)
		for _, skippedTest := range testsSkippedByRunner {
			fmt.Printf("- %s %s\n", skippedTest.Scope, skippedTest.Name)
		}
	}

	if len(testsSkippedByTestEngine) > 0 {
		fmt.Println("")
		fmt.Println("+++ Skipped by Test Engine:")
		for _, skippedTest := range testsSkippedByTestEngine {
			fmt.Printf("- %s %s\n", skippedTest.Scope, skippedTest.Name)
		}
	}

	fmt.Println("===================================================")
}

func createTimestamp() string {
	return time.Now().Format(time.RFC3339Nano)
}

func sendMetadata(ctx context.Context, apiClient *api.Client, cfg config.Config, timeline []api.Timeline, statistics runner.RunStatistics) {
	err := apiClient.PostTestPlanMetadata(ctx, cfg.SuiteSlug, cfg.Identifier, api.TestPlanMetadataParams{
		Timeline:   timeline,
		Env:        cfg.DumpEnv(),
		Version:    Version,
		Statistics: statistics,
	})

	// Error is suppressed because we don't want to fail the build if we can't send metadata.
	if err != nil {
		fmt.Printf("Failed to send metadata to Test Engine: %v\n", err)
	}
}

func runTestsWithRetry(testRunner TestRunner, testsCases *[]plan.TestCase, maxRetries int, mutedTests []plan.TestCase, timeline *[]api.Timeline) (runner.RunResult, error) {
	attemptCount := 0

	// Create a new run result with muted tests to keep track of the results.
	runResult := runner.NewRunResult(mutedTests)

	for attemptCount <= maxRetries {
		if attemptCount == 0 {
			fmt.Printf("+++ Buildkite Test Engine Client: Running tests\n")
			*timeline = append(*timeline, api.Timeline{
				Event:     "test_start",
				Timestamp: createTimestamp(),
			})
		} else {
			fmt.Printf("+++ Buildkite Test Engine Client: ♻️ Attempt %d of %d to retry failing tests\n", attemptCount, maxRetries)
			*timeline = append(*timeline, api.Timeline{
				Event:     fmt.Sprintf("retry_%d_start", attemptCount),
				Timestamp: createTimestamp(),
			})
		}

		err := testRunner.Run(runResult, *testsCases, attemptCount > 0)

		if attemptCount == 0 {
			*timeline = append(*timeline, api.Timeline{
				Event:     "test_end",
				Timestamp: createTimestamp(),
			})
		} else {
			*timeline = append(*timeline, api.Timeline{
				Event:     fmt.Sprintf("retry_%d_end", attemptCount),
				Timestamp: createTimestamp(),
			})
		}

		// Don't retry if there is an error that is not a test failure.
		if err != nil {
			return *runResult, err
		}

		// Don't retry if we've reached max retries.
		if attemptCount == maxRetries {
			return *runResult, nil
		}

		// Don't retry if tests are passed.
		if runResult.Status() == runner.RunStatusPassed {
			return *runResult, nil
		}

		// Retry only if there are failed tests.
		failedTests := runResult.FailedTests()

		if len(failedTests) == 0 {
			return *runResult, nil
		}

		*testsCases = failedTests
		attemptCount++
	}

	return *runResult, nil
}

func logSignalAndExit(name string, signal syscall.Signal) {
	fmt.Printf("Buildkite Test Engine Client: %s was terminated with signal: %v (%v)\n", name, unix.SignalName(signal), signal)

	exitCode := 128 + int(signal)
	os.Exit(exitCode)
}

// logErrorAndExit logs an error message and exits with the given exit code.
func logErrorAndExit(exitCode int, format string, v ...any) {
	fmt.Printf("Buildkite Test Engine Client: "+format+"\n", v...)
	os.Exit(exitCode)
}

// fetchOrCreateTestPlan fetches a test plan from the server, or creates a
// fallback plan if the server is unavailable or returns an error plan.
func fetchOrCreateTestPlan(ctx context.Context, apiClient *api.Client, cfg config.Config, files []string, testRunner TestRunner) (plan.TestPlan, error) {
	debug.Println("Fetching test plan")

	// Fetch the plan from the server's cache.
	cachedPlan, err := apiClient.FetchTestPlan(ctx, cfg.SuiteSlug, cfg.Identifier, cfg.JobRetryCount)

	handleError := func(err error) (plan.TestPlan, error) {
		if errors.Is(err, api.ErrRetryTimeout) {
			fmt.Println("⚠️ Could not fetch or create plan from server, falling back to non-intelligent splitting. Your build may take longer than usual.")
			p := plan.CreateFallbackPlan(files, cfg.Parallelism)
			return p, nil
		}

		if billingError := new(api.BillingError); errors.As(err, &billingError) {
			fmt.Println(billingError.Message)
			fmt.Println("⚠️ Falling back to non-intelligent splitting. Your build may take longer than usual.")
			p := plan.CreateFallbackPlan(files, cfg.Parallelism)
			return p, nil
		}

		return plan.TestPlan{}, err
	}

	if err != nil {
		return handleError(err)
	}

	if cachedPlan != nil {
		// The server can return an "error" plan indicated by an empty task list (i.e. `{"tasks": {}}`).
		// In this case, we should create a fallback plan.
		if len(cachedPlan.Tasks) == 0 {
			fmt.Println("⚠️ Error plan received, falling back to non-intelligent splitting. Your build may take longer than usual.")
			testPlan := plan.CreateFallbackPlan(files, cfg.Parallelism)
			return testPlan, nil
		}

		debug.Printf("Test plan found. Identifier: %q", cfg.Identifier)
		return *cachedPlan, nil
	}

	debug.Println("No test plan found, creating a new plan")
	// If the cache is empty, create a new plan.
	params, err := createRequestParam(ctx, cfg, files, *apiClient, testRunner)
	if err != nil {
		return handleError(err)
	}

	debug.Println("Creating test plan")
	testPlan, err := apiClient.CreateTestPlan(ctx, cfg.SuiteSlug, params)

	if err != nil {
		return handleError(err)
	}

	// The server can return an "error" plan indicated by an empty task list (i.e. `{"tasks": {}}`).
	// In this case, we should create a fallback plan.
	if len(testPlan.Tasks) == 0 {
		fmt.Println("⚠️ Error plan received, falling back to non-intelligent splitting. Your build may take longer than usual.")
		testPlan = plan.CreateFallbackPlan(files, cfg.Parallelism)
		return testPlan, nil
	}

	debug.Printf("Test plan created. Identifier: %q", cfg.Identifier)
	return testPlan, nil
}

// createRequestParam creates the request parameters for the test plan with the given configuration and files.
// The files should have been filtered by include/exclude patterns before passing to this function.
// If SplitByExample is disabled (default), it will return the default params that contain all the files.
// If SplitByExample is enabled, it will split the slow files into examples and return it along with the rest of the files.
//
// Error is returned if there is a failure to fetch test file timings or to get the test examples from test files when SplitByExample is enabled.
func createRequestParam(ctx context.Context, cfg config.Config, files []string, client api.Client, runner TestRunner) (api.TestPlanParams, error) {
	testFiles := []plan.TestCase{}
	for _, file := range files {
		testFiles = append(testFiles, plan.TestCase{
			Path: file,
		})
	}

	if cfg.SplitByExample {
		debug.Println("Splitting by example")
	}

	debug.Printf("Filtering %d files", len(files))
	filteredFiles, err := client.FilterTests(ctx, cfg.SuiteSlug, api.FilterTestsParams{
		Files: testFiles,
		Env:   cfg.DumpEnv(),
	})

	if err != nil {
		return api.TestPlanParams{}, fmt.Errorf("failed to filter tests: %w", err)
	}

	if len(filteredFiles) == 0 {
		debug.Println("No filtered files found")
		return api.TestPlanParams{
			Identifier:  cfg.Identifier,
			Parallelism: cfg.Parallelism,
			Branch:      cfg.Branch,
			Runner:      cfg.TestRunner,
			Tests: api.TestPlanParamsTest{
				Files: testFiles,
			},
		}, nil
	}

	debug.Printf("Filtered %d files", len(filteredFiles))

	debug.Printf("Getting examples for %d filtered files", len(filteredFiles))

	filteredFilesMap := map[string]bool{}
	filteredFilesPath := []string{}
	for _, file := range filteredFiles {
		filteredFilesMap[file.Path] = true
		filteredFilesPath = append(filteredFilesPath, file.Path)
	}

	examples, err := runner.GetExamples(filteredFilesPath)
	if err != nil {
		return api.TestPlanParams{}, fmt.Errorf("failed to get examples for filtered files: %w", err)
	}

	debug.Printf("Got %d examples within the filtered files", len(examples))

	unfilteredTestFiles := []plan.TestCase{}
	for _, file := range files {
		if _, ok := filteredFilesMap[file]; !ok {
			unfilteredTestFiles = append(unfilteredTestFiles, plan.TestCase{
				Path: file,
			})
		}
	}

	return api.TestPlanParams{
		Identifier:  cfg.Identifier,
		Parallelism: cfg.Parallelism,
		Branch:      cfg.Branch,
		Runner:      cfg.TestRunner,
		Tests: api.TestPlanParamsTest{
			Examples: examples,
			Files:    unfilteredTestFiles,
		},
	}, nil
}