From e41cf0e429135697079808f2c6c42fb8d66d8f32 Mon Sep 17 00:00:00 2001
From: Michael Vogt <michael.vogt@gmail.com>
Date: Wed, 29 Jan 2025 12:48:53 +0100
Subject: [PATCH] testutil: add new CaptureStdio helper

This commit adds a new testutil.CaptureStdio helper so that we
can test external go modules that use os.Std{out,err} but now
allow mocking or overwriting.
---
 internal/testutil/testutil.go      | 50 ++++++++++++++++++++++++++++++
 internal/testutil/testutil_test.go | 11 +++++++
 2 files changed, 61 insertions(+)

diff --git a/internal/testutil/testutil.go b/internal/testutil/testutil.go
index 284703c..92147d5 100644
--- a/internal/testutil/testutil.go
+++ b/internal/testutil/testutil.go
@@ -1,9 +1,12 @@
 package testutil
 
 import (
+	"bytes"
+	"io"
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -60,3 +63,50 @@ func (mc *MockCmd) Calls() [][]string {
 	}
 	return calls
 }
+
+// CaptureStdio runs the given function f() in an environment that
+// captures stdout, stderr and returns the the result as string.
+//
+// Use this very targeted to avoid real stdout/stderr output
+// from being displayed.
+func CaptureStdio(t *testing.T, f func()) (string, string) {
+	saved1 := os.Stdout
+	saved2 := os.Stderr
+
+	r1, w1, err := os.Pipe()
+	require.NoError(t, err)
+	defer r1.Close()
+	defer w1.Close()
+	r2, w2, err := os.Pipe()
+	require.NoError(t, err)
+	defer r2.Close()
+	defer w2.Close()
+
+	var wg sync.WaitGroup
+	var stdout, stderr bytes.Buffer
+	wg.Add(1)
+	// this needs to be a go-routines or we could deadlock
+	// when the pipe is full
+	go func() {
+		defer wg.Done()
+		io.Copy(&stdout, r1)
+	}()
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		io.Copy(&stderr, r2)
+	}()
+
+	os.Stdout = w1
+	os.Stderr = w2
+	defer func() {
+		os.Stdout = saved1
+		os.Stderr = saved2
+	}()
+
+	f()
+	w1.Close()
+	w2.Close()
+	wg.Wait()
+	return stdout.String(), stderr.String()
+}
diff --git a/internal/testutil/testutil_test.go b/internal/testutil/testutil_test.go
index 61d8192..0162281 100644
--- a/internal/testutil/testutil_test.go
+++ b/internal/testutil/testutil_test.go
@@ -1,6 +1,8 @@
 package testutil_test
 
 import (
+	"fmt"
+	"os"
 	"os/exec"
 	"testing"
 
@@ -23,3 +25,12 @@ func TestMockCommand(t *testing.T) {
 		{"run2-arg1", "run2-arg2"},
 	}, fakeCmd.Calls())
 }
+
+func TestCaptureStdout(t *testing.T) {
+	stdout, stderr := testutil.CaptureStdio(t, func() {
+		fmt.Fprintf(os.Stdout, "output on stdout")
+		fmt.Fprintf(os.Stderr, "output on stderr")
+	})
+	assert.Equal(t, "output on stdout", stdout)
+	assert.Equal(t, "output on stderr", stderr)
+}