workspace-server/internal/handlers/a2a_proxy_test.go

@@ -294,8 +294,7 @@ func TestProxyA2A_Upstream502_TriggersContainerDeadCheck(t *testing.T) {
 	c.Request.Header.Set("Content-Type", "application/json")
 
 	handler.ProxyA2A(c)
-
-	time.Sleep(80 * time.Millisecond)
+	handler.waitAsyncForTest()
 
 	// Caller sees a structured 503 (NOT the upstream 502 which CF would mask).
 	if w.Code != http.StatusServiceUnavailable {
@@ -350,7 +349,7 @@ func TestProxyA2A_Upstream502_AliveAgent_PropagatesAsIs(t *testing.T) {
 	c.Request.Header.Set("Content-Type", "application/json")
 
 	handler.ProxyA2A(c)
-	time.Sleep(50 * time.Millisecond)
+	handler.waitAsyncForTest()
 
 	if w.Code != http.StatusBadGateway {
 		t.Fatalf("alive agent 502 should propagate as 502; got %d: %s", w.Code, w.Body.String())

workspace-server/internal/handlers/restart_signals_test.go

@@ -273,7 +273,7 @@ func TestGracefulPreRestart_URLResolutionError(t *testing.T) {
 	}
 
 	hWrapper.gracefulPreRestart(context.Background(), "ws-url-err-111")
-	time.Sleep(200 * time.Millisecond)
+	hWrapper.waitAsyncForTest()
 	// No panic or error expected — proceeds with stop as documented
 }
 

workspace-server/internal/handlers/workspace_provision_auto_test.go

@@ -686,7 +686,7 @@ func TestRestartWorkspaceAuto_RoutesToDockerWhenOnlyDocker(t *testing.T) {
 	// recovered by logProvisionPanic. Without this wait, the goroutine
 	// outlives the test and writes to a sqlmock that the NEXT test
 	// owns, causing a `was not expected` race.
-	time.Sleep(200 * time.Millisecond)
+	h.waitAsyncForTest()
 
 	// Stop call is synchronous on the Docker leg.
 	if len(stub.stopped) == 0 || stub.stopped[0] != wsID {