From a16ccf19829dc5ca3b22f6d6febb12927f1271b0 Mon Sep 17 00:00:00 2001 From: apostasie Date: Sun, 20 Oct 2024 13:19:53 -0700 Subject: [PATCH] Add support for writing to ptys Signed-off-by: apostasie --- pkg/testutil/test/command.go | 18 +++++++++++++++--- pkg/testutil/test/test.go | 3 ++- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/pkg/testutil/test/command.go b/pkg/testutil/test/command.go index 5f072e647a8..68735968f70 100644 --- a/pkg/testutil/test/command.go +++ b/pkg/testutil/test/command.go @@ -49,6 +49,7 @@ type GenericCommand struct { stdin io.Reader async bool pty bool + ptyWriters []func(*os.File) error timeout time.Duration workingDir string @@ -69,8 +70,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) { @@ -105,8 +107,18 @@ func (gc *GenericCommand) Run(expect *Expected) { iCmdCmd.Stdin = tty iCmdCmd.Stdout = tty iCmdCmd.Stderr = tty - defer pty.Close() - defer tty.Close() + + for _, writer := range gc.ptyWriters { + go func() { + err := writer(pty) + assert.NilError(gc.t, err) + }() + } + + defer func() { + _ = tty.Close() + _ = pty.Close() + }() } // Run it diff --git a/pkg/testutil/test/test.go b/pkg/testutil/test/test.go index 07a7e3b1c60..a42a4c87dcc 100644 --- a/pkg/testutil/test/test.go +++ b/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