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  }