Expect exit code enhancement

ExpectProcess and ExpectFunc now take the exit code of the process into
account, not just the matching of the tty output.

This also refactors the many tests that were previously succeeding on
matching an output from a failing cmd execution.

Signed-off-by: Thomas Jungblut <tjungblu@redhat.com>

pkg/expect/expect.go

@@ -19,6 +19,7 @@ package expect
 import (
 	"bufio"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -33,18 +34,23 @@ import (
 
 const DEBUG_LINES_TAIL = 40
 
+var (
+	ErrProcessRunning = fmt.Errorf("process is still running")
+)
+
 type ExpectProcess struct {
 	cfg expectConfig
 
 	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
-	cur   int // current read position
-	err   error
+	mu       sync.Mutex // protects lines, count, cur, exitErr and exitCode
+	lines    []string
+	count    int   // increment whenever new line gets added
+	cur      int   // current read position
+	exitErr  error // process exit error
+	exitCode int
 }
 
 // NewExpect creates a new process for expect testing.
@@ -69,8 +75,9 @@ func NewExpectWithEnv(name string, args []string, env []string, serverProcessCon
 		return nil, err
 	}
 
-	ep.wg.Add(1)
+	ep.wg.Add(2)
 	go ep.read()
+	go ep.waitSaveExitErr()
 	return ep, nil
 }
 
@@ -95,46 +102,87 @@ func (ep *ExpectProcess) Pid() int {
 
 func (ep *ExpectProcess) read() {
 	defer ep.wg.Done()
-	printDebugLines := os.Getenv("EXPECT_DEBUG") != ""
+	defer func(fpty *os.File) {
+		err := fpty.Close()
+		if err != nil {
+			// we deliberately only log the error here, closing the PTY should mostly be (expected) broken pipes
+			fmt.Printf("error while closing fpty: %v", err)
+		}
+	}(ep.fpty)
+
 	r := bufio.NewReader(ep.fpty)
 	for {
-		l, err := r.ReadString('\n')
-		ep.mu.Lock()
-		if l != "" {
-			if printDebugLines {
-				fmt.Printf("%s (%s) (%d): %s", ep.cmd.Path, ep.cfg.name, ep.cmd.Process.Pid, l)
-			}
-			ep.lines = append(ep.lines, l)
-			ep.count++
-		}
+		err := ep.tryReadNextLine(r)
 		if err != nil {
-			ep.err = err
-			ep.mu.Unlock()
 			break
 		}
-		ep.mu.Unlock()
+	}
+}
+
+func (ep *ExpectProcess) tryReadNextLine(r *bufio.Reader) error {
+	printDebugLines := os.Getenv("EXPECT_DEBUG") != ""
+	l, err := r.ReadString('\n')
+
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
+
+	if l != "" {
+		if printDebugLines {
+			fmt.Printf("%s (%s) (%d): %s", ep.cmd.Path, ep.cfg.name, ep.cmd.Process.Pid, l)
+		}
+		ep.lines = append(ep.lines, l)
+		ep.count++
+	}
+
+	// we're checking the error here at the bottom to ensure any leftover reads are still taken into account
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (ep *ExpectProcess) waitSaveExitErr() {
+	defer ep.wg.Done()
+	err := ep.waitProcess()
+
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
+	if err != nil {
+		ep.exitErr = err
 	}
 }
 
 // ExpectFunc returns the first line satisfying the function f.
 func (ep *ExpectProcess) ExpectFunc(ctx context.Context, 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
+		line, errsFound := func() (string, bool) {
+			ep.mu.Lock()
+			defer ep.mu.Unlock()
+
+			// check if this expect has been already closed
+			if ep.cmd == nil {
+				return "", true
+			}
+
+			for i < len(ep.lines) {
+				line := ep.lines[i]
+				i++
+				if f(line) {
+					return line, false
+				}
 			}
+			return "", ep.exitErr != nil
+		}()
+
+		if line != "" {
+			return line, nil
 		}
-		if ep.err != nil {
-			ep.mu.Unlock()
+
+		if errsFound {
 			break
 		}
-		ep.mu.Unlock()
 
 		select {
 		case <-ctx.Done():
@@ -143,16 +191,18 @@ func (ep *ExpectProcess) ExpectFunc(ctx context.Context, f func(string) bool) (s
 			// continue loop
 		}
 	}
+
 	ep.mu.Lock()
+	defer ep.mu.Unlock()
+
 	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)
+	return "", fmt.Errorf("match not found. "+
+		" Set EXPECT_DEBUG for more info Errs: [%v], last lines:\n%s",
+		ep.exitErr, lastLines)
 }
 
 // ExpectWithContext returns the first line containing the given string.
@@ -174,63 +224,92 @@ func (ep *ExpectProcess) LineCount() int {
 	return ep.count
 }
 
-// Stop kills the expect process and waits for it to exit.
-func (ep *ExpectProcess) Stop() error { return ep.close(true) }
+// ExitCode returns the exit code of this process.
+// If the process is still running, it returns exit code 0 and ErrProcessRunning.
+func (ep *ExpectProcess) ExitCode() (int, error) {
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
 
-// Signal sends a signal to the expect process
-func (ep *ExpectProcess) Signal(sig os.Signal) error {
-	return ep.cmd.Process.Signal(sig)
+	if ep.cmd == nil {
+		return ep.exitCode, nil
+	}
+
+	return 0, ErrProcessRunning
+}
+
+// ExitError returns the exit error of this process (if any).
+// If the process is still running, it returns ErrProcessRunning instead.
+func (ep *ExpectProcess) ExitError() error {
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
+
+	if ep.cmd == nil {
+		return ep.exitErr
+	}
+
+	return ErrProcessRunning
 }
 
-func (ep *ExpectProcess) Wait() error {
-	_, err := ep.cmd.Process.Wait()
+// Stop signals the process to terminate via SIGTERM
+func (ep *ExpectProcess) Stop() error {
+	err := ep.Signal(syscall.SIGTERM)
+	if err != nil && strings.Contains(err.Error(), "os: process already finished") {
+		return nil
+	}
 	return err
 }
 
-// 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 process failure by default.
-func (ep *ExpectProcess) Close() error { return ep.close(false) }
+// Signal sends a signal to the expect process
+func (ep *ExpectProcess) Signal(sig os.Signal) error {
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
 
-func (ep *ExpectProcess) close(kill bool) error {
 	if ep.cmd == nil {
-		return ep.err
-	}
-	if kill {
-		ep.Signal(syscall.SIGTERM)
+		return errors.New("expect process already closed")
 	}
 
-	err := ep.cmd.Wait()
-	ep.fpty.Close()
-	ep.wg.Wait()
+	return ep.cmd.Process.Signal(sig)
+}
 
+func (ep *ExpectProcess) waitProcess() error {
+	state, err := ep.cmd.Process.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
-		}
+		return err
+	}
+
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
+	ep.exitCode = state.ExitCode()
+
+	if !state.Success() {
+		return fmt.Errorf("unexpected exit code [%d] after running [%s]", ep.exitCode, ep.cmd.String())
 	}
 
+	return nil
+}
+
+// Wait waits for the process to finish.
+func (ep *ExpectProcess) Wait() {
+	ep.wg.Wait()
+}
+
+// Close waits for the expect process to exit and return its error.
+func (ep *ExpectProcess) Close() error {
+	ep.wg.Wait()
+
+	ep.mu.Lock()
+	defer ep.mu.Unlock()
+
+	// this signals to other funcs that the process has finished
 	ep.cmd = nil
-	return err
+	return ep.exitErr
 }
 
 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()

pkg/expect/expect_test.go

@@ -19,9 +19,11 @@ package expect
 import (
 	"context"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -65,9 +67,57 @@ func TestExpectFuncTimeout(t *testing.T) {
 
 	require.ErrorAs(t, err, &context.DeadlineExceeded)
 
-	if err = ep.Stop(); err != nil {
+	if err := ep.Stop(); err != nil {
+		t.Fatal(err)
+	}
+
+	err = ep.Close()
+	require.ErrorContains(t, err, "unexpected exit code [-1] after running [/usr/bin/tail -f /dev/null]")
+	require.Equal(t, -1, ep.exitCode)
+}
+
+func TestExpectFuncExitFailure(t *testing.T) {
+	// tail -x should not exist and return a non-zero exit code
+	ep, err := NewExpect("tail", "-x")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	_, err = ep.ExpectFunc(ctx, func(s string) bool {
+		return strings.Contains(s, "something entirely unexpected")
+	})
+	require.ErrorContains(t, err, "unexpected exit code [1] after running [/usr/bin/tail -x]")
+	require.Equal(t, 1, ep.exitCode)
+}
+
+func TestExpectFuncExitFailureStop(t *testing.T) {
+	// tail -x should not exist and return a non-zero exit code
+	ep, err := NewExpect("tail", "-x")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
+	defer cancel()
+
+	_, err = ep.ExpectFunc(ctx, func(s string) bool {
+		return strings.Contains(s, "something entirely unexpected")
+	})
+	require.ErrorContains(t, err, "unexpected exit code [1] after running [/usr/bin/tail -x]")
+	exitCode, err := ep.ExitCode()
+	require.Equal(t, 0, exitCode)
+	require.Equal(t, err, ErrProcessRunning)
+	if err := ep.Stop(); err != nil {
 		t.Fatal(err)
 	}
+	err = ep.Close()
+	require.ErrorContains(t, err, "unexpected exit code [1] after running [/usr/bin/tail -x]")
+	exitCode, err = ep.ExitCode()
+	require.Equal(t, 1, exitCode)
+	require.NoError(t, err)
 }
 
 func TestEcho(t *testing.T) {
@@ -138,10 +188,8 @@ func TestSignal(t *testing.T) {
 	donec := make(chan struct{})
 	go func() {
 		defer close(donec)
-		werr := "signal: interrupt"
-		if cerr := ep.Close(); cerr == nil || cerr.Error() != werr {
-			t.Errorf("got error %v, wanted error %s", cerr, werr)
-		}
+		err = ep.Close()
+		assert.ErrorContains(t, err, "unexpected exit code [-1] after running [/usr/bin/sleep 100]")
 	}()
 	select {
 	case <-time.After(5 * time.Second):