From 9a2a01b34c226873ff46bd1aff1633793f6337e2 Mon Sep 17 00:00:00 2001
From: numinousmuses <103385201+numinousmuses@users.noreply.github.com>
Date: Sun, 6 Sep 2026 01:54:18 +0000
Subject: [PATCH] Make provider test fixtures wait for their lifecycle

---
 coordinator/api/failover_integration_test.go | 68 +++++++++++++++++++-
 coordinator/api/me_delete_test.go            |  9 +--
 2 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/coordinator/api/failover_integration_test.go b/coordinator/api/failover_integration_test.go
index 1df0dc750..0686cb11b 100644
--- a/coordinator/api/failover_integration_test.go
+++ b/coordinator/api/failover_integration_test.go
@@ -108,6 +108,7 @@ type failoverProvider struct {
 	dispatches atomic.Int32
 	bodies     chan []byte
 	done       chan struct{}
+	cancel     context.CancelFunc
 	closeOnce  sync.Once
 }
 
@@ -205,6 +206,7 @@ func startFailoverProvider(t *testing.T, ctx context.Context, ts *httptest.Serve
 	reg.SetTrustLevel(registryID, registry.TrustHardware)
 	reg.RecordChallengeSuccess(registryID)
 
+	runCtx, runCancel := context.WithCancel(ctx)
 	fp := &failoverProvider{
 		t:          t,
 		name:       cfg.Name,
@@ -215,8 +217,9 @@ func startFailoverProvider(t *testing.T, ctx context.Context, ts *httptest.Serve
 		script:     cfg.Script,
 		bodies:     make(chan []byte, 8),
 		done:       make(chan struct{}),
+		cancel:     runCancel,
 	}
-	go fp.run(ctx)
+	go fp.run(runCtx)
 	t.Cleanup(fp.close)
 	return fp
 }
@@ -285,9 +288,17 @@ func (fp *failoverProvider) dispatchCount() int {
 
 // close shuts the provider WebSocket down (idempotent; safe in t.Cleanup).
 func (fp *failoverProvider) close() {
+	fp.cancel()
 	fp.closeOnce.Do(func() {
 		_ = fp.conn.Close(websocket.StatusNormalClosure, "test done")
 	})
+	// A script can still log through testing.T after the socket closes.
+	// Join it before test cleanup returns, including after an earlier closeNow.
+	select {
+	case <-fp.done:
+	case <-time.After(5 * time.Second):
+		fp.t.Errorf("provider %s: read loop did not exit during cleanup", fp.name)
+	}
 }
 
 // closeNow abruptly drops the provider connection, simulating a crash /
@@ -977,3 +988,58 @@ func TestProviderClientError400_StopsAfterOne(t *testing.T) {
 		t.Errorf("body should surface invalid_request_error; got %s", body)
 	}
 }
+
+func TestFailoverProviderCleanupJoinsActiveScript(t *testing.T) {
+	reg, _, ts := setupFailoverServer(t)
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
+	entered := make(chan struct{})
+	release := make(chan struct{})
+	var releaseOnce sync.Once
+	releaseScript := func() { releaseOnce.Do(func() { close(release) }) }
+	defer releaseScript()
+	const model = "cleanup-model"
+	fp := startFailoverProvider(t, ctx, ts, reg, failoverProviderConfig{
+		Name: "cleanup-provider", Version: "0.6.4", DecodeTPS: 100,
+		Models: []failoverModelSpec{{ID: model}},
+		Script: func(context.Context, *failoverProvider, protocol.InferenceRequestMessage, []byte) {
+			close(entered)
+			<-release
+		},
+	})
+	clientDone := make(chan struct{})
+	body := buildChatBody(t, model, true, nil)
+	go func() {
+		defer close(clientDone)
+		_, _, _ = postChat(ctx, ts.URL, "test-key", body)
+	}()
+	select {
+	case <-entered:
+	case <-time.After(5 * time.Second):
+		t.Fatal("provider did not enter its script")
+	}
+	closed := make(chan struct{})
+	go func() { fp.close(); close(closed) }()
+	select {
+	case <-closed:
+		t.Fatal("provider cleanup returned while its script could still access testing.T")
+	case <-time.After(100 * time.Millisecond):
+	}
+	releaseScript()
+	select {
+	case <-closed:
+	case <-time.After(5 * time.Second):
+		t.Fatal("provider cleanup did not finish after its script exited")
+	}
+	select {
+	case <-fp.done:
+	default:
+		t.Fatal("provider cleanup returned before the read loop exited")
+	}
+	cancel()
+	select {
+	case <-clientDone:
+	case <-time.After(5 * time.Second):
+		t.Fatal("consumer request did not exit after cancellation")
+	}
+}
diff --git a/coordinator/api/me_delete_test.go b/coordinator/api/me_delete_test.go
index b5ab763f5..ea72d0684 100644
--- a/coordinator/api/me_delete_test.go
+++ b/coordinator/api/me_delete_test.go
@@ -110,14 +110,15 @@ func TestDeleteMyProvider_NotFound404(t *testing.T) {
 
 func TestDeleteMyProvider_OnlineConflict409(t *testing.T) {
 	srv, st := newKeyTestServer(t)
-	seedProviderRecord(t, st, "live-p", "SER-ON", "acct-1")
+	seedProviderRecord(t, st, "stored-p", "SER-ON", "acct-1")
 
-	// Register a live provider connection with a matching serial.
+	// Reconnecting sessions have distinct IDs. The live registration must not
+	// overwrite the owned record with its asynchronous pre-link snapshot.
 	live := srv.registry.Register("live-p", nil, &protocol.RegisterMessage{})
 	live.SetAttestationResult(&attestation.VerificationResult{SerialNumber: "SER-ON"})
 
-	r := reqWithUser(http.MethodDelete, "/v1/me/providers/live-p", "", "acct-1")
-	r.SetPathValue("id", "live-p")
+	r := reqWithUser(http.MethodDelete, "/v1/me/providers/stored-p", "", "acct-1")
+	r.SetPathValue("id", "stored-p")
 	w := httptest.NewRecorder()
 	srv.handleDeleteMyProvider(w, r)