From 259dee2ea899e846a327b9f7b317c1223bc95962 Mon Sep 17 00:00:00 2001
From: Danny Kopping <danny@coder.com>
Date: Fri, 5 Dec 2025 15:14:35 +0200
Subject: [PATCH] fix: move contexts to appropriate locations (#21121)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Closes https://github.com/coder/internal/issues/1173,
https://github.com/coder/internal/issues/1174

Currently these two tests are flaky because the contexts were created
before a potentially long-running process. By the time the context was
actually used, it may have timed out - leading to confusion.

Additionally, the `ExpectMatch` calls were not using the test context -
but rather a background context. I've marked that func as deprecated
because we should always tie these to the test context.

Special thanks to @mafredri for the brain probe 🧠

---------

Signed-off-by: Danny Kopping <danny@coder.com>
---
 cli/exp_rpty_test.go   | 6 +++---
 cli/ssh_test.go        | 6 +++---
 pty/ptytest/ptytest.go | 2 ++
 3 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/cli/exp_rpty_test.go b/cli/exp_rpty_test.go
index c7a0c47d18..9b5c0ffe22 100644
--- a/cli/exp_rpty_test.go
+++ b/cli/exp_rpty_test.go
@@ -90,7 +90,6 @@ func TestExpRpty(t *testing.T) {
 		wantLabel := "coder.devcontainers.TestExpRpty.Container"
 
 		client, workspace, agentToken := setupWorkspaceForAgent(t)
-		ctx := testutil.Context(t, testutil.WaitLong)
 		pool, err := dockertest.NewPool("")
 		require.NoError(t, err, "Could not connect to docker")
 		ct, err := pool.RunWithOptions(&dockertest.RunOptions{
@@ -128,14 +127,15 @@ func TestExpRpty(t *testing.T) {
 		clitest.SetupConfig(t, client, root)
 		pty := ptytest.New(t).Attach(inv)
 
+		ctx := testutil.Context(t, testutil.WaitLong)
 		cmdDone := tGo(t, func() {
 			err := inv.WithContext(ctx).Run()
 			assert.NoError(t, err)
 		})
 
-		pty.ExpectMatch(" #")
+		pty.ExpectMatchContext(ctx, " #")
 		pty.WriteLine("hostname")
-		pty.ExpectMatch(ct.Container.Config.Hostname)
+		pty.ExpectMatchContext(ctx, ct.Container.Config.Hostname)
 		pty.WriteLine("exit")
 		<-cmdDone
 	})
diff --git a/cli/ssh_test.go b/cli/ssh_test.go
index 7ce9d85258..2b3113a901 100644
--- a/cli/ssh_test.go
+++ b/cli/ssh_test.go
@@ -2052,7 +2052,6 @@ func TestSSH_Container(t *testing.T) {
 		t.Parallel()
 
 		client, workspace, agentToken := setupWorkspaceForAgent(t)
-		ctx := testutil.Context(t, testutil.WaitLong)
 		pool, err := dockertest.NewPool("")
 		require.NoError(t, err, "Could not connect to docker")
 		ct, err := pool.RunWithOptions(&dockertest.RunOptions{
@@ -2087,14 +2086,15 @@ func TestSSH_Container(t *testing.T) {
 		clitest.SetupConfig(t, client, root)
 		ptty := ptytest.New(t).Attach(inv)
 
+		ctx := testutil.Context(t, testutil.WaitLong)
 		cmdDone := tGo(t, func() {
 			err := inv.WithContext(ctx).Run()
 			assert.NoError(t, err)
 		})
 
-		ptty.ExpectMatch(" #")
+		ptty.ExpectMatchContext(ctx, " #")
 		ptty.WriteLine("hostname")
-		ptty.ExpectMatch(ct.Container.Config.Hostname)
+		ptty.ExpectMatchContext(ctx, ct.Container.Config.Hostname)
 		ptty.WriteLine("exit")
 		<-cmdDone
 	})
diff --git a/pty/ptytest/ptytest.go b/pty/ptytest/ptytest.go
index 3991bdeb04..5d15078094 100644
--- a/pty/ptytest/ptytest.go
+++ b/pty/ptytest/ptytest.go
@@ -145,6 +145,8 @@ type outExpecter struct {
 	runeReader *bufio.Reader
 }
 
+// Deprecated: use ExpectMatchContext instead.
+// This uses a background context, so will not respect the test's context.
 func (e *outExpecter) ExpectMatch(str string) string {
 	return e.expectMatchContextFunc(str, e.ExpectMatchContext)
 }