From 8bee99d7f267530e097ba60cdbb9c00442e1395c Mon Sep 17 00:00:00 2001
From: Ngo Quoc Dat <datlechin@gmail.com>
Date: Fri, 18 Sep 2026 13:58:29 +0700
Subject: [PATCH] test(plugins): wait on the stream's producer instead of on
 the clock

---
 .../Database/PluginStreamAbortTests.swift     | 47 ++++++++++++++++---
 1 file changed, 41 insertions(+), 6 deletions(-)

diff --git a/TableProTests/Core/Database/PluginStreamAbortTests.swift b/TableProTests/Core/Database/PluginStreamAbortTests.swift
index 0f51c5f3c..12089bffd 100644
--- a/TableProTests/Core/Database/PluginStreamAbortTests.swift
+++ b/TableProTests/Core/Database/PluginStreamAbortTests.swift
@@ -87,10 +87,14 @@ struct PluginStreamAbortTests {
         #expect(observed.value == true)
     }
 
-    @Test("A stream still held in scope has not terminated, however long the consumer has stopped")
+    @Test(
+        "A stream still held in scope has not terminated, however long the consumer has stopped",
+        .timeLimit(.minutes(1))
+    )
     func holdingTheStreamKeepsItAlive() async throws {
         let queue = DispatchQueue(label: "test.stream.abort.held")
         let counter = ProducedRowCounter()
+        let producerFinished = ProducerCompletion()
 
         let stream = PluginRowStream.make { continuation, abort in
             queue.async {
@@ -101,6 +105,7 @@ struct PluginStreamAbortTests {
                     usleep(20_000)
                 }
                 continuation.finish()
+                producerFinished.signal()
             }
         }
 
@@ -110,11 +115,11 @@ struct PluginStreamAbortTests {
             if seen >= 2 { break }
         }
 
-        /// Waited for rather than slept through. The producer paces itself at 20ms a row, so a
-        /// fixed wait asserts the runner's load as much as the stream's behaviour: CI saw 19 of 20.
-        for _ in 0 ..< 100 where counter.value < 20 {
-            try await Task.sleep(for: .milliseconds(50))
-        }
+        /// Waited on the producer rather than on the clock. It paces itself at 20ms a row, so a
+        /// fixed sleep asserts the runner's load as much as the stream's behaviour: CI saw 19 of 20.
+        /// An abort would leave the loop early, and the signal still arrives, so the count is what
+        /// says whether the stream stayed alive. The time limit is what fails a producer that hangs.
+        await producerFinished.wait()
         #expect(counter.value == 20)
         _ = stream
     }
@@ -156,6 +161,36 @@ private final class ProducedRowCounter: @unchecked Sendable {
     }
 }
 
+/// One shot, and safe to await after the fact: a producer that finished before the consumer asked
+/// resumes the waiter immediately rather than leaving it there.
+private final class ProducerCompletion: @unchecked Sendable {
+    private let lock = NSLock()
+    private var hasFinished = false
+    private var waiter: CheckedContinuation<Void, Never>?
+
+    func signal() {
+        lock.lock()
+        hasFinished = true
+        let pending = waiter
+        waiter = nil
+        lock.unlock()
+        pending?.resume()
+    }
+
+    func wait() async {
+        await withCheckedContinuation { continuation in
+            lock.lock()
+            if hasFinished {
+                lock.unlock()
+                continuation.resume()
+                return
+            }
+            waiter = continuation
+            lock.unlock()
+        }
+    }
+}
+
 private final class AbortObservation: @unchecked Sendable {
     private let lock = NSLock()
     private var observed: Bool?