Add pty support for testing

Signed-off-by: apostasie <[email protected]>

Author: apostasie

pkg/testutil/test/command.go

@@ -17,15 +17,20 @@
 package test
 
 import (
+	"bytes"
 	"fmt"
 	"io"
 	"os"
 	"strings"
 	"testing"
 	"time"
 
+	"golang.org/x/sync/errgroup"
+	"golang.org/x/term"
 	"gotest.tools/v3/assert"
 	"gotest.tools/v3/icmd"
+
+	"github.com/containerd/nerdctl/v2/pkg/testutil/test/internal/pty"
 )
 
 const ExitCodeGenericFail = -1
@@ -49,6 +54,7 @@ type GenericCommand struct {
 	stdin        io.Reader
 	async        bool
 	pty          bool
+	ptyWriters   []func(*os.File) error
 	timeout      time.Duration
 	workingDir   string
 
@@ -69,8 +75,9 @@ func (gc *GenericCommand) WithWrapper(binary string, args ...string) {
 	gc.helperArgs = args
 }
 
-func (gc *GenericCommand) WithPseudoTTY() {
+func (gc *GenericCommand) WithPseudoTTY(writers ...func(*os.File) error) {
 	gc.pty = true
+	gc.ptyWriters = writers
 }
 
 func (gc *GenericCommand) WithStdin(r io.Reader) {
@@ -81,36 +88,58 @@ func (gc *GenericCommand) WithCwd(path string) {
 	gc.workingDir = path
 }
 
-// TODO: it should be possible to timeout execution
-// Primitives (gc.timeout) is here, it is just a matter of exposing a WithTimeout method
-// - UX to be decided
-// - validate use case: would we ever need this?
-
 func (gc *GenericCommand) Run(expect *Expected) {
 	if gc.t != nil {
 		gc.t.Helper()
 	}
 
 	var result *icmd.Result
 	var env []string
-	if gc.async {
-		result = icmd.WaitOnCmd(gc.timeout, gc.result)
-		env = gc.result.Cmd.Env
-	} else {
+	output := &bytes.Buffer{}
+	stdout := ""
+	errorGroup := &errgroup.Group{}
+	var tty *os.File
+	var psty *os.File
+	if !gc.async {
 		iCmdCmd := gc.boot()
 		env = iCmdCmd.Env
 
 		if gc.pty {
-			pty, tty, _ := Open()
+			psty, tty, _ = pty.Open()
+			_, _ = term.MakeRaw(int(tty.Fd()))
+
 			iCmdCmd.Stdin = tty
 			iCmdCmd.Stdout = tty
-			iCmdCmd.Stderr = tty
-			defer pty.Close()
-			defer tty.Close()
+
+			gc.result = icmd.StartCmd(iCmdCmd)
+
+			for _, writer := range gc.ptyWriters {
+				_ = writer(psty)
+			}
+
+			// Copy from the master
+			errorGroup.Go(func() error {
+				_, _ = io.Copy(output, psty)
+				return nil
+			})
+		} else {
+			// Run it
+			gc.result = icmd.StartCmd(iCmdCmd)
 		}
+	}
+
+	result = icmd.WaitOnCmd(gc.timeout, gc.result)
+	env = gc.result.Cmd.Env
+
+	if gc.pty {
+		_ = tty.Close()
+		_ = psty.Close()
+		_ = errorGroup.Wait()
+	}
 
-		// Run it
-		result = icmd.RunCmd(iCmdCmd)
+	stdout = result.Stdout()
+	if stdout == "" {
+		stdout = output.String()
 	}
 
 	gc.rawStdErr = result.Stderr()
@@ -137,7 +166,7 @@ func (gc *GenericCommand) Run(expect *Expected) {
 		}
 		// Finally, check the output if we are asked to
 		if expect.Output != nil {
-			expect.Output(result.Stdout(), debug, gc.t)
+			expect.Output(stdout, debug, gc.t)
 		}
 	}
 }

pkg/testutil/test/pty.go renamed to pkg/testutil/test/internal/pty/pty.go

@@ -14,9 +14,11 @@
    limitations under the License.
 */
 
-package test
+package pty
 
 import "errors"
 
-var ErrPTYFailure = errors.New("pty failure")
-var ErrPTYUnsupportedPlatform = errors.New("pty not supported on this platform")
+var (
+	ErrPTYFailure             = errors.New("pty failure")
+	ErrPTYUnsupportedPlatform = errors.New("pty not supported on this platform")
+)

pkg/testutil/test/pty_freebsd.go renamed to pkg/testutil/test/internal/pty/pty_freebsd.go

@@ -14,7 +14,7 @@
    limitations under the License.
 */
 
-package test
+package pty
 
 import (
 	"os"

pkg/testutil/test/pty_linux.go renamed to pkg/testutil/test/internal/pty/pty_linux.go

@@ -14,7 +14,7 @@
    limitations under the License.
 */
 
-package test
+package pty
 
 import (
 	"errors"
@@ -27,6 +27,7 @@ import (
 // Inspiration from https://github.com/creack/pty/tree/2cde18bfb702199728dd43bf10a6c15c7336da0a
 
 func Open() (pty, tty *os.File, err error) {
+	// Wrap errors
 	defer func() {
 		if err != nil && pty != nil {
 			err = errors.Join(pty.Close(), err)
@@ -36,37 +37,33 @@ func Open() (pty, tty *os.File, err error) {
 		}
 	}()
 
+	// Open the pty
 	pty, err = os.OpenFile("/dev/ptmx", os.O_RDWR, 0)
 	if err != nil {
 		return nil, nil, err
 	}
 
+	// Get the slave unit number
 	var n uint32
-	err = ioctl(pty, syscall.TIOCGPTN, uintptr(unsafe.Pointer(&n)))
-	if err != nil {
+	_, _, e := syscall.Syscall(syscall.SYS_IOCTL, pty.Fd(), syscall.TIOCGPTN, uintptr(unsafe.Pointer(&n)))
+	if e != 0 {
 		return nil, nil, err
 	}
 
 	sname := "/dev/pts/" + strconv.Itoa(int(n))
 
+	// Unlock
 	var u int32
-	err = ioctl(pty, syscall.TIOCSPTLCK, uintptr(unsafe.Pointer(&u)))
-	if err != nil {
+	_, _, e = syscall.Syscall(syscall.SYS_IOCTL, pty.Fd(), syscall.TIOCSPTLCK, uintptr(unsafe.Pointer(&u)))
+	if e != 0 {
 		return nil, nil, err
 	}
 
+	// Open the slave, preventing it from becoming the controlling terminal
 	tty, err = os.OpenFile(sname, os.O_RDWR|syscall.O_NOCTTY, 0)
 	if err != nil {
 		return nil, nil, err
 	}
 
 	return pty, tty, nil
 }
-
-func ioctl(f *os.File, cmd, ptr uintptr) error {
-	_, _, e := syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), cmd, ptr)
-	if e != 0 {
-		return e
-	}
-	return nil
-}

pkg/testutil/test/pty_windows.go renamed to pkg/testutil/test/internal/pty/pty_windows.go

@@ -14,7 +14,7 @@
    limitations under the License.
 */
 
-package test
+package pty
 
 import (
 	"os"

pkg/testutil/test/test.go

@@ -18,6 +18,7 @@ package test
 
 import (
 	"io"
+	"os"
 	"testing"
 	"time"
 )
@@ -98,7 +99,7 @@ type TestableCommand interface {
 	// WithWrapper allows wrapping a command with another command (for example: `time`, `unbuffer`)
 	WithWrapper(binary string, args ...string)
 	// WithPseudoTTY
-	WithPseudoTTY()
+	WithPseudoTTY(writers ...func(*os.File) error)
 	// WithStdin allows passing a reader to be used for stdin for the command
 	WithStdin(r io.Reader)
 	// WithCwd allows specifying the working directory for the command