From 9bf47b46e441cdf984e0773fa98aab2d14286641 Mon Sep 17 00:00:00 2001 From: joshuaaferguson Date: Sat, 25 Apr 2026 11:00:34 -0700 Subject: [PATCH] fix(websocket): race in TestUpdateAgentHeartbeat between sqlmock writes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run() executes inside a goroutine and calls handleRegister, which issues a sqlmock-backed Exec for the registration UPDATE. The test then calls mock.ExpectExec to declare the heartbeat expectation. go-sqlmock <=1.5.2 does not synchronize ExpectExec writes against exec reads on its internal expected[] slice — the original time.Sleep(100ms) was not enough to guarantee the registration Exec had finished, so -race intermittently flagged a write-vs-read race. Replace the sleep with a deadline poll on mock.ExpectationsWereMet(). The poll completes only after the registration UPDATE has been consumed, so subsequent ExpectExec runs without a concurrent reader. Verified with `go test -race -count=20`: 20/20 clean (was failing ~50% before). --- api/internal/websocket/agent_hub_test.go | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/api/internal/websocket/agent_hub_test.go b/api/internal/websocket/agent_hub_test.go index 173acf74..1e94f5d0 100644 --- a/api/internal/websocket/agent_hub_test.go +++ b/api/internal/websocket/agent_hub_test.go @@ -226,7 +226,22 @@ func TestUpdateAgentHeartbeat(t *testing.T) { t.Fatalf("Failed to register agent: %v", err) } - time.Sleep(100 * time.Millisecond) + // Wait for the registration UPDATE to be consumed by the Run() goroutine + // before declaring the next expectation. Calling mock.ExpectExec while + // the registration goroutine is still inside sqlmock.exec races on + // sqlmock's internal expected[] slice (go-sqlmock <=1.5.2 is not safe + // for concurrent ExpectExec/exec). + deadline := time.After(2 * time.Second) + for { + if mock.ExpectationsWereMet() == nil { + break + } + select { + case <-deadline: + t.Fatal("timeout waiting for registration UPDATE to be consumed") + case <-time.After(5 * time.Millisecond): + } + } // Mock database update for heartbeat (includes status update) // Note: Query uses $1 for both last_heartbeat and updated_at, so only 2 args (timestamp, agentID)