fix: Fix goleak in cli TestSSH tests

mafredri · mafredri · commit eeb74e1b8958 · 2022-07-27T20:11:04.000+03:00
Commands are now also run with contexts that time out. Work towards #3221.
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
@@ -68,15 +68,20 @@ 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)
 		pty := ptytest.New(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,14 +118,22 @@ 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)
 		cmd.SetIn(clientOutput)
 		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)
 		})
 
