testing: add Output

Output is a method on T, B and F. It provides an io.Writer that writes
to the same test output stream as TB.Log. The new output writer is
used to refactor the implementation of Log. It maintains the formatting
provided by Log while making call site information optional.
Additionally, it provides buffering of log messages.

Co-authored-by: Aleks Fazlieva <britishrum@users.noreply.github.com>

Fixes #59928.

Change-Id: I29090b3d4f61f7334388b373ec18750d5637aafa
GitHub-Last-Rev: 18af0e15262494f2074d545a6042b079d62301a2
GitHub-Pull-Request: golang/go#71575
Reviewed-on: https://go-review.googlesource.com/c/go/+/646956
Reviewed-by: Arati <artichaut2023@gmail.com>
Auto-Submit: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Jonathan Amsterdam <jba@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Michael Knyszek <mknyszek@google.com>

1 files changed, 150 insertions, 64 deletions

diff --git a/src/testing/testing.go b/src/testing/testing.go
index d50abea32f..efbcd59dc0 100644
--- a/src/testing/testing.go
+++ b/src/testing/testing.go
@@ -631,6 +631,7 @@ type common struct {
 	mu          sync.RWMutex         // guards this group of fields
 	output      []byte               // Output generated by test or benchmark.
 	w           io.Writer            // For flushToParent.
+	o           *outputWriter        // Writes output.
 	ran         bool                 // Test or benchmark (or one of its subtests) was executed.
 	failed      bool                 // Test or benchmark has failed.
 	skipped     bool                 // Test or benchmark has been skipped.
@@ -799,44 +800,6 @@ func (c *common) frameSkip(skip int) runtime.Frame {
 	return firstFrame
 }
 
-// decorate prefixes the string with the file and line of the call site
-// and inserts the final newline if needed and indentation spaces for formatting.
-// This function must be called with c.mu held.
-func (c *common) decorate(s string, skip int) string {
-	frame := c.frameSkip(skip)
-	file := frame.File
-	line := frame.Line
-	if file != "" {
-		if *fullPath {
-			// If relative path, truncate file name at last file name separator.
-		} else if index := strings.LastIndexAny(file, `/\`); index >= 0 {
-			file = file[index+1:]
-		}
-	} else {
-		file = "???"
-	}
-	if line == 0 {
-		line = 1
-	}
-	buf := new(strings.Builder)
-	// Every line is indented at least 4 spaces.
-	buf.WriteString("    ")
-	fmt.Fprintf(buf, "%s:%d: ", file, line)
-	lines := strings.Split(s, "\n")
-	if l := len(lines); l > 1 && lines[l-1] == "" {
-		lines = lines[:l-1]
-	}
-	for i, line := range lines {
-		if i > 0 {
-			// Second and subsequent lines are indented an additional 4 spaces.
-			buf.WriteString("\n        ")
-		}
-		buf.WriteString(line)
-	}
-	buf.WriteByte('\n')
-	return buf.String()
-}
-
 // flushToParent writes c.output to the parent after first writing the header
 // with the given format and arguments.
 func (c *common) flushToParent(testName, format string, args ...any) {
@@ -882,6 +845,8 @@ type indenter struct {
 	c *common
 }
 
+const indent = "    "
+
 func (w indenter) Write(b []byte) (n int, err error) {
 	n = len(b)
 	for len(b) > 0 {
@@ -898,7 +863,6 @@ func (w indenter) Write(b []byte) (n int, err error) {
 			w.c.output = append(w.c.output, marker)
 			line = line[1:]
 		}
-		const indent = "    "
 		w.c.output = append(w.c.output, indent...)
 		w.c.output = append(w.c.output, line...)
 		b = b[end:]
@@ -1043,41 +1007,156 @@ func (c *common) FailNow() {
 	runtime.Goexit()
 }
 
-// log generates the output. It's always at the same stack depth.
+// log generates the output. It is always at the same stack depth. log inserts
+// indentation and the final newline if necessary. It prefixes the string
+// with the file and line of the call site.
 func (c *common) log(s string) {
-	c.logDepth(s, 3) // logDepth + log + public function
+	s = strings.TrimSuffix(s, "\n")
+
+	// Second and subsequent lines are indented 4 spaces. This is in addition to
+	// the indentation provided by outputWriter.
+	s = strings.ReplaceAll(s, "\n", "\n"+indent)
+	s += "\n"
+
+	n := c.destination()
+	if n == nil {
+		// The test and all its parents are done. The log cannot be output.
+		panic("Log in goroutine after " + c.name + " has completed: " + s)
+	}
+
+	// Prefix with the call site. It is located by skipping 3 functions:
+	// callSite + log + public function
+	s = n.callSite(3) + s
+
+	// Output buffered logs.
+	n.flushPartial()
+
+	n.o.Write([]byte(s))
 }
 
-// logDepth generates the output at an arbitrary stack depth.
-func (c *common) logDepth(s string, depth int) {
+// destination selects the test to which output should be appended. It returns the
+// test if it is incomplete. Otherwise, it finds its closest incomplete parent.
+func (c *common) destination() *common {
 	c.mu.Lock()
 	defer c.mu.Unlock()
-	if c.done {
-		// This test has already finished. Try and log this message
-		// with our parent. If we don't have a parent, panic.
-		for parent := c.parent; parent != nil; parent = parent.parent {
-			parent.mu.Lock()
-			defer parent.mu.Unlock()
-			if !parent.done {
-				parent.output = append(parent.output, parent.decorate(s, depth+1)...)
-				return
-			}
+
+	if !c.done {
+		return c
+	}
+	for parent := c.parent; parent != nil; parent = parent.parent {
+		parent.mu.Lock()
+		defer parent.mu.Unlock()
+		if !parent.done {
+			return parent
+		}
+	}
+	return nil
+}
+
+// callSite retrieves and formats the file and line of the call site.
+func (c *common) callSite(skip int) string {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	frame := c.frameSkip(skip)
+	file := frame.File
+	line := frame.Line
+	if file != "" {
+		if *fullPath {
+			// If relative path, truncate file name at last file name separator.
+		} else {
+			file = filepath.Base(file)
 		}
-		panic("Log in goroutine after " + c.name + " has completed: " + s)
 	} else {
-		if c.chatty != nil {
-			if c.bench {
-				// Benchmarks don't print === CONT, so we should skip the test
-				// printer and just print straight to stdout.
-				fmt.Print(c.decorate(s, depth+1))
-			} else {
-				c.chatty.Printf(c.name, "%s", c.decorate(s, depth+1))
-			}
+		file = "???"
+	}
+	if line == 0 {
+		line = 1
+	}
 
-			return
+	return fmt.Sprintf("%s:%d: ", file, line)
+}
+
+// flushPartial checks the buffer for partial logs and outputs them.
+func (c *common) flushPartial() {
+	partial := func() bool {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		return (c.o != nil) && (len(c.o.partial) > 0)
+	}
+
+	if partial() {
+		c.o.Write([]byte("\n"))
+	}
+}
+
+// Output returns a Writer that writes to the same test output stream as TB.Log.
+// The output is indented like TB.Log lines, but Output does not
+// add source locations or newlines. The output is internally line
+// buffered, and a call to TB.Log or the end of the test will implicitly
+// flush the buffer, followed by a newline. After a test function and all its
+// parents return, neither Output nor the Write method may be called.
+func (c *common) Output() io.Writer {
+	c.checkFuzzFn("Output")
+	n := c.destination()
+	if n == nil {
+		panic("Output called after " + c.name + " has completed")
+	}
+	return n.o
+}
+
+// setOutputWriter initializes an outputWriter and sets it as a common field.
+func (c *common) setOutputWriter() {
+	c.o = &outputWriter{c: c}
+}
+
+// outputWriter buffers, formats and writes log messages.
+type outputWriter struct {
+	c       *common
+	partial []byte // incomplete ('\n'-free) suffix of last Write
+}
+
+// Write writes a log message to the test's output stream, properly formatted and
+// indented. It may not be called after a test function and all its parents return.
+func (o *outputWriter) Write(p []byte) (int, error) {
+	if o.c.destination() == nil {
+		panic("Write called after " + o.c.name + " has completed")
+	}
+
+	o.c.mu.Lock()
+	defer o.c.mu.Unlock()
+
+	// The last element is a partial line.
+	lines := bytes.SplitAfter(p, []byte("\n"))
+	last := len(lines) - 1 // Inv: 0 <= last
+	for i, line := range lines[:last] {
+		// Emit partial line from previous call.
+		if i == 0 && len(o.partial) > 0 {
+			line = slices.Concat(o.partial, line)
+			o.partial = o.partial[:0]
+		}
+		o.writeLine(line)
+	}
+	// Save partial line for next call.
+	o.partial = append(o.partial, lines[last]...)
+
+	return len(p), nil
+}
+
+// writeLine generates the output for a given line.
+func (o *outputWriter) writeLine(b []byte) {
+	if !o.c.done && (o.c.chatty != nil) {
+		if o.c.bench {
+			// Benchmarks don't print === CONT, so we should skip the test
+			// printer and just print straight to stdout.
+			fmt.Printf("%s%s", indent, b)
+		} else {
+			o.c.chatty.Printf(o.c.name, "%s%s", indent, b)
 		}
-		c.output = append(c.output, c.decorate(s, depth+1)...)
+		return
 	}
+	o.c.output = append(o.c.output, indent...)
+	o.c.output = append(o.c.output, b...)
 }
 
 // Log formats its arguments using default formatting, analogous to [fmt.Println],
@@ -1739,6 +1818,8 @@ func tRunner(t *T, fn func(t *T)) {
 				root.duration += highPrecisionTimeSince(root.start)
 				d := root.duration
 				root.mu.Unlock()
+				// Output buffered logs.
+				root.flushPartial()
 				root.flushToParent(root.name, "--- FAIL: %s (%s)\n", root.name, fmtDuration(d))
 				if r := root.parent.runCleanup(recoverAndReturnPanic); r != nil {
 					fmt.Fprintf(root.parent.w, "cleanup panicked with %v", r)
@@ -1786,6 +1867,10 @@ func tRunner(t *T, fn func(t *T)) {
 			// test. See comment in Run method.
 			t.tstate.release()
 		}
+		// Output buffered logs.
+		for root := &t.common; root.parent != nil; root = root.parent {
+			root.flushPartial()
+		}
 		t.report() // Report after all subtests have finished.
 
 		// Do not lock t.done to allow race detector to detect race in case
@@ -1851,6 +1936,7 @@ func (t *T) Run(name string, f func(t *T)) bool {
 		tstate: t.tstate,
 	}
 	t.w = indenter{&t.common}
+	t.setOutputWriter()
 
 	if t.chatty != nil {
 		t.chatty.Updatef(t.name, "=== RUN   %s\n", t.name)