test(scaletest/createworkspaces): fix ctx init in multiple tests (#16177)

mafredri · web-flow · commit ea8cd5540468 · 2025-01-17T14:43:22.000Z
diff --git a/scaletest/createworkspaces/run_test.go b/scaletest/createworkspaces/run_test.go
@@ -52,9 +52,6 @@ func Test_Runner(t *testing.T) {
 	t.Run("OK", func(t *testing.T) {
 		t.Parallel()
 
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
-		defer cancel()
-
 		client := coderdtest.New(t, &coderdtest.Options{
 			IncludeProvisionerDaemon: true,
 		})
@@ -109,6 +106,8 @@ func Test_Runner(t *testing.T) {
 		version = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
 		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
 
+		ctx := testutil.Context(t, testutil.WaitLong)
+
 		closerCh := goEventuallyStartFakeAgent(ctx, t, client, authToken)
 
 		const (
@@ -199,9 +198,6 @@ func Test_Runner(t *testing.T) {
 	t.Run("CleanupPendingBuild", func(t *testing.T) {
 		t.Parallel()
 
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
-		defer cancel()
-
 		// need to include our own logger because the provisioner (rightly) drops error logs when we shut down the
 		// test with a build in progress.
 		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).Leveled(slog.LevelDebug)
@@ -253,7 +249,9 @@ func Test_Runner(t *testing.T) {
 			},
 		})
 
+		ctx := testutil.Context(t, testutil.WaitLong)
 		cancelCtx, cancelFunc := context.WithCancel(ctx)
+
 		done := make(chan struct{})
 		logs := bytes.NewBuffer(nil)
 		go func() {
@@ -287,6 +285,8 @@ func Test_Runner(t *testing.T) {
 		cancelFunc()
 		<-done
 
+		ctx = testutil.Context(t, testutil.WaitLong) // Reset ctx to avoid timeouts.
+
 		// When we run the cleanup, it should be canceled
 		cleanupLogs := bytes.NewBuffer(nil)
 		cancelCtx, cancelFunc = context.WithCancel(ctx)
@@ -340,9 +340,6 @@ func Test_Runner(t *testing.T) {
 	t.Run("NoCleanup", func(t *testing.T) {
 		t.Parallel()
 
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
-		defer cancel()
-
 		client := coderdtest.New(t, &coderdtest.Options{
 			IncludeProvisionerDaemon: true,
 		})
@@ -397,6 +394,7 @@ func Test_Runner(t *testing.T) {
 		version = coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID)
 		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
 
+		ctx := testutil.Context(t, testutil.WaitLong)
 		closeCh := goEventuallyStartFakeAgent(ctx, t, client, authToken)
 
 		const (
@@ -484,9 +482,6 @@ func Test_Runner(t *testing.T) {
 	t.Run("FailedBuild", func(t *testing.T) {
 		t.Parallel()
 
-		ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
-		defer cancel()
-
 		logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
 		client := coderdtest.New(t, &coderdtest.Options{
 			IncludeProvisionerDaemon: true,
@@ -534,6 +529,8 @@ func Test_Runner(t *testing.T) {
 			},
 		})
 
+		ctx := testutil.Context(t, testutil.WaitLong)
+
 		logs := bytes.NewBuffer(nil)
 		err := runner.Run(ctx, "1", logs)
 		logsStr := logs.String()
