github.com/ava-labs/subnet-evm@v0.6.4/internal/cmdtest/test_cmd.go (about) 1 // (c) 2023, Ava Labs, Inc. 2 // 3 // This file is a derived work, based on the go-ethereum library whose original 4 // notices appear below. 5 // 6 // It is distributed under a license compatible with the licensing terms of the 7 // original code from which it is derived. 8 // 9 // Much love to the original authors for their work. 10 // ********** 11 // Copyright 2017 The go-ethereum Authors 12 // This file is part of the go-ethereum library. 13 // 14 // The go-ethereum library is free software: you can redistribute it and/or modify 15 // it under the terms of the GNU Lesser General Public License as published by 16 // the Free Software Foundation, either version 3 of the License, or 17 // (at your option) any later version. 18 // 19 // The go-ethereum library is distributed in the hope that it will be useful, 20 // but WITHOUT ANY WARRANTY; without even the implied warranty of 21 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 22 // GNU Lesser General Public License for more details. 23 // 24 // You should have received a copy of the GNU Lesser General Public License 25 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>. 26 27 package cmdtest 28 29 import ( 30 "bufio" 31 "bytes" 32 "fmt" 33 "io" 34 "os" 35 "os/exec" 36 "regexp" 37 "strings" 38 "sync" 39 "sync/atomic" 40 "syscall" 41 "testing" 42 "text/template" 43 "time" 44 45 "github.com/docker/docker/pkg/reexec" 46 ) 47 48 func NewTestCmd(t *testing.T, data interface{}) *TestCmd { 49 return &TestCmd{T: t, Data: data} 50 } 51 52 type TestCmd struct { 53 // For total convenience, all testing methods are available. 54 *testing.T 55 56 Func template.FuncMap 57 Data interface{} 58 Cleanup func() 59 60 cmd *exec.Cmd 61 stdout *bufio.Reader 62 stdin io.WriteCloser 63 stderr *testlogger 64 // Err will contain the process exit error or interrupt signal error 65 Err error 66 } 67 68 var id atomic.Int32 69 70 // Run exec's the current binary using name as argv[0] which will trigger the 71 // reexec init function for that name (e.g. "geth-test" in cmd/geth/run_test.go) 72 func (tt *TestCmd) Run(name string, args ...string) { 73 id.Add(1) 74 tt.stderr = &testlogger{t: tt.T, name: fmt.Sprintf("%d", id.Load())} 75 tt.cmd = &exec.Cmd{ 76 Path: reexec.Self(), 77 Args: append([]string{name}, args...), 78 Stderr: tt.stderr, 79 } 80 stdout, err := tt.cmd.StdoutPipe() 81 if err != nil { 82 tt.Fatal(err) 83 } 84 tt.stdout = bufio.NewReader(stdout) 85 if tt.stdin, err = tt.cmd.StdinPipe(); err != nil { 86 tt.Fatal(err) 87 } 88 if err := tt.cmd.Start(); err != nil { 89 tt.Fatal(err) 90 } 91 } 92 93 // InputLine writes the given text to the child's stdin. 94 // This method can also be called from an expect template, e.g.: 95 // 96 // geth.expect(`Passphrase: {{.InputLine "password"}}`) 97 func (tt *TestCmd) InputLine(s string) string { 98 io.WriteString(tt.stdin, s+"\n") 99 return "" 100 } 101 102 func (tt *TestCmd) SetTemplateFunc(name string, fn interface{}) { 103 if tt.Func == nil { 104 tt.Func = make(map[string]interface{}) 105 } 106 tt.Func[name] = fn 107 } 108 109 // Expect runs its argument as a template, then expects the 110 // child process to output the result of the template within 5s. 111 // 112 // If the template starts with a newline, the newline is removed 113 // before matching. 114 func (tt *TestCmd) Expect(tplsource string) { 115 // Generate the expected output by running the template. 116 tpl := template.Must(template.New("").Funcs(tt.Func).Parse(tplsource)) 117 wantbuf := new(bytes.Buffer) 118 if err := tpl.Execute(wantbuf, tt.Data); err != nil { 119 panic(err) 120 } 121 // Trim exactly one newline at the beginning. This makes tests look 122 // much nicer because all expect strings are at column 0. 123 want := bytes.TrimPrefix(wantbuf.Bytes(), []byte("\n")) 124 if err := tt.matchExactOutput(want); err != nil { 125 tt.Fatal(err) 126 } 127 tt.Logf("Matched stdout text:\n%s", want) 128 } 129 130 // Output reads all output from stdout, and returns the data. 131 func (tt *TestCmd) Output() []byte { 132 var buf []byte 133 tt.withKillTimeout(func() { buf, _ = io.ReadAll(tt.stdout) }) 134 return buf 135 } 136 137 func (tt *TestCmd) matchExactOutput(want []byte) error { 138 buf := make([]byte, len(want)) 139 n := 0 140 tt.withKillTimeout(func() { n, _ = io.ReadFull(tt.stdout, buf) }) 141 buf = buf[:n] 142 if n < len(want) || !bytes.Equal(buf, want) { 143 // Grab any additional buffered output in case of mismatch 144 // because it might help with debugging. 145 buf = append(buf, make([]byte, tt.stdout.Buffered())...) 146 tt.stdout.Read(buf[n:]) 147 // Find the mismatch position. 148 for i := 0; i < n; i++ { 149 if want[i] != buf[i] { 150 return fmt.Errorf("output mismatch at ā:\n---------------- (stdout text)\n%sā%s\n---------------- (expected text)\n%s", 151 buf[:i], buf[i:n], want) 152 } 153 } 154 if n < len(want) { 155 return fmt.Errorf("not enough output, got until ā:\n---------------- (stdout text)\n%s\n---------------- (expected text)\n%sā%s", 156 buf, want[:n], want[n:]) 157 } 158 } 159 return nil 160 } 161 162 // ExpectRegexp expects the child process to output text matching the 163 // given regular expression within 5s. 164 // 165 // Note that an arbitrary amount of output may be consumed by the 166 // regular expression. This usually means that expect cannot be used 167 // after ExpectRegexp. 168 func (tt *TestCmd) ExpectRegexp(regex string) (*regexp.Regexp, []string) { 169 regex = strings.TrimPrefix(regex, "\n") 170 var ( 171 re = regexp.MustCompile(regex) 172 rtee = &runeTee{in: tt.stdout} 173 matches []int 174 ) 175 tt.withKillTimeout(func() { matches = re.FindReaderSubmatchIndex(rtee) }) 176 output := rtee.buf.Bytes() 177 if matches == nil { 178 tt.Fatalf("Output did not match:\n---------------- (stdout text)\n%s\n---------------- (regular expression)\n%s", 179 output, regex) 180 return re, nil 181 } 182 tt.Logf("Matched stdout text:\n%s", output) 183 var submatches []string 184 for i := 0; i < len(matches); i += 2 { 185 submatch := string(output[matches[i]:matches[i+1]]) 186 submatches = append(submatches, submatch) 187 } 188 return re, submatches 189 } 190 191 // ExpectExit expects the child process to exit within 5s without 192 // printing any additional text on stdout. 193 func (tt *TestCmd) ExpectExit() { 194 var output []byte 195 tt.withKillTimeout(func() { 196 output, _ = io.ReadAll(tt.stdout) 197 }) 198 tt.WaitExit() 199 if tt.Cleanup != nil { 200 tt.Cleanup() 201 } 202 if len(output) > 0 { 203 tt.Errorf("Unmatched stdout text:\n%s", output) 204 } 205 } 206 207 func (tt *TestCmd) WaitExit() { 208 tt.Err = tt.cmd.Wait() 209 } 210 211 func (tt *TestCmd) Interrupt() { 212 tt.Err = tt.cmd.Process.Signal(os.Interrupt) 213 } 214 215 // ExitStatus exposes the process' OS exit code 216 // It will only return a valid value after the process has finished. 217 func (tt *TestCmd) ExitStatus() int { 218 if tt.Err != nil { 219 exitErr := tt.Err.(*exec.ExitError) 220 if exitErr != nil { 221 if status, ok := exitErr.Sys().(syscall.WaitStatus); ok { 222 return status.ExitStatus() 223 } 224 } 225 } 226 return 0 227 } 228 229 // StderrText returns any stderr output written so far. 230 // The returned text holds all log lines after ExpectExit has 231 // returned. 232 func (tt *TestCmd) StderrText() string { 233 tt.stderr.mu.Lock() 234 defer tt.stderr.mu.Unlock() 235 return tt.stderr.buf.String() 236 } 237 238 func (tt *TestCmd) CloseStdin() { 239 tt.stdin.Close() 240 } 241 242 func (tt *TestCmd) Kill() { 243 tt.cmd.Process.Kill() 244 if tt.Cleanup != nil { 245 tt.Cleanup() 246 } 247 } 248 249 func (tt *TestCmd) withKillTimeout(fn func()) { 250 timeout := time.AfterFunc(30*time.Second, func() { 251 tt.Log("killing the child process (timeout)") 252 tt.Kill() 253 }) 254 defer timeout.Stop() 255 fn() 256 } 257 258 // testlogger logs all written lines via t.Log and also 259 // collects them for later inspection. 260 type testlogger struct { 261 t *testing.T 262 mu sync.Mutex 263 buf bytes.Buffer 264 name string 265 } 266 267 func (tl *testlogger) Write(b []byte) (n int, err error) { 268 lines := bytes.Split(b, []byte("\n")) 269 for _, line := range lines { 270 if len(line) > 0 { 271 tl.t.Logf("(stderr:%v) %s", tl.name, line) 272 } 273 } 274 tl.mu.Lock() 275 tl.buf.Write(b) 276 tl.mu.Unlock() 277 return len(b), err 278 } 279 280 // runeTee collects text read through it into buf. 281 type runeTee struct { 282 in interface { 283 io.Reader 284 io.ByteReader 285 io.RuneReader 286 } 287 buf bytes.Buffer 288 } 289 290 func (rtee *runeTee) Read(b []byte) (n int, err error) { 291 n, err = rtee.in.Read(b) 292 rtee.buf.Write(b[:n]) 293 return n, err 294 } 295 296 func (rtee *runeTee) ReadRune() (r rune, size int, err error) { 297 r, size, err = rtee.in.ReadRune() 298 if err == nil { 299 rtee.buf.WriteRune(r) 300 } 301 return r, size, err 302 } 303 304 func (rtee *runeTee) ReadByte() (b byte, err error) { 305 b, err = rtee.in.ReadByte() 306 if err == nil { 307 rtee.buf.WriteByte(b) 308 } 309 return b, err 310 }