-
Notifications
You must be signed in to change notification settings - Fork 0
/
fakeIO.go
179 lines (149 loc) · 4.18 KB
/
fakeIO.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
package testhelper
import (
"bytes"
"errors"
"fmt"
"io"
"os"
"sync"
)
var OrigStdout = os.Stdout
// FakeIO holds the details needed to redirect and restore Stdin, Stdout and
// Stderr for the duration of a test. Create this just before the test with a
// NewStdio... func and after the test is complete restore the original
// settings by calling the Done method which will also return the contents of
// anything written to stdout and stderr.
type FakeIO struct {
sync.Mutex
finished bool
origStdin *os.File
origStdout *os.File
origStderr *os.File
stdoutCh chan []byte
stdoutErrCh chan error
stderrCh chan []byte
stderrErrCh chan error
stdinErrCh chan error
stdinWriter *os.File
stdoutReader *os.File
stderrReader *os.File
}
// closeIfNotNil closes the passed File pointer if it isn't closeIfNotNil
func closeIfNotNil(f *os.File) {
if f != nil {
f.Close()
}
}
func (fio *FakeIO) resetFakeIO() {
closeIfNotNil(os.Stdin)
closeIfNotNil(os.Stdout)
closeIfNotNil(os.Stderr)
os.Stdin = fio.origStdin
os.Stdout = fio.origStdout
os.Stderr = fio.origStderr
closeIfNotNil(fio.stdinWriter)
closeIfNotNil(fio.stdoutReader)
closeIfNotNil(fio.stderrReader)
}
// reader will copy the contents of the passed file into a buffer passing any
// errors back over the errCh and any bytes read over the byteCh.
func reader(name string, r *os.File, byteCh chan []byte, errCh chan error) {
var b bytes.Buffer
if _, err := io.Copy(&b, r); err != nil {
errCh <- fmt.Errorf("Error copying from %s: %w", name, err)
}
byteCh <- b.Bytes()
r.Close()
close(byteCh)
close(errCh)
}
// writer will write the given byte slice to the passed File passing any
// errors back over the errCh.
func writer(name string, w *os.File, b []byte, errCh chan error) {
if _, err := w.Write(b); err != nil {
errCh <- fmt.Errorf("Error writing to %s: %w", name, err)
}
w.Close()
close(errCh)
}
// NewStdioFromString will create a Stdio object which will provide access to
// the contents of anything written to stdout or stderr. After this has been
// called any code reading from stdin will get the contents of the passed
// string. Any output to stdout or stderr will be captured
func NewStdioFromString(input string) (fio *FakeIO, err error) {
fio = &FakeIO{
origStdin: os.Stdin,
origStdout: os.Stdout,
origStderr: os.Stderr,
}
os.Stdin = nil
os.Stdout = nil
os.Stderr = nil
defer func() {
if err != nil {
fio.resetFakeIO()
fio = nil
}
}()
os.Stdin, fio.stdinWriter, err = os.Pipe()
if err != nil {
err = fmt.Errorf("Cannot create the stdin pipe: %w", err)
return
}
fio.stdinErrCh = make(chan error)
go writer("stdin", fio.stdinWriter, []byte(input), fio.stdinErrCh)
fio.stdoutReader, os.Stdout, err = os.Pipe()
if err != nil {
err = fmt.Errorf("Cannot create the stdout pipe: %w", err)
return
}
fio.stdoutCh = make(chan []byte)
fio.stdoutErrCh = make(chan error)
go reader("stdout", fio.stdoutReader, fio.stdoutCh, fio.stdoutErrCh)
fio.stderrReader, os.Stderr, err = os.Pipe()
if err != nil {
err = fmt.Errorf("Cannot create the stderr pipe: %w", err)
return
}
fio.stderrCh = make(chan []byte)
fio.stderrErrCh = make(chan error)
go reader("stderr", fio.stderrReader, fio.stderrCh, fio.stderrErrCh)
return
}
// Done tidies up, restores the std IO values to their previous settings and
// returns anything written to stdout and stderr. It is an error to call this
// twice on the same FakeIO.
func (fio *FakeIO) Done() (stdout, stderr []byte, err error) {
if fio == nil {
err = errors.New("FakeIO.Done - nil pointer")
return
}
fio.Lock()
if fio.finished {
err = errors.New("FakeIO.Done - already called")
}
fio.finished = true
fio.Unlock()
fio.stdinWriter.Close()
os.Stdout.Close()
stdout = <-fio.stdoutCh
if tmpErr, ok := <-fio.stdoutErrCh; ok {
err = tmpErr
}
os.Stderr.Close()
stderr = <-fio.stderrCh
if tmpErr, ok := <-fio.stderrErrCh; ok {
err = errors.Join(err, tmpErr)
}
if tmpErr, ok := <-fio.stdinErrCh; ok && !errors.Is(tmpErr, os.ErrClosed) {
err = errors.Join(err, tmpErr)
}
fio.stdinWriter.Close()
fio.stdoutReader.Close()
fio.stderrReader.Close()
fio.stdinWriter = nil
fio.stdoutReader = nil
fio.stderrReader = nil
fio.resetFakeIO()
return
}