diff --git a/workspace-server/internal/handlers/a2a_proxy_test.go b/workspace-server/internal/handlers/a2a_proxy_test.go index 29be97a6..0bf84a15 100644 --- a/workspace-server/internal/handlers/a2a_proxy_test.go +++ b/workspace-server/internal/handlers/a2a_proxy_test.go @@ -13,6 +13,7 @@ import ( "os" "os/exec" "strings" + "sync" "testing" "time" @@ -309,8 +310,8 @@ func TestProxyA2A_Upstream502_TriggersContainerDeadCheck(t *testing.T) { if w.Header().Get("Retry-After") != "15" { t.Errorf("Retry-After header should be 15 to throttle canvas-side retry loop; got %q", w.Header().Get("Retry-After")) } - if cp.calls != 1 { - t.Errorf("cpProv.IsRunning must be consulted exactly once; got %d calls", cp.calls) + if cp.Calls() != 1 { + t.Errorf("cpProv.IsRunning must be consulted exactly once; got %d calls", cp.Calls()) } } @@ -2199,8 +2200,8 @@ func TestMaybeMarkContainerDead_CPOnly_NotRunning(t *testing.T) { if !dead { t.Fatal("expected true (cpProv reports not running) — without cpProv consultation, SaaS dead-agent recovery is impossible") } - if cp.calls != 1 { - t.Errorf("expected exactly 1 IsRunning call on cpProv; got %d", cp.calls) + if cp.Calls() != 1 { + t.Errorf("expected exactly 1 IsRunning call on cpProv; got %d", cp.Calls()) } if err := mock.ExpectationsWereMet(); err != nil { t.Errorf("unmet sqlmock expectations: %v", err) @@ -2228,8 +2229,8 @@ func TestMaybeMarkContainerDead_CPOnly_Running(t *testing.T) { if dead { t.Error("expected false when cpProv reports running — must not recycle a healthy agent") } - if cp.calls != 1 { - t.Errorf("expected exactly 1 IsRunning call on cpProv; got %d", cp.calls) + if cp.Calls() != 1 { + t.Errorf("expected exactly 1 IsRunning call on cpProv; got %d", cp.Calls()) } } @@ -2259,11 +2260,11 @@ func TestStopForRestart_SaaSPath_DispatchesViaCPProv(t *testing.T) { handler.stopForRestart(context.Background(), "ws-saas-restart") - if cp.stopCalls != 1 { - t.Fatalf("expected cpProv.Stop to be called once on SaaS auto-restart; got %d", cp.stopCalls) + if cp.StopCalls() != 1 { + t.Fatalf("expected cpProv.Stop to be called once on SaaS auto-restart; got %d", cp.StopCalls()) } - if cp.startCalls != 0 { - t.Fatalf("expected cpProv.Start NOT to be called by stopForRestart; got %d", cp.startCalls) + if cp.StartCalls() != 0 { + t.Fatalf("expected cpProv.Start NOT to be called by stopForRestart; got %d", cp.StartCalls()) } } @@ -2294,6 +2295,7 @@ func TestStopForRestart_NoProvisioner_NoOp(t *testing.T) { // after the assertions ran). Tests that want to ASSERT a method is unused // can check `calls == 0` after a sync barrier. type fakeCPProv struct { + mu sync.Mutex running bool err error calls int @@ -2301,15 +2303,45 @@ type fakeCPProv struct { startCalls int } +func (f *fakeCPProv) setRunning(v bool) { + f.mu.Lock() + defer f.mu.Unlock() + f.running = v +} + +func (f *fakeCPProv) Calls() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.calls +} + +func (f *fakeCPProv) StopCalls() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.stopCalls +} + +func (f *fakeCPProv) StartCalls() int { + f.mu.Lock() + defer f.mu.Unlock() + return f.startCalls +} + func (f *fakeCPProv) Start(_ context.Context, _ provisioner.WorkspaceConfig) (string, error) { + f.mu.Lock() + defer f.mu.Unlock() f.startCalls++ return "", nil } func (f *fakeCPProv) Stop(_ context.Context, _ string) error { + f.mu.Lock() + defer f.mu.Unlock() f.stopCalls++ return nil } func (f *fakeCPProv) StopAndPrune(_ context.Context, _ string) error { + f.mu.Lock() + defer f.mu.Unlock() f.stopCalls++ return nil } @@ -2317,6 +2349,8 @@ func (f *fakeCPProv) GetConsoleOutput(_ context.Context, _ string) (string, erro return "", nil } func (f *fakeCPProv) IsRunning(_ context.Context, _ string) (bool, error) { + f.mu.Lock() + defer f.mu.Unlock() f.calls++ return f.running, f.err } @@ -2369,8 +2403,8 @@ func TestMaybeMarkContainerDead_RecentHeartbeat_DoesNotRestart(t *testing.T) { t.Fatal("dead observation with recent heartbeat should not declare container dead") } // It re-probes, so two IsRunning calls even with delay=0. - if cp.calls != 2 { - t.Errorf("expected 2 IsRunning calls (probe + re-probe), got %d", cp.calls) + if cp.Calls() != 2 { + t.Errorf("expected 2 IsRunning calls (probe + re-probe), got %d", cp.Calls()) } // If EnqueueA2A fails (no DB expectations here) it returns status=0 and // lets the caller fall back to its normal error path. The critical @@ -2403,8 +2437,8 @@ func TestMaybeMarkContainerDead_NoRecentHeartbeat_DeclaresDead(t *testing.T) { if !dead { t.Fatal("expected dead when there is no recent heartbeat and IsRunning=false") } - if cp.calls != 1 { - t.Errorf("expected 1 IsRunning call, got %d", cp.calls) + if cp.Calls() != 1 { + t.Errorf("expected 1 IsRunning call, got %d", cp.Calls()) } } @@ -2453,8 +2487,8 @@ func TestMaybeMarkContainerDead_SettleWindow_DoesNotClearURL(t *testing.T) { if dead { t.Fatal("single IsRunning=false in post-restart settle window must not declare container dead") } - if cp.calls != 2 { - t.Errorf("expected 2 IsRunning calls (probe + re-probe), got %d", cp.calls) + if cp.Calls() != 2 { + t.Errorf("expected 2 IsRunning calls (probe + re-probe), got %d", cp.Calls()) } } @@ -2511,7 +2545,7 @@ func TestMaybeMarkContainerDead_RunningTrueAfterReprobe_Resets(t *testing.T) { } // Next observation says running=true → no restart. - cp.running = true + cp.setRunning(true) mock.ExpectQuery(`SELECT COALESCE\(runtime, 'claude-code'\) FROM workspaces WHERE id =`). WithArgs(wsid). WillReturnRows(sqlmock.NewRows([]string{"runtime"}).AddRow("hermes"))