Expect exit code enhancement

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,83 @@ 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
+	return err
+}
+
+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 +187,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 +220,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
 }
 
-func (ep *ExpectProcess) Wait() error {
-	_, err := ep.cmd.Process.Wait()
+// 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
+}
+
+// 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):