test: use t.Context() instead of context.Background() (#1306)

The test context is cancelled when the test ends, so anything still holding it stops instead of leaking into the next test. Uses inside `t.Cleanup`, `defer` and `go func` blocks are left alone, as the context is already cancelled there.

Stacked on https://gitea.com/gitea/runner/pulls/1305.

Assisted-by: Codet
Co-authored-by: bircni <[email protected]>
Reviewed-on: https://gitea.com/gitea/runner/pulls/1306
Reviewed-by: bircni <[email protected]>

internal/app/cmd/exec_test.go

@@ -5,7 +5,6 @@ package cmd
 
 import (
 	"bytes"
-	"context"
 	"io"
 	"os"
 	"path/filepath"
@@ -217,7 +216,7 @@ func TestPrintListReportsDuplicateJobIDs(t *testing.T) {
 }
 
 func TestLoadExecCmdDefinesExpectedFlags(t *testing.T) {
-	cmd := loadExecCmd(context.Background())
+	cmd := loadExecCmd(t.Context())
 
 	for _, name := range []string{
 		"list",

internal/app/poll/poller_test.go

@@ -38,7 +38,7 @@ func TestPoller_WorkerStateCounters(t *testing.T) {
 	require.NoError(t, err)
 	p := &Poller{client: client, cfg: cfg}
 
-	ctx := context.Background()
+	ctx := t.Context()
 	s1 := &workerState{}
 	s2 := &workerState{}
 
@@ -69,7 +69,7 @@ func TestPoller_FetchErrorIncrementsErrorsOnly(t *testing.T) {
 	p := &Poller{client: client, cfg: cfg}
 
 	s := &workerState{}
-	_, ok := p.fetchTask(context.Background(), s)
+	_, ok := p.fetchTask(t.Context(), s)
 	require.False(t, ok)
 	assert.Equal(t, int64(1), s.consecutiveErrors)
 	assert.Equal(t, int64(0), s.consecutiveEmpty)
@@ -90,14 +90,14 @@ func TestPoller_FetchTimeoutIsNoSignal(t *testing.T) {
 	hook := logtest.SetDefault(t)
 
 	s := &workerState{consecutiveErrors: 2}
-	_, ok := p.fetchTask(context.Background(), s)
+	_, ok := p.fetchTask(t.Context(), s)
 	require.False(t, ok)
 	assert.Equal(t, int64(2), s.consecutiveErrors)
 	assert.Equal(t, int64(1), s.consecutiveEmpty)
 	assert.True(t, p.lastPollFailed.Load(), "a timeout must not clear a known failure")
 
 	// An idle runner against a server holding the request open must not flood.
-	_, _ = p.fetchTask(context.Background(), s)
+	_, _ = p.fetchTask(t.Context(), s)
 	assert.Len(t, hook.AllRecords(), 1)
 }
 
@@ -120,7 +120,7 @@ func TestPoller_FetchUnauthenticatedStopsPolling(t *testing.T) {
 			require.NoError(t, err)
 			p := New(cfg, client, nil)
 			s := &workerState{}
-			_, ok := p.fetchTask(context.Background(), s)
+			_, ok := p.fetchTask(t.Context(), s)
 			require.False(t, ok)
 
 			assert.True(t, p.Unregistered(), "runner should be marked unregistered")
@@ -261,7 +261,7 @@ func TestPollerPausesAndResumesForLocalAvailability(t *testing.T) {
 		return fetches.Load() > 0
 	}, time.Second, 10*time.Millisecond, "polling should resume after local recovery")
 
-	shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Second)
+	shutdownCtx, cancel := context.WithTimeout(t.Context(), time.Second)
 	defer cancel()
 	require.NoError(t, poller.Shutdown(shutdownCtx))
 	wg.Wait()
@@ -280,7 +280,7 @@ func (m *mockRunner) Run(ctx context.Context, _ *runnerv1.Task) error {
 
 func TestPollerRunIdleMaintenance(t *testing.T) {
 	runner := &idleAwareRunner{}
-	p := &Poller{runner: runner, jobsCtx: context.Background()}
+	p := &Poller{runner: runner, jobsCtx: t.Context()}
 
 	p.runIdleMaintenance()
 
@@ -315,7 +315,7 @@ func TestPollerPollCallsOnIdle(t *testing.T) {
 		return runner.idleCalls.Load() > 0
 	}, time.Second, 10*time.Millisecond)
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
 	defer cancel()
 	require.NoError(t, poller.Shutdown(ctx))
 	wg.Wait()
@@ -344,7 +344,7 @@ func TestPollerPollOnceCallsOnIdle(t *testing.T) {
 		return runner.idleCalls.Load() > 0
 	}, time.Second, 10*time.Millisecond)
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
 	defer cancel()
 	require.NoError(t, poller.Shutdown(ctx))
 	wg.Wait()
@@ -411,7 +411,7 @@ func TestPoller_ConcurrencyLimitedByCapacity(t *testing.T) {
 		return runner.totalCompleted.Load() >= totalTasks
 	}, 2*time.Second, 10*time.Millisecond, "all tasks should complete")
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	ctx, cancel := context.WithTimeout(t.Context(), time.Second)
 	defer cancel()
 	err = poller.Shutdown(ctx)
 	require.NoError(t, err)
@@ -467,7 +467,7 @@ func TestPoller_ShutdownForcesJobsOnTimeout(t *testing.T) {
 		return runner.running.Load() == 1
 	}, time.Second, 10*time.Millisecond, "task should start running")
 
-	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
 	defer cancel()
 	start := time.Now()
 	err = poller.Shutdown(ctx)

internal/app/run/post_task_script_test.go

@@ -25,7 +25,7 @@ func TestRunPostTaskScriptSkippedWhenEmpty(t *testing.T) {
 	r := &Runner{
 		cfg: &config.Config{},
 	}
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 
 	taskCtx, err := structpb.NewStruct(map[string]any{})
@@ -48,7 +48,7 @@ func TestRunPostTaskScriptNonZeroExitDoesNotPanic(t *testing.T) {
 	cfg.Runner.PostTaskScript = scriptPath
 
 	r := &Runner{cfg: cfg}
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 
 	taskCtx, err := structpb.NewStruct(map[string]any{})
@@ -66,7 +66,7 @@ func TestPostTaskScriptContextUsesFullTimeout(t *testing.T) {
 
 	// A task context that finished close to its own deadline must not truncate the
 	// script's budget: the script should still get its full configured timeout.
-	near, cancelNear := context.WithTimeout(context.Background(), time.Second)
+	near, cancelNear := context.WithTimeout(t.Context(), time.Second)
 	defer cancelNear()
 	scriptCtx, cancel := postTaskScriptContext(near, timeout)
 	defer cancel()
@@ -75,7 +75,7 @@ func TestPostTaskScriptContextUsesFullTimeout(t *testing.T) {
 	assert.Greater(t, time.Until(deadline), time.Minute, "script timeout truncated to task deadline")
 
 	// An already-cancelled task context must not cancel the script either.
-	cancelledCtx, cancelIt := context.WithCancel(context.Background())
+	cancelledCtx, cancelIt := context.WithCancel(t.Context())
 	cancelIt()
 	scriptCtx2, cancel2 := postTaskScriptContext(cancelledCtx, timeout)
 	defer cancel2()
@@ -91,7 +91,7 @@ func TestPostTaskScriptEnv(t *testing.T) {
 		envs: map[string]string{"BASE": "1"},
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 
 	taskCtx, err := structpb.NewStruct(map[string]any{
@@ -128,7 +128,7 @@ func TestRunPostTaskScriptIntegration(t *testing.T) {
 		envs: map[string]string{"CUSTOM": "runner-env"},
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 
 	taskCtx, err := structpb.NewStruct(map[string]any{})

internal/app/run/runner_idle_cleanup_test.go

@@ -46,7 +46,7 @@ func TestRunnerCleanupStaleTaskDirs(t *testing.T) {
 		now: func() time.Time { return now },
 	}
 
-	r.cleanupStaleDirs(context.Background(), workdirRoot, isTaskIDDir)
+	r.cleanupStaleDirs(t.Context(), workdirRoot, isTaskIDDir)
 
 	assert.NoDirExists(t, oldTask)
 	assert.DirExists(t, freshTask)
@@ -85,7 +85,7 @@ func TestRunnerOnIdleCleansStaleHostScratchDirs(t *testing.T) {
 		now: func() time.Time { return now },
 	}
 
-	r.OnIdle(context.Background())
+	r.OnIdle(t.Context())
 
 	assert.NoDirExists(t, staleScratch) // stale scratch reclaimed
 	assert.DirExists(t, freshScratch)   // within cleanup age, kept
@@ -113,7 +113,7 @@ func TestRunnerCleanupStaleTaskDirsMissingRoot(t *testing.T) {
 
 	// Must be a silent no-op rather than a warning or panic when the root
 	// has not yet been created (e.g. the runner has never executed a task).
-	r.cleanupStaleDirs(context.Background(), filepath.Join(t.TempDir(), "missing"), isTaskIDDir)
+	r.cleanupStaleDirs(t.Context(), filepath.Join(t.TempDir(), "missing"), isTaskIDDir)
 }
 
 func TestRunnerCleanupStaleTaskDirsHonorsContext(t *testing.T) {
@@ -127,7 +127,7 @@ func TestRunnerCleanupStaleTaskDirsHonorsContext(t *testing.T) {
 		require.NoError(t, os.Chtimes(dir, now.Add(-3*time.Hour), now.Add(-3*time.Hour)))
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	cancel()
 
 	r := &Runner{
@@ -265,7 +265,7 @@ func TestRunnerOnIdleIntegratesCleanup(t *testing.T) {
 		now: func() time.Time { return now },
 	}
 
-	r.OnIdle(context.Background())
+	r.OnIdle(t.Context())
 
 	assert.NoDirExists(t, stale)
 }
@@ -293,7 +293,7 @@ func TestRunnerOnIdleSkipsWhenAlreadyCancelled(t *testing.T) {
 		now: func() time.Time { return now },
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	cancel()
 	r.OnIdle(ctx)
 
@@ -346,7 +346,7 @@ func TestRunnerOnIdleRemovesOrphanNetworks(t *testing.T) {
 
 	r := &Runner{uuid: "runner-1", cfg: cfg, labels: labels.Labels{kubernetesLabel}, now: func() time.Time { return now }}
 	r.isolatedCacheContainer = func() string { return "a1b2c3d4e5f6" }
-	r.OnIdle(context.Background())
+	r.OnIdle(t.Context())
 	assert.Equal(t, []string{"runner-1"}, swept)
 	assert.Equal(t, swept, sweptVolumes)
 	// a network of a job starting during the pass is younger than this and so out of scope
@@ -358,13 +358,13 @@ func TestRunnerOnIdleRemovesOrphanNetworks(t *testing.T) {
 	t.Cleanup(func() { dockerReachable = origDockerReachable })
 	hostOnly := &Runner{uuid: "runner-2", cfg: &config.Config{Runner: cfg.Runner}, now: func() time.Time { return now }}
 	hostOnly.isolatedCacheContainer = func() string { return "" }
-	hostOnly.OnIdle(context.Background())
+	hostOnly.OnIdle(t.Context())
 	assert.Equal(t, []string{"runner-1"}, swept)
 	assert.Equal(t, swept, sweptVolumes)
 
 	dockerReachable = func(context.Context) bool { return true }
 	now = now.Add(time.Minute)
-	hostOnly.OnIdle(context.Background())
+	hostOnly.OnIdle(t.Context())
 	assert.Equal(t, []string{"runner-1", "runner-2"}, sweptVolumes)
 	assert.Equal(t, []string{"a1b2c3d4e5f6", ""}, detached)
 	assert.Equal(t, []string{"runner-1"}, sweptPods)

internal/app/run/runner_test.go

@@ -47,7 +47,7 @@ func TestRunnerCapabilitiesAndDeclare(t *testing.T) {
 	})).Return(connect.NewResponse(&runnerv1.DeclareResponse{}), nil)
 
 	r := &Runner{client: cli}
-	_, err := r.Declare(context.Background(), []string{"ubuntu"})
+	_, err := r.Declare(t.Context(), []string{"ubuntu"})
 	require.NoError(t, err)
 }
 

internal/pkg/envcheck/docker_test.go

@@ -11,7 +11,7 @@ import (
 )
 
 func TestCheckIfDockerRunningReturnsPingError(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	cancel()
 
 	err := CheckIfDockerRunning(ctx, "unix:///definitely/missing/docker.sock")

internal/pkg/metrics/metrics_test.go

@@ -96,7 +96,7 @@ func TestRegisterRunningJobsFuncZeroCapacity(t *testing.T) {
 }
 
 func TestStartServerCanBeCancelled(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	StartServer(ctx, "127.0.0.1:0", nil)
 	cancel()
 }

internal/pkg/process/leftover_windows_test.go

@@ -4,7 +4,6 @@
 package process
 
 import (
-	"context"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -70,7 +69,7 @@ func TestKillProcessesWithCWDUnderTerminatesMatch(t *testing.T) {
 	cmd := startSleeperIn(t, dir)
 	awaitCWD(t, cmd.Process.Pid)
 
-	killed, err := KillProcessesWithCWDUnder(context.Background(), []string{dir})
+	killed, err := KillProcessesWithCWDUnder(t.Context(), []string{dir})
 	require.NoError(t, err)
 	require.GreaterOrEqual(t, killed, 1)
 
@@ -85,7 +84,7 @@ func TestKillProcessesWithCWDUnderSparesUnrelated(t *testing.T) {
 	cmd := startSleeperIn(t, elsewhere)
 	awaitCWD(t, cmd.Process.Pid)
 
-	_, err := KillProcessesWithCWDUnder(context.Background(), []string{workspace})
+	_, err := KillProcessesWithCWDUnder(t.Context(), []string{workspace})
 	require.NoError(t, err)
 	require.True(t, processAlive(cmd.Process.Pid), "process outside owned dirs must be spared")
 }

internal/pkg/process/treekill_test.go

@@ -4,7 +4,6 @@
 package process
 
 import (
-	"context"
 	"os/exec"
 	"testing"
 
@@ -12,7 +11,7 @@ import (
 )
 
 func TestNewTreeKillConfiguresCommand(t *testing.T) {
-	cmd := exec.CommandContext(context.Background(), "sleep", "1")
+	cmd := exec.CommandContext(t.Context(), "sleep", "1")
 	tk := NewTreeKill(cmd)
 
 	require.NotNil(t, tk)
@@ -22,7 +21,7 @@ func TestNewTreeKillConfiguresCommand(t *testing.T) {
 }
 
 func TestTreeKillCaptureStoresKiller(t *testing.T) {
-	cmd := exec.CommandContext(context.Background(), "sleep", "10")
+	cmd := exec.CommandContext(t.Context(), "sleep", "10")
 	cmd.SysProcAttr = SysProcAttr("", false)
 	tk := NewTreeKill(cmd)
 	require.NoError(t, cmd.Start())

internal/pkg/report/joblog_test.go

@@ -43,7 +43,7 @@ func TestJobLog_MirrorsUploadedRows(t *testing.T) {
 	require.NoError(t, err)
 	cfg.Log.Job.Dir = t.TempDir()
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	t.Cleanup(cancel)
 	task := &runnerv1.Task{Id: 41, Context: &structpb.Struct{}, Secrets: map[string]string{"TOKEN": "s3cret-value"}}
 	reporter := NewReporter(ctx, cancel, client, task, cfg)

internal/pkg/report/reporter_test.go

@@ -244,7 +244,7 @@ func TestReporter_Handle(t *testing.T) {
 			t.Logf("Received UpdateTask: %s", req.Msg.String())
 			return connect_go.NewResponse(&runnerv1.UpdateTaskResponse{}), nil
 		})
-		ctx, cancel := context.WithCancel(context.Background())
+		ctx, cancel := context.WithCancel(t.Context())
 		taskCtx, err := structpb.NewStruct(map[string]any{})
 		require.NoError(t, err)
 		cfg, _ := config.LoadDefault("")
@@ -305,7 +305,7 @@ func TestReporter_LogLevelFiltering(t *testing.T) {
 		return connect_go.NewResponse(&runnerv1.UpdateTaskResponse{}), nil
 	})
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
 	cfg, _ := config.LoadDefault("")
@@ -374,7 +374,7 @@ func TestReporter_EphemeralRunnerDeletion(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -439,7 +439,7 @@ func TestReporter_RunDaemonClose_Race(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
 	cfg, _ := config.LoadDefault("")
@@ -481,7 +481,7 @@ func TestReporter_MaxLatencyTimer(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -532,7 +532,7 @@ func TestReporter_BatchSizeFlush(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -585,7 +585,7 @@ func TestReporter_StateChangedNotLostDuringReport(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -646,7 +646,7 @@ func TestReporter_StateChangedRestoredOnError(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -691,7 +691,7 @@ func TestReporter_StateNotifyFlush(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -748,7 +748,7 @@ func TestReporter_CloseAlwaysSendsRowsWithNoMore(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -799,7 +799,7 @@ func TestReporter_StateHeartbeat(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -871,7 +871,7 @@ func TestReporter_ServerCancelStillFlushesFinal(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -940,7 +940,7 @@ func TestReporter_CloseReportsCancelledOnCanceledCtx(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
 	cfg, _ := config.LoadDefault("")
@@ -999,7 +999,7 @@ func TestReporter_StopHeartbeats(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -1112,7 +1112,7 @@ func TestReporter_DebugSettingsAreNotMasked(t *testing.T) {
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
 
-	reporter := NewReporter(context.Background(), nil, nil, &runnerv1.Task{
+	reporter := NewReporter(t.Context(), nil, nil, &runnerv1.Task{
 		Context: taskCtx,
 		Secrets: map[string]string{"ACTIONS_STEP_DEBUG": "true", "ACTIONS_RUNNER_DEBUG": "true", "TOKEN": "s3cr3t"},
 	}, &config.Config{})
@@ -1136,7 +1136,7 @@ func TestReporter_OutputsSentOnce(t *testing.T) {
 			return connect_go.NewResponse(&runnerv1.UpdateTaskResponse{SentOutputs: slices.Collect(maps.Keys(req.Msg.Outputs))}), nil
 		})
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -1210,7 +1210,7 @@ func TestReporter_AckIndexBeyondBuffer(t *testing.T) {
 		return connect_go.NewResponse(&runnerv1.UpdateLogResponse{AckIndex: 1000}), nil
 	})
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -1247,7 +1247,7 @@ func TestReporter_CloseWithStuckDaemon(t *testing.T) {
 		return connect_go.NewResponse(&runnerv1.UpdateTaskResponse{}), nil
 	}).Maybe()
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -1294,7 +1294,7 @@ func TestReporter_ServerFailureResultIsReportedAsFailure(t *testing.T) {
 		}), nil
 	}).Maybe()
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -1351,7 +1351,7 @@ func TestReporter_CloseMasksLastWords(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	const idToken = "id-token-request-secret"
 	taskCtx, err := structpb.NewStruct(map[string]any{"actions_id_token_request_token": idToken})
@@ -1420,7 +1420,7 @@ func TestReporter_StepRangeCoversAckedRows(t *testing.T) {
 		},
 	)
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)
@@ -1466,7 +1466,7 @@ func TestReporter_StepRangeCoversAckedRows(t *testing.T) {
 // A flush the server took no rows for describes nothing new, so it must not spend an
 // UpdateTask. UpdateTask is never registered on the client, so calling it fails the test.
 func TestReporter_FlushWithoutAckedRowsSkipsState(t *testing.T) {
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 	taskCtx, err := structpb.NewStruct(map[string]any{})
 	require.NoError(t, err)