// Copyright 2016 The etcd Authors // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // Package expect implements a small expect-style interface // TODO(ptab): Consider migration to https://github.com/google/goexpect. package expect import ( "bufio" "fmt" "io" "os" "os/exec" "strings" "sync" "syscall" "time" "github.com/creack/pty" ) const DEBUG_LINES_TAIL = 40 type ExpectProcess struct { cmd *exec.Cmd fpty *os.File wg sync.WaitGroup mu sync.Mutex // protects lines and err lines []string count int // increment whenever new line gets added err error // StopSignal is the signal Stop sends to the process; defaults to SIGKILL. StopSignal os.Signal } // NewExpect creates a new process for expect testing. func NewExpect(name string, arg ...string) (ep *ExpectProcess, err error) { // if env[] is nil, use current system env return NewExpectWithEnv(name, arg, nil) } // NewExpectWithEnv creates a new process with user defined env variables for expect testing. func NewExpectWithEnv(name string, args []string, env []string) (ep *ExpectProcess, err error) { cmd := exec.Command(name, args...) cmd.Env = env ep = &ExpectProcess{ cmd: cmd, StopSignal: syscall.SIGKILL, } ep.cmd.Stderr = ep.cmd.Stdout ep.cmd.Stdin = nil if ep.fpty, err = pty.Start(ep.cmd); err != nil { return nil, err } ep.wg.Add(1) go ep.read() return ep, nil } func (ep *ExpectProcess) read() { defer ep.wg.Done() printDebugLines := os.Getenv("EXPECT_DEBUG") != "" r := bufio.NewReader(ep.fpty) for { l, err := r.ReadString('\n') ep.mu.Lock() if l != "" { if printDebugLines { fmt.Printf("%s-%d: %s", ep.cmd.Path, ep.cmd.Process.Pid, l) } ep.lines = append(ep.lines, l) ep.count++ } if err != nil { ep.err = err ep.mu.Unlock() break } ep.mu.Unlock() } } // ExpectFunc returns the first line satisfying the function f. func (ep *ExpectProcess) ExpectFunc(f func(string) bool) (string, error) { i := 0 for { ep.mu.Lock() for i < len(ep.lines) { line := ep.lines[i] i++ if f(line) { ep.mu.Unlock() return line, nil } } if ep.err != nil { ep.mu.Unlock() break } ep.mu.Unlock() time.Sleep(time.Millisecond * 100) } ep.mu.Lock() lastLinesIndex := len(ep.lines) - DEBUG_LINES_TAIL if lastLinesIndex < 0 { lastLinesIndex = 0 } lastLines := strings.Join(ep.lines[lastLinesIndex:], "") ep.mu.Unlock() return "", fmt.Errorf("match not found."+ " Set EXPECT_DEBUG for more info Err: %v, last lines:\n%s", ep.err, lastLines) } // Expect returns the first line containing the given string. func (ep *ExpectProcess) Expect(s string) (string, error) { return ep.ExpectFunc(func(txt string) bool { return strings.Contains(txt, s) }) } // LineCount returns the number of recorded lines since // the beginning of the process. func (ep *ExpectProcess) LineCount() int { ep.mu.Lock() defer ep.mu.Unlock() return ep.count } // Stop kills the expect process and waits for it to exit. func (ep *ExpectProcess) Stop() error { return ep.close(true) } // Signal sends a signal to the expect process func (ep *ExpectProcess) Signal(sig os.Signal) error { return ep.cmd.Process.Signal(sig) } // Wait waits for the process to finish. func (ep *ExpectProcess) Wait() { ep.wg.Wait() } // Close waits for the expect process to exit. // Close currently does not return error if process exited with !=0 status. // TODO: Close should expose underlying proces failure by default. func (ep *ExpectProcess) Close() error { return ep.close(false) } func (ep *ExpectProcess) close(kill bool) error { if ep.cmd == nil { return ep.err } if kill { ep.Signal(ep.StopSignal) } err := ep.cmd.Wait() ep.fpty.Close() ep.wg.Wait() if err != nil { if !kill && strings.Contains(err.Error(), "exit status") { // non-zero exit code err = nil } else if kill && strings.Contains(err.Error(), "signal:") { err = nil } } ep.cmd = nil return err } func (ep *ExpectProcess) Send(command string) error { _, err := io.WriteString(ep.fpty, command) return err } func (ep *ExpectProcess) ProcessError() error { if strings.Contains(ep.err.Error(), "input/output error") { // TODO: The expect library should not return // `/dev/ptmx: input/output error` when process just exits. return nil } return ep.err } func (ep *ExpectProcess) Lines() []string { ep.mu.Lock() defer ep.mu.Unlock() return ep.lines } func (ep *ExpectProcess) IsRunning() bool { return ep.cmd != nil }