1package terminal
2
3import (
4	"fmt"
5	"io"
6	"os"
7	"strings"
8	"sync"
9	"testing"
10)
11
12// StreamsForTesting is a helper for test code that is aiming to test functions
13// that interact with the input and output streams.
14//
15// This particular function is for the simple case of a function that only
16// produces output: the returned input stream is connected to the system's
17// "null device", as if a user had run Terraform with I/O redirection like
18// </dev/null on Unix. It also configures the output as a pipe rather than
19// as a terminal, and so can't be used to test whether code is able to adapt
20// to different terminal widths.
21//
22// The return values are a Streams object ready to pass into a function under
23// test, and a callback function for the test itself to call afterwards
24// in order to obtain any characters that were written to the streams. Once
25// you call the close function, the Streams object becomes invalid and must
26// not be used anymore. Any caller of this function _must_ call close before
27// its test concludes, even if it doesn't intend to check the output, or else
28// it will leak resources.
29//
30// Since this function is for testing only, for convenience it will react to
31// any setup errors by logging a message to the given testing.T object and
32// then failing the test, preventing any later code from running.
33func StreamsForTesting(t *testing.T) (streams *Streams, close func(*testing.T) *TestOutput) {
34	stdinR, err := os.Open(os.DevNull)
35	if err != nil {
36		t.Fatalf("failed to open /dev/null to represent stdin: %s", err)
37	}
38
39	// (Although we only have StreamsForTesting right now, it seems plausible
40	// that we'll want some other similar helpers for more complicated
41	// situations, such as codepaths that need to read from Stdin or
42	// tests for whether a function responds properly to terminal width.
43	// In that case, we'd probably want to factor out the core guts of this
44	// which set up the pipe *os.File values and the goroutines, but then
45	// let each caller produce its own Streams wrapping around those. For
46	// now though, it's simpler to just have this whole implementation together
47	// in one function.)
48
49	// Our idea of streams is only a very thin wrapper around OS-level file
50	// descriptors, so in order to produce a realistic implementation for
51	// the code under test while still allowing us to capture the output
52	// we'll OS-level pipes and concurrently copy anything we read from
53	// them into the output object.
54	outp := &TestOutput{}
55	var lock sync.Mutex // hold while appending to outp
56	stdoutR, stdoutW, err := os.Pipe()
57	if err != nil {
58		t.Fatalf("failed to create stdout pipe: %s", err)
59	}
60	stderrR, stderrW, err := os.Pipe()
61	if err != nil {
62		t.Fatalf("failed to create stderr pipe: %s", err)
63	}
64	var wg sync.WaitGroup // for waiting until our goroutines have exited
65
66	// We need an extra goroutine for each of the pipes so we can block
67	// on reading both of them alongside the caller hopefully writing to
68	// the write sides.
69	wg.Add(2)
70	consume := func(r *os.File, isErr bool) {
71		var buf [1024]byte
72		for {
73			n, err := r.Read(buf[:])
74			if err != nil {
75				if err != io.EOF {
76					// We aren't allowed to write to the testing.T from
77					// a different goroutine than it was created on, but
78					// encountering other errors would be weird here anyway
79					// so we'll just panic. (If we were to just ignore this
80					// and then drop out of the loop then we might deadlock
81					// anyone still trying to write to the write end.)
82					panic(fmt.Sprintf("failed to read from pipe: %s", err))
83				}
84				break
85			}
86			lock.Lock()
87			outp.parts = append(outp.parts, testOutputPart{
88				isErr: isErr,
89				bytes: append(([]byte)(nil), buf[:n]...), // copy so we can reuse the buffer
90			})
91			lock.Unlock()
92		}
93		wg.Done()
94	}
95	go consume(stdoutR, false)
96	go consume(stderrR, true)
97
98	close = func(t *testing.T) *TestOutput {
99		err := stdinR.Close()
100		if err != nil {
101			t.Errorf("failed to close stdin handle: %s", err)
102		}
103
104		// We'll close both of the writer streams now, which should in turn
105		// cause both of the "consume" goroutines above to terminate by
106		// encountering io.EOF.
107		err = stdoutW.Close()
108		if err != nil {
109			t.Errorf("failed to close stdout pipe: %s", err)
110		}
111		err = stderrW.Close()
112		if err != nil {
113			t.Errorf("failed to close stderr pipe: %s", err)
114		}
115
116		// The above error cases still allow this to complete and thus
117		// potentially allow the test to report its own result, but will
118		// ensure that the test doesn't pass while also leaking resources.
119
120		// Wait for the stream-copying goroutines to finish anything they
121		// are working on before we return, or else we might miss some
122		// late-arriving writes.
123		wg.Wait()
124		return outp
125	}
126
127	return &Streams{
128		Stdout: &OutputStream{
129			File: stdoutW,
130		},
131		Stderr: &OutputStream{
132			File: stderrW,
133		},
134		Stdin: &InputStream{
135			File: stdinR,
136		},
137	}, close
138}
139
140// TestOutput is a type used to return the results from the various stream
141// testing helpers. It encapsulates any captured writes to the output and
142// error streams, and has methods to consume that data in some different ways
143// to allow for a few different styles of testing.
144type TestOutput struct {
145	parts []testOutputPart
146}
147
148type testOutputPart struct {
149	// isErr is true if this part was written to the error stream, or false
150	// if it was written to the output stream.
151	isErr bool
152
153	// bytes are the raw bytes that were written
154	bytes []byte
155}
156
157// All returns the output written to both the Stdout and Stderr streams,
158// interleaved together in the order of writing in a single string.
159func (o TestOutput) All() string {
160	buf := &strings.Builder{}
161	for _, part := range o.parts {
162		buf.Write(part.bytes)
163	}
164	return buf.String()
165}
166
167// Stdout returns the output written to just the Stdout stream, ignoring
168// anything that was written to the Stderr stream.
169func (o TestOutput) Stdout() string {
170	buf := &strings.Builder{}
171	for _, part := range o.parts {
172		if part.isErr {
173			continue
174		}
175		buf.Write(part.bytes)
176	}
177	return buf.String()
178}
179
180// Stderr returns the output written to just the Stderr stream, ignoring
181// anything that was written to the Stdout stream.
182func (o TestOutput) Stderr() string {
183	buf := &strings.Builder{}
184	for _, part := range o.parts {
185		if !part.isErr {
186			continue
187		}
188		buf.Write(part.bytes)
189	}
190	return buf.String()
191}
192