blob: 11b907fd0bca23f799fe1e9079d62dc6ae3500db [file] [log] [blame]
// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
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()
}