diff --git a/cli/ssh_test.go b/cli/ssh_test.go index c5849e99eada6..f88339c103fd9 100644 --- a/cli/ssh_test.go +++ b/cli/ssh_test.go @@ -68,6 +68,7 @@ func TestSSH(t *testing.T) { t.Parallel() t.Run("ImmediateExit", func(t *testing.T) { t.Parallel() + client, workspace, agentToken := setupWorkspaceForSSH(t) cmd, root := clitest.New(t, "ssh", workspace.Name) clitest.SetupConfig(t, client, root) @@ -75,8 +76,12 @@ func TestSSH(t *testing.T) { cmd.SetIn(pty.Input()) cmd.SetErr(pty.Output()) cmd.SetOut(pty.Output()) + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + cmdDone := tGo(t, func() { - err := cmd.Execute() + err := cmd.ExecuteContext(ctx) assert.NoError(t, err) }) pty.ExpectMatch("Waiting") @@ -85,9 +90,9 @@ func TestSSH(t *testing.T) { agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{ Logger: slogtest.Make(t, nil).Leveled(slog.LevelDebug), }) - t.Cleanup(func() { + defer func() { _ = agentCloser.Close() - }) + }() // Shells on Mac, Windows, and Linux all exit shells with the "exit" command. pty.WriteLine("exit") @@ -113,6 +118,14 @@ func TestSSH(t *testing.T) { clientOutput, clientInput := io.Pipe() serverOutput, serverInput := io.Pipe() + defer func() { + for _, c := range []io.Closer{clientOutput, clientInput, serverOutput, serverInput} { + _ = c.Close() + } + }() + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() cmd, root := clitest.New(t, "ssh", "--stdio", workspace.Name) clitest.SetupConfig(t, client, root) @@ -120,7 +133,7 @@ func TestSSH(t *testing.T) { cmd.SetOut(serverInput) cmd.SetErr(io.Discard) cmdDone := tGo(t, func() { - err := cmd.Execute() + err := cmd.ExecuteContext(ctx) assert.NoError(t, err) }) @@ -132,9 +145,13 @@ func TestSSH(t *testing.T) { HostKeyCallback: ssh.InsecureIgnoreHostKey(), }) require.NoError(t, err) + defer conn.Close() + sshClient := ssh.NewClient(conn, channels, requests) session, err := sshClient.NewSession() require.NoError(t, err) + defer session.Close() + command := "sh -c exit" if runtime.GOOS == "windows" { command = "cmd.exe /c exit" @@ -198,6 +215,9 @@ func TestSSH(t *testing.T) { } }) + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + cmd, root := clitest.New(t, "ssh", workspace.Name, @@ -210,7 +230,7 @@ func TestSSH(t *testing.T) { cmd.SetOut(pty.Output()) cmd.SetErr(io.Discard) cmdDone := tGo(t, func() { - err := cmd.Execute() + err := cmd.ExecuteContext(ctx) assert.NoError(t, err) })