pytest: refactor testing infrastructure

This CL also exposes a -regen flag to easily regenerate golden files.

Co-authored-by: Drew O'Meara <drew@plan-systems.org>
Signed-off-by: Sebastien Binet <binet@cern.ch>

Author: sbinet

pytest/pytest.go

@@ -6,11 +6,14 @@ package pytest
 
 import (
 	"bytes"
+	"flag"
+	"fmt"
 	"io"
 	"os"
 	"path"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"github.com/go-python/gpython/compile"
@@ -20,6 +23,8 @@ import (
 	_ "github.com/go-python/gpython/stdlib"
 )
 
+var RegenTestData = flag.Bool("regen", false, "Regenerate golden files from current testdata.")
+
 var gContext = py.NewContext(py.DefaultContextOpts())
 
 // Compile the program in the file prog to code in the module that is returned
@@ -132,58 +137,148 @@ func RunBenchmarks(b *testing.B, testDir string) {
 
 // RunScript runs the provided path to a script.
 // RunScript captures the stdout and stderr while executing the script
-// and compares it to a golden file:
+// and compares it to a golden file, blocking until completion.
 //
 //	RunScript("./testdata/foo.py")
 //
 // will compare the output with "./testdata/foo_golden.txt".
 func RunScript(t *testing.T, fname string) {
+
+	RunTestTasks(t, []*Task{
+		{
+			PyFile: fname,
+		},
+	})
+}
+
+// RunTestTasks runs each given task in a newly created py.Context concurrently.
+// If a fatal error is encountered, the given testing.T is signaled.
+func RunTestTasks(t *testing.T, tasks []*Task) {
+	onCompleted := make(chan *Task)
+
+	numTasks := len(tasks)
+	for ti := 0; ti < numTasks; ti++ {
+		task := tasks[ti]
+		go func() {
+			err := task.run()
+			task.Err = err
+			onCompleted <- task
+		}()
+	}
+
+	tasks = tasks[:0]
+	for ti := 0; ti < numTasks; ti++ {
+		task := <-onCompleted
+		if task.Err != nil {
+			t.Error(task.Err)
+		}
+		tasks = append(tasks, task)
+	}
+}
+
+var (
+	taskCounter int32
+)
+
+type Task struct {
+	num      int32                      // Assigned when this task is run
+	ID       string                     // unique key identifying this task.  If empty, autogenerated from the basename of PyFile
+	PyFile   string                     // If set, this file pathname is executed in a newly created ctx
+	PyTask   func(ctx py.Context) error // If set, a new created ctx is created and this blocks until completion
+	GoldFile string                     // Filename containing the "gold standard" stdout+stderr.  If empty, autogenerated from PyFile or ID
+	Err      error                      // Non-nil if a fatal error is encountered with this task
+}
+
+func (task *Task) run() error {
+	fileBase := ""
+
 	opts := py.DefaultContextOpts()
-	opts.SysArgs = []string{fname}
+	if task.PyFile != "" {
+		opts.SysArgs = []string{task.PyFile}
+		if task.ID == "" {
+			ext := filepath.Ext(task.PyFile)
+			fileBase = task.PyFile[0 : len(task.PyFile)-len(ext)]
+		}
+	}
+
+	task.num = atomic.AddInt32(&taskCounter, 1)
+	if task.ID == "" {
+		if fileBase == "" {
+			task.ID = fmt.Sprintf("task-%04d", atomic.AddInt32(&taskCounter, 1))
+		} else {
+			task.ID = strings.TrimPrefix(fileBase, "./")
+		}
+	}
+
+	if task.GoldFile == "" {
+		task.GoldFile = fileBase + "_golden.txt"
+	}
+
 	ctx := py.NewContext(opts)
 	defer ctx.Close()
 
 	sys := ctx.Store().MustGetModule("sys")
 	tmp, err := os.MkdirTemp("", "gpython-pytest-")
 	if err != nil {
-		t.Fatal(err)
+		return err
 	}
 	defer os.RemoveAll(tmp)
 
 	out, err := os.Create(filepath.Join(tmp, "combined"))
 	if err != nil {
-		t.Fatalf("could not create stdout/stderr: %+v", err)
+		return fmt.Errorf("could not create stdout+stderr output file: %w", err)
 	}
 	defer out.Close()
 
 	sys.Globals["stdout"] = &py.File{File: out, FileMode: py.FileWrite}
 	sys.Globals["stderr"] = &py.File{File: out, FileMode: py.FileWrite}
 
-	_, err = py.RunFile(ctx, fname, py.CompileOpts{}, nil)
-	if err != nil {
-		t.Fatalf("could not run script %q: %+v", fname, err)
+	if task.PyFile != "" {
+		_, err := py.RunFile(ctx, task.PyFile, py.CompileOpts{}, nil)
+		if err != nil {
+			return fmt.Errorf("could not run target script %q: %w", task.PyFile, err)
+		}
 	}
 
+	if task.PyTask != nil {
+		err := task.PyTask(ctx)
+		if err != nil {
+			return fmt.Errorf("PyTask %q failed: %w", task.ID, err)
+		}
+	}
+
+	// Close the ctx explicitly as it may legitimately generate output
+	ctx.Close()
+	<-ctx.Done()
+
 	err = out.Close()
 	if err != nil {
-		t.Fatalf("could not close stdout/stderr: %+v", err)
+		return fmt.Errorf("could not close output file: %w", err)
 	}
 
 	got, err := os.ReadFile(out.Name())
 	if err != nil {
-		t.Fatalf("could not read script output: %+v", err)
+		return fmt.Errorf("could not read script output file: %w", err)
 	}
 
-	ref := fname[:len(fname)-len(".py")] + "_golden.txt"
-	want, err := os.ReadFile(ref)
+	if *RegenTestData {
+		err := os.WriteFile(task.GoldFile, got, 0644)
+		if err != nil {
+			return fmt.Errorf("could not write golden output %q: %w", task.GoldFile, err)
+		}
+	}
+
+	want, err := os.ReadFile(task.GoldFile)
 	if err != nil {
-		t.Fatalf("could not read golden output %q: %+v", ref, err)
+		return fmt.Errorf("could not read golden output %q: %w", task.GoldFile, err)
 	}
 
 	diff := cmp.Diff(string(want), string(got))
 	if !bytes.Equal(got, want) {
-		out := fname[:len(fname)-len(".py")] + ".txt"
+		out := fileBase + ".txt"
 		_ = os.WriteFile(out, got, 0644)
-		t.Fatalf("output differ: -- (-ref +got)\n%s", diff)
+		return fmt.Errorf("output differ: -- (-ref +got)\n%s", diff)
 	}
+
+	return nil
 }