expect.go

Documentation: go.etcd.io/etcd/pkg/v3/expect

     1  // Copyright 2016 The etcd Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package expect implements a small expect-style interface
    16  // TODO(ptab): Consider migration to https://github.com/google/goexpect.
    17  package expect
    18  
    19  import (
    20  	"bufio"
    21  	"fmt"
    22  	"io"
    23  	"os"
    24  	"os/exec"
    25  	"strings"
    26  	"sync"
    27  	"syscall"
    28  	"time"
    29  
    30  	"github.com/creack/pty"
    31  )
    32  
    33  const DEBUG_LINES_TAIL = 40
    34  
    35  type ExpectProcess struct {
    36  	cmd  *exec.Cmd
    37  	fpty *os.File
    38  	wg   sync.WaitGroup
    39  
    40  	mu    sync.Mutex // protects lines and err
    41  	lines []string
    42  	count int // increment whenever new line gets added
    43  	err   error
    44  
    45  	// StopSignal is the signal Stop sends to the process; defaults to SIGKILL.
    46  	StopSignal os.Signal
    47  }
    48  
    49  // NewExpect creates a new process for expect testing.
    50  func NewExpect(name string, arg ...string) (ep *ExpectProcess, err error) {
    51  	// if env[] is nil, use current system env
    52  	return NewExpectWithEnv(name, arg, nil)
    53  }
    54  
    55  // NewExpectWithEnv creates a new process with user defined env variables for expect testing.
    56  func NewExpectWithEnv(name string, args []string, env []string) (ep *ExpectProcess, err error) {
    57  	cmd := exec.Command(name, args...)
    58  	cmd.Env = env
    59  	ep = &ExpectProcess{
    60  		cmd:        cmd,
    61  		StopSignal: syscall.SIGKILL,
    62  	}
    63  	ep.cmd.Stderr = ep.cmd.Stdout
    64  	ep.cmd.Stdin = nil
    65  
    66  	if ep.fpty, err = pty.Start(ep.cmd); err != nil {
    67  		return nil, err
    68  	}
    69  
    70  	ep.wg.Add(1)
    71  	go ep.read()
    72  	return ep, nil
    73  }
    74  
    75  func (ep *ExpectProcess) read() {
    76  	defer ep.wg.Done()
    77  	printDebugLines := os.Getenv("EXPECT_DEBUG") != ""
    78  	r := bufio.NewReader(ep.fpty)
    79  	for {
    80  		l, err := r.ReadString('\n')
    81  		ep.mu.Lock()
    82  		if l != "" {
    83  			if printDebugLines {
    84  				fmt.Printf("%s-%d: %s", ep.cmd.Path, ep.cmd.Process.Pid, l)
    85  			}
    86  			ep.lines = append(ep.lines, l)
    87  			ep.count++
    88  		}
    89  		if err != nil {
    90  			ep.err = err
    91  			ep.mu.Unlock()
    92  			break
    93  		}
    94  		ep.mu.Unlock()
    95  	}
    96  }
    97  
    98  // ExpectFunc returns the first line satisfying the function f.
    99  func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) {
   100  	i := 0
   101  
   102  	for {
   103  		ep.mu.Lock()
   104  		for i < len(ep.lines) {
   105  			line := ep.lines[i]
   106  			i++
   107  			if f(line) {
   108  				ep.mu.Unlock()
   109  				return line, nil
   110  			}
   111  		}
   112  		if ep.err != nil {
   113  			ep.mu.Unlock()
   114  			break
   115  		}
   116  		ep.mu.Unlock()
   117  		time.Sleep(time.Millisecond * 100)
   118  	}
   119  	ep.mu.Lock()
   120  	lastLinesIndex := len(ep.lines) - DEBUG_LINES_TAIL
   121  	if lastLinesIndex < 0 {
   122  		lastLinesIndex = 0
   123  	}
   124  	lastLines := strings.Join(ep.lines[lastLinesIndex:], "")
   125  	ep.mu.Unlock()
   126  	return "", fmt.Errorf("match not found."+
   127  		" Set EXPECT_DEBUG for more info Err: %v, last lines:\n%s",
   128  		ep.err, lastLines)
   129  }
   130  
   131  // Expect returns the first line containing the given string.
   132  func (ep *ExpectProcess) Expect(s string) (string, error) {
   133  	return ep.ExpectFunc(func(txt string) bool { return strings.Contains(txt, s) })
   134  }
   135  
   136  // LineCount returns the number of recorded lines since
   137  // the beginning of the process.
   138  func (ep *ExpectProcess) LineCount() int {
   139  	ep.mu.Lock()
   140  	defer ep.mu.Unlock()
   141  	return ep.count
   142  }
   143  
   144  // Stop kills the expect process and waits for it to exit.
   145  func (ep *ExpectProcess) Stop() error { return ep.close(true) }
   146  
   147  // Signal sends a signal to the expect process
   148  func (ep *ExpectProcess) Signal(sig os.Signal) error {
   149  	return ep.cmd.Process.Signal(sig)
   150  }
   151  
   152  // Wait waits for the process to finish.
   153  func (ep *ExpectProcess) Wait() {
   154  	ep.wg.Wait()
   155  }
   156  
   157  // Close waits for the expect process to exit.
   158  // Close currently does not return error if process exited with !=0 status.
   159  // TODO: Close should expose underlying proces failure by default.
   160  func (ep *ExpectProcess) Close() error { return ep.close(false) }
   161  
   162  func (ep *ExpectProcess) close(kill bool) error {
   163  	if ep.cmd == nil {
   164  		return ep.err
   165  	}
   166  	if kill {
   167  		ep.Signal(ep.StopSignal)
   168  	}
   169  
   170  	err := ep.cmd.Wait()
   171  	ep.fpty.Close()
   172  	ep.wg.Wait()
   173  
   174  	if err != nil {
   175  		if !kill && strings.Contains(err.Error(), "exit status") {
   176  			// non-zero exit code
   177  			err = nil
   178  		} else if kill && strings.Contains(err.Error(), "signal:") {
   179  			err = nil
   180  		}
   181  	}
   182  
   183  	ep.cmd = nil
   184  	return err
   185  }
   186  
   187  func (ep *ExpectProcess) Send(command string) error {
   188  	_, err := io.WriteString(ep.fpty, command)
   189  	return err
   190  }
   191  
   192  func (ep *ExpectProcess) ProcessError() error {
   193  	if strings.Contains(ep.err.Error(), "input/output error") {
   194  		// TODO: The expect library should not return
   195  		// `/dev/ptmx: input/output error` when process just exits.
   196  		return nil
   197  	}
   198  	return ep.err
   199  }
   200  
   201  func (ep *ExpectProcess) Lines() []string {
   202  	ep.mu.Lock()
   203  	defer ep.mu.Unlock()
   204  	return ep.lines
   205  }
   206  
   207  func (ep *ExpectProcess) IsRunning() bool {
   208  	return ep.cmd != nil
   209  }
   210
View as plain text