enhance capture functions and add comprehensive tests

Author: umputun

capture.go

@@ -8,25 +8,69 @@ import (
 	"testing"
 )
 
-// CaptureStdout captures the output of a function that writes to stdout.
-// All Capture functions are not thread-safe if used in parallel tests.
-// Usually it is better to pass a custom io.Writer to the function under test instead.
+// CaptureStdout captures os.Stdout output from the provided function.
 func CaptureStdout(t *testing.T, f func()) string {
 	t.Helper()
-	return capture(t, os.Stdout, f)
+	old := os.Stdout
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	os.Stdout = w
+	defer func() { os.Stdout = old }()
+
+	var buf bytes.Buffer
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go func() {
+		defer wg.Done()
+		_, err := io.Copy(&buf, r)
+		if err != nil {
+			t.Errorf("failed to read captured stdout: %v", err)
+		}
+	}()
+
+	f()
+	_ = w.Close()
+	wg.Wait()
+	return buf.String()
 }
 
-// CaptureStderr captures the output of a function that writes to stderr.
+// CaptureStderr captures os.Stderr output from the provided function.
 func CaptureStderr(t *testing.T, f func()) string {
 	t.Helper()
-	return capture(t, os.Stderr, f)
+	old := os.Stderr
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	os.Stderr = w
+	defer func() { os.Stderr = old }()
+
+	var buf bytes.Buffer
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	go func() {
+		defer wg.Done()
+		_, err := io.Copy(&buf, r)
+		if err != nil {
+			t.Errorf("failed to read captured stderr: %v", err)
+		}
+	}()
+
+	f()
+	_ = w.Close()
+	wg.Wait()
+	return buf.String()
 }
 
-// CaptureStdoutAndStderr captures the output of a function that writes to
-// stdout and stderr.
-func CaptureStdoutAndStderr(t *testing.T, f func()) (o, e string) {
+// CaptureStdoutAndStderr captures os.Stdout and os.Stderr output from the provided function.
+func CaptureStdoutAndStderr(t *testing.T, f func()) (stdout, stderr string) {
 	t.Helper()
-
 	oldout, olderr := os.Stdout, os.Stderr
 	rOut, wOut, err := os.Pipe()
 	if err != nil {
@@ -36,65 +80,38 @@ func CaptureStdoutAndStderr(t *testing.T, f func()) (o, e string) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	os.Stdout, os.Stderr = wOut, wErr
+
+	os.Stdout = wOut
+	os.Stderr = wErr
 	defer func() {
-		os.Stdout, os.Stderr = oldout, olderr
+		os.Stdout = oldout
+		os.Stderr = olderr
 	}()
-	outCh, errCh := make(chan string), make(chan string)
 
+	var outBuf, errBuf bytes.Buffer
 	var wg sync.WaitGroup
 	wg.Add(2)
 
-	go func() { //nolint
-		var buf bytes.Buffer
-		wg.Done()
-		if _, err := io.Copy(&buf, rOut); err != nil {
-			t.Fatal(err) //nolint
+	go func() {
+		defer wg.Done()
+		_, err := io.Copy(&outBuf, rOut)
+		if err != nil {
+			t.Errorf("failed to read captured stdout: %v", err)
 		}
-		outCh <- buf.String()
 	}()
 
-	go func() { //nolint
-		var buf bytes.Buffer
-		wg.Done()
-		if _, err := io.Copy(&buf, rErr); err != nil {
-			t.Fatal(err) //nolint
+	go func() {
+		defer wg.Done()
+		_, err := io.Copy(&errBuf, rErr)
+		if err != nil {
+			t.Errorf("failed to read captured stderr: %v", err)
 		}
-		errCh <- buf.String()
 	}()
 
-	wg.Wait()
 	f()
+	_ = wOut.Close()
+	_ = wErr.Close()
+	wg.Wait()
 
-	if err := wOut.Close(); err != nil {
-		t.Fatal(err)
-	}
-	if err := wErr.Close(); err != nil {
-		t.Fatal(err)
-	}
-
-	stdout, stderr := <-outCh, <-errCh
-	return stdout, stderr
-}
-
-func capture(t *testing.T, out *os.File, f func()) string {
-	old := out
-	r, w, err := os.Pipe()
-	if err != nil {
-		t.Fatal(err)
-	}
-	*out = *w
-	defer func() { *out = *old }()
-
-	f()
-
-	_ = w.Close()
-
-	var buf bytes.Buffer
-	_, err = io.Copy(&buf, r)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	return buf.String()
+	return outBuf.String(), errBuf.String()
 }

capture_test.go

@@ -3,40 +3,138 @@ package testutils
 import (
 	"fmt"
 	"os"
+	"strings"
 	"testing"
 )
 
 func TestCaptureStdout(t *testing.T) {
-	want := "hello world\n"
-	got := CaptureStdout(t, func() {
-		fmt.Fprintf(os.Stdout, want)
-	})
-	if want != got {
-		t.Errorf("want %q, got %q", want, got)
+	tests := []struct {
+		name string
+		want string
+		f    func()
+	}{
+		{
+			name: "simple output",
+			want: "hello world\n",
+			f: func() {
+				fmt.Fprintf(os.Stdout, "hello world\n")
+			},
+		},
+		{
+			name: "multiple lines",
+			want: "line1\nline2\n",
+			f: func() {
+				fmt.Fprintln(os.Stdout, "line1")
+				fmt.Fprintln(os.Stdout, "line2")
+			},
+		},
+		{
+			name: "empty output",
+			want: "",
+			f:    func() {},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := CaptureStdout(t, tt.f)
+			if got != tt.want {
+				t.Errorf("CaptureStdout() = %q, want %q", got, tt.want)
+			}
+		})
 	}
 }
 
 func TestCaptureStderr(t *testing.T) {
-	want := "hello world\n"
-	got := CaptureStderr(t, func() {
-		fmt.Fprintf(os.Stderr, want)
-	})
-	if want != got {
-		t.Errorf("want %q, got %q", want, got)
+	tests := []struct {
+		name string
+		want string
+		f    func()
+	}{
+		{
+			name: "simple output",
+			want: "hello world\n",
+			f: func() {
+				fmt.Fprintf(os.Stderr, "hello world\n")
+			},
+		},
+		{
+			name: "multiple lines",
+			want: "line1\nline2\n",
+			f: func() {
+				fmt.Fprintln(os.Stderr, "line1")
+				fmt.Fprintln(os.Stderr, "line2")
+			},
+		},
+		{
+			name: "empty output",
+			want: "",
+			f:    func() {},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := CaptureStderr(t, tt.f)
+			if got != tt.want {
+				t.Errorf("CaptureStderr() = %q, want %q", got, tt.want)
+			}
+		})
 	}
 }
 
 func TestCaptureStdoutAndStderr(t *testing.T) {
-	wantOut := "hello world\n"
-	wantErr := "hello world\n"
-	gotOut, gotErr := CaptureStdoutAndStderr(t, func() {
-		fmt.Fprintf(os.Stdout, wantOut)
-		fmt.Fprintf(os.Stderr, wantErr)
-	})
-	if wantOut != gotOut {
-		t.Errorf("want %q, got %q", wantOut, gotOut)
+	tests := []struct {
+		name    string
+		wantOut string
+		wantErr string
+		f       func()
+	}{
+		{
+			name:    "both outputs",
+			wantOut: "stdout\n",
+			wantErr: "stderr\n",
+			f: func() {
+				fmt.Fprintln(os.Stdout, "stdout")
+				fmt.Fprintln(os.Stderr, "stderr")
+			},
+		},
+		{
+			name:    "only stdout",
+			wantOut: "stdout\n",
+			wantErr: "",
+			f: func() {
+				fmt.Fprintln(os.Stdout, "stdout")
+			},
+		},
+		{
+			name:    "only stderr",
+			wantOut: "",
+			wantErr: "stderr\n",
+			f: func() {
+				fmt.Fprintln(os.Stderr, "stderr")
+			},
+		},
+		{
+			name:    "large output",
+			wantOut: strings.Repeat("a", 100000) + "\n",
+			wantErr: strings.Repeat("b", 100000) + "\n",
+			f: func() {
+				fmt.Fprintln(os.Stdout, strings.Repeat("a", 100000))
+				fmt.Fprintln(os.Stderr, strings.Repeat("b", 100000))
+			},
+		},
 	}
-	if wantErr != gotErr {
-		t.Errorf("want %q, got %q", wantErr, gotErr)
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gotOut, gotErr := CaptureStdoutAndStderr(t, tt.f)
+			if gotOut != tt.wantOut {
+				t.Errorf("CaptureStdoutAndStderr() stdout = %q, want %q", gotOut, tt.wantOut)
+			}
+			if gotErr != tt.wantErr {
+				t.Errorf("CaptureStdoutAndStderr() stderr = %q, want %q", gotErr, tt.wantErr)
+			}
+		})
 	}
 }