From ffba8064edb6c56f0106638026859b9f2e8ce579 Mon Sep 17 00:00:00 2001
From: Martin Atkins <mart@degeneration.co.uk>
Date: Wed, 10 Feb 2021 11:11:50 -0800
Subject: [PATCH] terminal: StreamsForTesting helper

This is to allow convenient testing of functions that are designed to work
directly with *terminal.Streams or the individual stream objects inside.

Because the InputStream and OutputStream APIs expose directly an *os.File,
this does some extra work to set up OS-level pipes so we can capture the
output into local buffers to make test assertions against. The idea here
is to keep the tricky stuff we need for testing confined to the test
codepaths, so that the "real" codepaths don't end up needing to work
around abstractions that are otherwise unnecessary.
---
 internal/terminal/testing.go | 191 +++++++++++++++++++++++++++++++++++
 1 file changed, 191 insertions(+)
 create mode 100644 internal/terminal/testing.go

diff --git a/internal/terminal/testing.go b/internal/terminal/testing.go
new file mode 100644
index 000000000..2830b5d0b
--- /dev/null
+++ b/internal/terminal/testing.go
@@ -0,0 +1,191 @@
+package terminal
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"strings"
+	"sync"
+	"testing"
+)
+
+// StreamsForTesting is a helper for test code that is aiming to test functions
+// that interact with the input and output streams.
+//
+// This particular function is for the simple case of a function that only
+// produces output: the returned input stream is connected to the system's
+// "null device", as if a user had run Terraform with I/O redirection like
+// </dev/null on Unix. It also configures the output as a pipe rather than
+// as a terminal, and so can't be used to test whether code is able to adapt
+// to different terminal widths.
+//
+// The return values are a Streams object ready to pass into a function under
+// test, and a callback function for the test itself to call afterwards
+// in order to obtain any characters that were written to the streams. Once
+// you call the close function, the Streams object becomes invalid and must
+// not be used anymore. Any caller of this function _must_ call close before
+// its test concludes, even if it doesn't intend to check the output, or else
+// it will leak resources.
+//
+// Since this function is for testing only, for convenience it will react to
+// any setup errors by logging a message to the given testing.T object and
+// then failing the test, preventing any later code from running.
+func StreamsForTesting(t *testing.T) (streams *Streams, close func(*testing.T) *TestOutput) {
+	stdinR, err := os.Open(os.DevNull)
+	if err != nil {
+		t.Fatalf("failed to open /dev/null to represent stdin: %s", err)
+	}
+
+	// (Although we only have StreamsForTesting right now, it seems plausible
+	// that we'll want some other similar helpers for more complicated
+	// situations, such as codepaths that need to read from Stdin or
+	// tests for whether a function responds properly to terminal width.
+	// In that case, we'd probably want to factor out the core guts of this
+	// which set up the pipe *os.File values and the goroutines, but then
+	// let each caller produce its own Streams wrapping around those. For
+	// now though, it's simpler to just have this whole implementation together
+	// in one function.)
+
+	// Our idea of streams is only a very thin wrapper around OS-level file
+	// descriptors, so in order to produce a realistic implementation for
+	// the code under test while still allowing us to capture the output
+	// we'll OS-level pipes and concurrently copy anything we read from
+	// them into the output object.
+	outp := &TestOutput{}
+	var lock sync.Mutex // hold while appending to outp
+	stdoutR, stdoutW, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to create stdout pipe: %s", err)
+	}
+	stderrR, stderrW, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to create stderr pipe: %s", err)
+	}
+	var wg sync.WaitGroup // for waiting until our goroutines have exited
+
+	// We need an extra goroutine for each of the pipes so we can block
+	// on reading both of them alongside the caller hopefully writing to
+	// the write sides.
+	wg.Add(2)
+	consume := func(r *os.File, isErr bool) {
+		var buf [1024]byte
+		for {
+			n, err := r.Read(buf[:])
+			if err != nil {
+				if err != io.EOF {
+					// We aren't allowed to write to the testing.T from
+					// a different goroutine than it was created on, but
+					// encountering other errors would be weird here anyway
+					// so we'll just panic. (If we were to just ignore this
+					// and then drop out of the loop then we might deadlock
+					// anyone still trying to write to the write end.)
+					panic(fmt.Sprintf("failed to read from pipe: %s", err))
+				}
+				break
+			}
+			lock.Lock()
+			outp.parts = append(outp.parts, testOutputPart{
+				isErr: isErr,
+				bytes: append(([]byte)(nil), buf[:n]...), // copy so we can reuse the buffer
+			})
+			lock.Unlock()
+		}
+		wg.Done()
+	}
+	go consume(stdoutR, false)
+	go consume(stderrR, true)
+
+	close = func(t *testing.T) *TestOutput {
+		err := stdinR.Close()
+		if err != nil {
+			t.Errorf("failed to close stdin handle: %s", err)
+		}
+
+		// We'll close both of the writer streams now, which should in turn
+		// cause both of the "consume" goroutines above to terminate by
+		// encountering io.EOF.
+		err = stdoutW.Close()
+		if err != nil {
+			t.Errorf("failed to close stdout pipe: %s", err)
+		}
+		err = stderrW.Close()
+		if err != nil {
+			t.Errorf("failed to close stderr pipe: %s", err)
+		}
+
+		// The above error cases still allow this to complete and thus
+		// potentially allow the test to report its own result, but will
+		// ensure that the test doesn't pass while also leaking resources.
+
+		// Wait for the stream-copying goroutines to finish anything they
+		// are working on before we return, or else we might miss some
+		// late-arriving writes.
+		wg.Wait()
+		return outp
+	}
+
+	return &Streams{
+		Stdout: &OutputStream{
+			File: stdoutW,
+		},
+		Stderr: &OutputStream{
+			File: stderrW,
+		},
+		Stdin: &InputStream{
+			File: stdinR,
+		},
+	}, close
+}
+
+// TestOutput is a type used to return the results from the various stream
+// testing helpers. It encapsulates any captured writes to the output and
+// error streams, and has methods to consume that data in some different ways
+// to allow for a few different styles of testing.
+type TestOutput struct {
+	parts []testOutputPart
+}
+
+type testOutputPart struct {
+	// isErr is true if this part was written to the error stream, or false
+	// if it was written to the output stream.
+	isErr bool
+
+	// bytes are the raw bytes that were written
+	bytes []byte
+}
+
+// All returns the output written to both the Stdout and Stderr streams,
+// interleaved together in the order of writing in a single string.
+func (o TestOutput) All() string {
+	buf := &strings.Builder{}
+	for _, part := range o.parts {
+		buf.Write(part.bytes)
+	}
+	return buf.String()
+}
+
+// Stdout returns the output written to just the Stdout stream, ignoring
+// anything that was written to the Stderr stream.
+func (o TestOutput) Stdout() string {
+	buf := &strings.Builder{}
+	for _, part := range o.parts {
+		if part.isErr {
+			continue
+		}
+		buf.Write(part.bytes)
+	}
+	return buf.String()
+}
+
+// Stderr returns the output written to just the Stderr stream, ignoring
+// anything that was written to the Stdout stream.
+func (o TestOutput) Stderr() string {
+	buf := &strings.Builder{}
+	for _, part := range o.parts {
+		if !part.isErr {
+			continue
+		}
+		buf.Write(part.bytes)
+	}
+	return buf.String()
+}