From 55a524e51eef129b020d97dd0978bb9b5344b4e7 Mon Sep 17 00:00:00 2001 From: Irvin Lim Date: Mon, 10 Feb 2025 20:50:25 +0800 Subject: [PATCH] chore: Fix flaky PTY tests (#137) --- pkg/cli/cmd/cmd_list_job_test.go | 17 +++++-- pkg/cli/console/console.go | 35 +++++++++++-- pkg/cli/prompt/prompt_test.go | 18 +++---- pkg/runtime/testing/command.go | 85 ++++++++++++++++++++++++-------- 4 files changed, 116 insertions(+), 39 deletions(-) diff --git a/pkg/cli/cmd/cmd_list_job_test.go b/pkg/cli/cmd/cmd_list_job_test.go index 7929bf1..a25340b 100644 --- a/pkg/cli/cmd/cmd_list_job_test.go +++ b/pkg/cli/cmd/cmd_list_job_test.go @@ -516,7 +516,7 @@ func TestListJobCommand(t *testing.T) { }, { Name: "watch jobs belonging to jobconfig with --for and --finished", - Args: []string{"list", "job", "--for", adhocJobConfig.Name}, + Args: []string{"list", "job", "-w", "--for", adhocJobConfig.Name}, Fixtures: []runtime.Object{ jobRunning, jobFinished, @@ -524,15 +524,24 @@ func TestListJobCommand(t *testing.T) { adhocJobConfig, }, Procedure: func(t *testing.T, rc runtimetesting.RunContext) { + // Expect that headers and rows are printed. + rc.Console.ExpectString("NAME") + rc.Console.ExpectString(jobRunning.Name) + rc.Console.ExpectString(string(jobRunning.Status.Phase)) + // Update the Job's state. newJob := jobRunning.DeepCopy() newJob.Status.Phase = v1alpha1.JobSucceeded newJob.Status.State = v1alpha1.JobStateFinished + newJob.Status.Condition = v1alpha1.JobCondition{ + Finished: &v1alpha1.JobConditionFinished{ + FinishTimestamp: testutils.Mkmtime(taskFinishTime), + }, + } _, err := rc.CtrlContext.Clientsets().Furiko().ExecutionV1alpha1().Jobs(newJob.Namespace).Update(rc.Context, newJob, metav1.UpdateOptions{}) assert.NoError(t, err) - // Expect that headers and rows are printed. - rc.Console.ExpectString("NAME") + // Expect that we got the update. rc.Console.ExpectString(newJob.Name) rc.Console.ExpectString(string(newJob.Status.Phase)) @@ -540,7 +549,7 @@ func TestListJobCommand(t *testing.T) { rc.Cancel() }, Stdout: runtimetesting.Output{ - NumLines: pointer.Int(2), + NumLines: pointer.Int(3), }, WantActions: runtimetesting.CombinedActions{ Ignore: true, diff --git a/pkg/cli/console/console.go b/pkg/cli/console/console.go index a0d4ae2..ae47773 100644 --- a/pkg/cli/console/console.go +++ b/pkg/cli/console/console.go @@ -17,21 +17,28 @@ package console import ( + "bytes" "io" + "testing" "github.com/Netflix/go-expect" "github.com/creack/pty" "github.com/hinshun/vt10x" "github.com/pkg/errors" + "github.com/stretchr/testify/assert" ) type Console struct { *expect.Console + T *testing.T + + // The bytes that were read during Expect. + read bytes.Buffer } // NewConsole returns a Console for testing PTY handling. // All PTY output will be piped to w. -func NewConsole(w io.Writer) (*Console, error) { +func NewConsole(t *testing.T, w io.Writer) (*Console, error) { ptty, tty, err := pty.Open() if err != nil { return nil, errors.Wrapf(err, "failed to open pty") @@ -48,6 +55,7 @@ func NewConsole(w io.Writer) (*Console, error) { } c := &Console{ + T: t, Console: console, } @@ -68,10 +76,27 @@ func (c *Console) Run(f func(c *Console)) <-chan struct{} { return done } -func (c *Console) ExpectString(s string) { - _, _ = c.Console.ExpectString(s) +// ExpectString advances the PTY buffer until the given string is found. +// If not, a test error will be thrown, and returns false. +func (c *Console) ExpectString(s string) bool { + got, err := c.Console.Expect(expect.String(s)) + c.read.WriteString(got) + return assert.NoError(c.T, err, `did not find expected string: "%v", got "%v"`, s, got) +} + +// SendLine writes to the PTY buffer. +// Blocks until the line is written. +func (c *Console) SendLine(s string) bool { + _, err := c.Console.SendLine(s) + return assert.NoError(c.T, err, `cannot send line: "%v"`, s) } -func (c *Console) SendLine(s string) { - _, _ = c.Console.SendLine(s) +// Close the TTY, unblocking all Expect and ExpectEOF calls. +func (c *Console) Close() error { + err := c.Console.Close() + output := c.read.String() + if len(output) > 0 { + c.T.Logf("Console output:\n%s\n\n", output) + } + return err } diff --git a/pkg/cli/prompt/prompt_test.go b/pkg/cli/prompt/prompt_test.go index 9e0ed4f..0547c47 100644 --- a/pkg/cli/prompt/prompt_test.go +++ b/pkg/cli/prompt/prompt_test.go @@ -99,7 +99,7 @@ func TestNewBoolPrompt(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - pt, ps, err := NewPromptTest() + pt, ps, err := NewPromptTest(t) if err != nil { t.Fatalf("cannot initialize test: %v", err) } @@ -214,7 +214,7 @@ func TestNewStringPrompt(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - pt, ps, err := NewPromptTest() + pt, ps, err := NewPromptTest(t) if err != nil { t.Fatalf("cannot initialize test: %v", err) } @@ -324,7 +324,7 @@ func TestNewSelectPrompt(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - pt, ps, err := NewPromptTest() + pt, ps, err := NewPromptTest(t) if err != nil { t.Fatalf("cannot initialize test: %v", err) } @@ -418,7 +418,7 @@ func TestNewMultiPrompt(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - pt, ps, err := NewPromptTest() + pt, ps, err := NewPromptTest(t) if err != nil { t.Fatalf("cannot initialize test: %v", err) } @@ -472,7 +472,7 @@ func TestNewDatePrompt(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - pt, ps, err := NewPromptTest() + pt, ps, err := NewPromptTest(t) if err != nil { t.Fatalf("cannot initialize test: %v", err) } @@ -492,9 +492,9 @@ type PromptTest struct { console *console.Console } -func NewPromptTest() (*PromptTest, *streams.Streams, error) { +func NewPromptTest(t *testing.T) (*PromptTest, *streams.Streams, error) { iostreams, _, _, _ := genericclioptions.NewTestIOStreams() - c, err := console.NewConsole(iostreams.Out) + c, err := console.NewConsole(t, iostreams.Out) if err != nil { return nil, nil, errors.Wrap(err, "failed to init pseudo TTY") } @@ -546,6 +546,6 @@ func (pt *PromptTest) Run(t *testing.T, p prompt.Prompt, procedure func(c *conso return resp.retval, resp.err } -func (pt *PromptTest) Close() { - _ = pt.console.Close() +func (pt *PromptTest) Close() error { + return pt.console.Close() } diff --git a/pkg/runtime/testing/command.go b/pkg/runtime/testing/command.go index 368ebee..264f317 100644 --- a/pkg/runtime/testing/command.go +++ b/pkg/runtime/testing/command.go @@ -17,14 +17,17 @@ package testing import ( + "bytes" "context" "fmt" "regexp" "strings" + "sync" "testing" "time" "github.com/google/go-cmp/cmp" + "github.com/pkg/errors" "github.com/stretchr/testify/assert" "k8s.io/apimachinery/pkg/runtime" "k8s.io/cli-runtime/pkg/genericclioptions" @@ -39,6 +42,10 @@ import ( "github.com/furiko-io/furiko/pkg/utils/testutils" ) +const ( + testTimeout = time.Second * 5 +) + // RunCommandTests executes all CommandTest cases. func RunCommandTests(t *testing.T, cases []CommandTest) { for _, tt := range cases { @@ -101,7 +108,7 @@ type Output struct { // If specified, expects the output to match the specified regexp. Matches *regexp.Regexp - // If specified, expects the output to match all of the given regular expressions. + // If specified, expects the output to match all the given regular expressions. MatchesAll []*regexp.Regexp // If specified, expects that the output to NOT contain the given string. @@ -129,27 +136,32 @@ type RunContext struct { } func (c *CommandTest) Run(t *testing.T) { - ctx, cancel := context.WithTimeout(context.Background(), time.Second) + ctx, cancel := context.WithTimeout(context.Background(), testTimeout) defer cancel() + ctrlContext := mock.NewContext() + common.SetCtrlContext(ctrlContext) + iostreams, _, stdout, stderr := genericclioptions.NewTestIOStreams() + done := make(chan struct{}) go func() { - c.run(ctx, t) + c.run(ctx, ctrlContext, t, iostreams) close(done) }() select { case <-done: + c.verify(t, ctrlContext, stdout, stderr) return case <-ctx.Done(): - t.Errorf("test did not complete: %v", ctx.Err()) + t.Errorf("test did not complete after %v: %v", testTimeout, ctx.Err()) + t.Logf("command stdout =\n%s", stdout.String()) + t.Logf("command stderr =\n%s", stderr.String()) } } -func (c *CommandTest) run(ctx context.Context, t *testing.T) { +func (c *CommandTest) run(ctx context.Context, ctrlContext *mock.Context, t *testing.T, iostreams genericclioptions.IOStreams) bool { // Override the shared context. - ctrlContext := mock.NewContext() - common.SetCtrlContext(ctrlContext) client := ctrlContext.MockClientsets() assert.NoError(t, InitFixtures(ctx, client, c.Fixtures)) client.ClearActions() @@ -162,10 +174,11 @@ func (c *CommandTest) run(ctx context.Context, t *testing.T) { ktime.Clock = fakeclock.NewFakeClock(now) // Run command with I/O. - iostreams, _, stdout, stderr := genericclioptions.NewTestIOStreams() - if c.runCommand(ctx, t, ctrlContext, iostreams) { - return - } + return c.runCommand(ctx, t, ctrlContext, iostreams) +} + +func (c *CommandTest) verify(t *testing.T, ctrlContext *mock.Context, stdout, stderr *bytes.Buffer) { + client := ctrlContext.MockClientsets() // Ensure that output matches. c.checkOutput(t, "stdout", stdout.String(), c.Stdout) @@ -176,7 +189,7 @@ func (c *CommandTest) run(ctx context.Context, t *testing.T) { } // runCommand will execute the command, setting up all I/O streams and blocking -// until the streams are done. +// until the streams are done. Returns true if an error was encountered. // // Reference: // https://github.com/AlecAivazis/survey/blob/93657ef69381dd1ffc7a4a9cfe5a2aefff4ca4ad/survey_posix_test.go#L15 @@ -184,11 +197,10 @@ func (c *CommandTest) runCommand(ctx context.Context, t *testing.T, ctrlContext ctx, cancel := context.WithCancel(ctx) defer cancel() - console, err := console.NewConsole(iostreams.Out) + console, err := console.NewConsole(t, iostreams.Out) if err != nil { t.Fatalf("failed to create console: %v", err) } - defer console.Close() // Prepare root command. command := cmd.NewRootCommand(streams.NewTTYStreams(console.Tty())) @@ -196,7 +208,7 @@ func (c *CommandTest) runCommand(ctx context.Context, t *testing.T, ctrlContext var done <-chan struct{} - // Run procedure if specified. + // Run procedure in the background. if c.Procedure != nil { done = c.runProcedure(t, c.Procedure, RunContext{ Console: console, @@ -208,6 +220,41 @@ func (c *CommandTest) runCommand(ctx context.Context, t *testing.T, ctrlContext done = console.Run(c.Stdin.Procedure) } + // Only close the PTY once. + closeConsole := sync.OnceFunc(func() { + if err := console.Close(); err != nil { + t.Errorf("cannot close PTY: %v", err) + } + }) + + // Close the console if context had deadline exceeded. + // Note that the context could possibly be canceled by tests in order to terminate the command execution. + var wg sync.WaitGroup + wg.Add(1) + go func() { + defer wg.Done() + select { + case <-ctx.Done(): + if errors.Is(ctx.Err(), context.DeadlineExceeded) { + t.Errorf("Context was canceled: %v", ctx.Err().Error()) + closeConsole() + } + case <-done: + // Already closed. + } + }() + + defer func() { + // Always make sure to explicitly close the PTY. + closeConsole() + + // Wait for the procedure to complete and EOF to be read. + <-done + + // Wait for . + wg.Wait() + }() + // Execute command. if testutils.WantError(t, c.WantError, command.ExecuteContext(ctx), fmt.Sprintf("Run error with args: %v", c.Args)) { return true @@ -216,15 +263,11 @@ func (c *CommandTest) runCommand(ctx context.Context, t *testing.T, ctrlContext // TODO(irvinlim): We need a sleep here otherwise tests will be flaky time.Sleep(time.Millisecond * 5) - // Wait for tty to be closed. - if err := console.Tty().Close(); err != nil { - t.Errorf("error closing Tty: %v", err) - } - <-done - return false } +// runProcedure starts the procedure in the background, and returns a channel +// that will be closed once the PTY output is closed. func (c *CommandTest) runProcedure(t *testing.T, procedure func(t *testing.T, rc RunContext), rc RunContext) <-chan struct{} { done := make(chan struct{}) go func() {