From 8a578de9a9b4ac72e5d5fb99761a472a8a9d3512 Mon Sep 17 00:00:00 2001
From: Jan Rose <jan.rose@databricks.com>
Date: Wed, 2 Sep 2026 22:26:45 +0000
Subject: [PATCH] Fix race condition in TestJobRunWaitAbandonedLinksTheRun

---
 bundle/direct/dresources/job_run_test.go | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/bundle/direct/dresources/job_run_test.go b/bundle/direct/dresources/job_run_test.go
index 0ffc1c48040..a2efdc0e6d3 100644
--- a/bundle/direct/dresources/job_run_test.go
+++ b/bundle/direct/dresources/job_run_test.go
@@ -8,7 +8,6 @@ import (
 	"strings"
 	"sync/atomic"
 	"testing"
-	"time"
 
 	"github.com/databricks/cli/libs/structs/structpath"
 	"github.com/databricks/cli/libs/testserver"
@@ -232,11 +231,23 @@ func TestJobRunWaitReportsOnlyTheLastAttemptOfATask(t *testing.T) {
 }
 
 func TestJobRunWaitAbandonedLinksTheRun(t *testing.T) {
-	client := jobRunClient(t, &jobs.RunState{LifeCycleState: jobs.RunLifeCycleStateRunning})
-
-	ctx, cancel := context.WithTimeout(t.Context(), time.Millisecond)
+	ctx, cancel := context.WithCancel(t.Context())
 	defer cancel()
 
+	// The first poll reports the run page and finds the run still going; the
+	// second cancels the wait. Driving the interrupt from the handler keeps it
+	// deterministic: the run page URL is always captured before the wait ends,
+	// instead of racing a wall-clock timeout against the first GetRun.
+	var gets atomic.Int32
+	client := jobRunServer(t, func(req testserver.Request) any {
+		if gets.Add(1) >= 2 {
+			cancel()
+		}
+		return jobs.Run{RunId: 123, JobId: 456, State: &jobs.RunState{
+			LifeCycleState: jobs.RunLifeCycleStateRunning,
+		}, RunPageUrl: testRunPageURL}
+	})
+
 	_, err := waitForTestRun(t, ctx, client)
 
 	// The run keeps going, so the error links to it and names the interrupt.