Fix waitUntil, which was broken in the new implementation of polling expectations

This one is wild, because apparently, RunLoop.run(mode:before:) wasn't actually blocking for the amount of time specified. It was basically spinning the runloop once, then returning.
This change now continuously spins the runloop until either the timeout passes, or a result is set.

Author: younata

Sources/Nimble/Utils/AsyncAwait.swift

@@ -44,7 +44,7 @@ internal enum AsyncPollResult<T> {
     }
 }
 
-final class BlockingTask: Sendable {
+private final class BlockingTask: Sendable {
     private nonisolated(unsafe) var finished = false
     private nonisolated(unsafe) var continuation: CheckedContinuation<Void, Never>? = nil
     let sourceLocation: SourceLocation

Sources/Nimble/Utils/PollAwait.swift

@@ -81,6 +81,30 @@ internal enum PollStatus {
     case incomplete
 }
 
+private final class LockBox<T>: Sendable {
+    private nonisolated(unsafe) var value: T
+
+    private let lock = NSLock()
+
+    init(value: T) {
+        self.value = value
+    }
+
+    var currentValue: T {
+        lock.lock()
+        defer {
+            lock.unlock()
+        }
+        return value
+    }
+
+    func set(_ newValue: T) {
+        lock.lock()
+        value = newValue
+        lock.unlock()
+    }
+}
+
 func synchronousWaitUntil(
     timeout: NimbleTimeInterval,
     fnName: String,
@@ -100,23 +124,20 @@ Please use Swift Testing's `confirmation(...)` APIs to accomplish (nearly) the s
     return guaranteeNotNested(fnName: fnName, sourceLocation: sourceLocation) {
         let runloop = RunLoop.current
 
-        nonisolated(unsafe) var result = PollResult<Void>.timedOut
-        let lock = NSLock()
+        let resultBox = LockBox(value: PollResult<Void>.timedOut)
 
         let doneBlock: () -> Void = {
             let onFinish = {
-                lock.lock()
-                defer { lock.unlock() }
-                if case .completed = result {
+                if case .completed = resultBox.currentValue {
                     fail("waitUntil(...) expects its completion closure to be only called once", location: sourceLocation)
                     return
                 }
+                resultBox.set(.completed(()))
 #if canImport(CoreFoundation)
                 CFRunLoopStop(CFRunLoopGetCurrent())
 #else
                 RunLoop.main._stop()
 #endif
-                result = .completed(())
             }
             if Thread.isMainThread {
                 onFinish()
@@ -127,33 +148,30 @@ Please use Swift Testing's `confirmation(...)` APIs to accomplish (nearly) the s
 
         let capture = NMBExceptionCapture(
             handler: ({ exception in
-                lock.lock()
-                defer { lock.unlock() }
-                result = .raisedException(exception)
+                resultBox.set(.raisedException(exception))
             }),
             finally: ({ })
         )
         capture.tryBlock {
             do {
                 try closure(doneBlock)
             } catch {
-                lock.lock()
-                defer { lock.unlock() }
-                result = .errorThrown(error)
+                resultBox.set(.errorThrown(error))
             }
         }
 
-        if Thread.isMainThread {
-            runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
-        } else {
-            DispatchQueue.main.sync {
-                _ = runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
+        let start = Date()
+        while case .timedOut = resultBox.currentValue, abs(start.timeIntervalSinceNow) < timeout.timeInterval {
+            if Thread.isMainThread {
+                runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
+            } else {
+                DispatchQueue.main.sync {
+                    _ = runloop.run(mode: .default, before: Date(timeIntervalSinceNow: timeout.timeInterval))
+                }
             }
         }
 
-        lock.lock()
-        defer { lock.unlock() }
-        return result
+        return resultBox.currentValue
     }
 }
 

Tests/NimbleTests/AsyncAwaitTest.swift

@@ -298,6 +298,74 @@ final class AsyncAwaitTest: XCTestCase { // swiftlint:disable:this type_body_len
         timer.cancel()
     }
 
+    func testWaitUntilNestedMainQueueAsyncCalls() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            DispatchQueue.main.async {
+                DispatchQueue.main.async {
+                    DispatchQueue.main.async {
+                        done()
+                    }
+                }
+            }
+        }
+    }
+
+    func testWaitUntilOperationQueueWithMainUnderlyingQueueAndBarrier() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = .main
+
+            let operation = BlockOperation {}
+
+            operationQueue.addOperation(operation)
+            operationQueue.addBarrierBlock {
+                done()
+            }
+        }
+    }
+
+    func testNestedMainAsyncWithOperationQueue() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue.main
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    func testNestedBackgroundAsyncWithOperationQueue() async {
+        await waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue(label: "test.queue")
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     final class ClassUnderTest {
         var deinitCalled: (() -> Void)?
         var count = 0

Tests/NimbleTests/PollingTest.swift

@@ -219,6 +219,74 @@ final class PollingTest: XCTestCase {
 #endif // canImport(Darwin)
     }
 
+    func testWaitUntilNestedMainQueueAsyncCalls() {
+        waitUntil(timeout: .seconds(1)) { done in
+            DispatchQueue.main.async {
+                DispatchQueue.main.async {
+                    DispatchQueue.main.async {
+                        done()
+                    }
+                }
+            }
+        }
+    }
+
+    func testWaitUntilOperationQueueWithMainUnderlyingQueueAndBarrier() {
+        waitUntil(timeout: .seconds(1)) { done in
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = .main
+
+            let operation = BlockOperation {}
+
+            operationQueue.addOperation(operation)
+            operationQueue.addBarrierBlock {
+                done()
+            }
+        }
+    }
+
+    func testNestedMainAsyncWithOperationQueue() {
+        waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue.main
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    func testNestedBackgroundAsyncWithOperationQueue() {
+        waitUntil(timeout: .seconds(1)) { done in
+            let indexingQueue = DispatchQueue(label: "test.queue")
+            let operationQueue = OperationQueue()
+            operationQueue.underlyingQueue = indexingQueue
+
+            indexingQueue.async {
+                indexingQueue.async {
+                    indexingQueue.async {
+                        let operation = BlockOperation {}
+
+                        operationQueue.addOperation(operation)
+                        operationQueue.addBarrierBlock {
+                            done()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
     func testWaitUntilAllowsInBackgroundThread() {
 #if !SWIFT_PACKAGE
         var executedAsyncBlock: Bool = false