keep test output separate, support -json

main.go

@@ -2,7 +2,10 @@ package main
 
 import (
 	"bufio"
+	"bytes"
+	"encoding/json"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"os/exec"
@@ -13,18 +16,55 @@ import (
 	"syscall"
 )
 
+func usage() {
+	log.Fatalf("Usage:\n\t%s <test-binary> [args...]\n\t%s test <package> [-json]", os.Args[0], os.Args[0])
+}
+
 func main() {
 	if len(os.Args) < 2 {
-		log.Fatalf("Usage: %s <test-binary> [args...]\n", os.Args[0])
+		usage()
 	}
 
-	testBinary := os.Args[1]
-	testArgs := os.Args[2:]
+	var testBinary string
+	var testArgs []string
+	var tests []string
+	var useJson bool = false
 
-	// Get all tests from the binary
-	tests, err := getTestList(testBinary)
-	if err != nil {
-		log.Fatalf("Error getting test list: %v\n", err)
+	if os.Args[1] != "test" {
+		testBinary = os.Args[1]
+		testArgs = os.Args[2:]
+
+		// Get all tests from the binary
+		var err error
+		tests, err = getTestListFromBinary(testBinary)
+		if err != nil {
+			log.Fatalf("Error getting test list: %v\n", err)
+		}
+	} else {
+		if len(os.Args) < 3 {
+			usage()
+		}
+		testPackage := os.Args[2]
+		if len(os.Args) > 3 {
+			if len(os.Args) > 4 || os.Args[3] != "-json" {
+				usage()
+			}
+			useJson = true
+		}
+
+		var err error
+		tests, err = getTestListFromPackage(testPackage)
+		if err != nil {
+			log.Fatalf("Error getting test list: %v\n", err)
+		}
+
+		testBinary = "go"
+		testArgs = []string{"test"}
+		if useJson {
+			testArgs = append(testArgs, "-json")
+		} else {
+			testArgs = append(testArgs, "-v")
+		}
 	}
 
 	// Check if we should parallelize (>20 tests).
@@ -38,12 +78,20 @@ func main() {
 	n := runtime.GOMAXPROCS(0)
 	testGroups := divideTests(tests, n)
 
-	runTestsInParallel(testBinary, testArgs, testGroups)
+	runTestsInParallel(testBinary, testArgs, testGroups, useJson)
+}
+
+// getTestListFromPackage gets all test function names using a 'go test' call
+func getTestListFromPackage(testPackage string) ([]string, error) {
+	return getTestListFromCommand(exec.Command("go", "test", "-list", ".*", testPackage))
+}
+
+// getTestListFromBinary gets all test function names from the test binary
+func getTestListFromBinary(testBinary string) ([]string, error) {
+	return getTestListFromCommand(exec.Command(testBinary, "-test.list", ".*"))
 }
 
-// getTestList gets all test function names from the test binary
-func getTestList(testBinary string) ([]string, error) {
-	cmd := exec.Command(testBinary, "-test.list", ".*")
+func getTestListFromCommand(cmd *exec.Cmd) ([]string, error) {
 	output, err := cmd.Output()
 	if err != nil {
 		return nil, fmt.Errorf("failed to list tests: %v", err)
@@ -100,7 +148,7 @@ func divideTests(tests []string, n int) [][]string {
 }
 
 // runTestsInParallel runs test groups in parallel
-func runTestsInParallel(testBinary string, testArgs []string, testGroups [][]string) {
+func runTestsInParallel(testBinary string, testArgs []string, testGroups [][]string, useJson bool) {
 	var wg sync.WaitGroup
 	var exitCode atomic.Int64
 
@@ -113,13 +161,26 @@ func runTestsInParallel(testBinary string, testArgs []string, testGroups [][]str
 			// Create regex pattern for this group
 			pattern := "^(" + strings.Join(group, "|") + ")$"
 
-			args := append([]string{"-test.run", pattern}, testArgs...)
+			args := append([]string{}, testArgs...)
+			args = append(args, "-test.run", pattern)
 
 			cmd := exec.Command(testBinary, args...)
-			cmd.Stdout = os.Stdout
-			cmd.Stderr = os.Stderr
+			pr, pw, err := os.Pipe()
+			if err != nil {
+				panic(fmt.Errorf("can't pipe: %v", err))
+			}
+			cmd.StdoutPipe()
+			cmd.Stdout = pw
+			cmd.Stderr = pw
+			wg.Go(func() { reader(pr, useJson) })
+
+			err = cmd.Start()
+			pw.Close()
+			if err == nil {
+				err = cmd.Wait()
+			}
 
-			if err := cmd.Run(); err != nil {
+			if err != nil {
 				exitCode.Store(1)
 
 				if exitErr, ok := err.(*exec.ExitError); ok {
@@ -155,3 +216,46 @@ func runTestBinary(testBinary string, args []string) {
 		os.Exit(1)
 	}
 }
+
+type testEvent struct {
+	Action string
+	Test   string
+}
+
+var mu sync.Mutex
+
+func reader(in io.ReadCloser, useJson bool) {
+	rdr := bufio.NewScanner(in)
+	buf := new(bytes.Buffer)
+
+	flush := func() {
+		mu.Lock()
+		io.Copy(os.Stdout, buf)
+		mu.Unlock()
+		buf.Reset()
+	}
+
+	for rdr.Scan() {
+		text := rdr.Text()
+		// We flush the output buffer when we se the start of a new test, this is
+		// easier than doing it at the *end* of a test because tests that have
+		// sub-tests will print a 'PASS' line for the main test, followed by a
+		// 'PASS' line for each sub-test.
+		if useJson {
+			var ev testEvent
+			err := json.Unmarshal([]byte(text), &ev)
+			if err == nil {
+				if ev.Action == "run" && strings.Index(ev.Test, "/") < 0 {
+					flush()
+				}
+			}
+		} else {
+			if strings.HasPrefix(text, "=== RUN ") {
+				flush()
+			}
+		}
+		buf.WriteString(text)
+		buf.WriteByte('\n')
+	}
+	flush()
+}