Add a new testing interface - wait(for:timeout:until) (#80)

* Add a new testing functionality wait(until:)

* Update interface to wait for an specified atom

* Update documentation

* Remove unnecessary @mainactor

README.md

@@ -1061,7 +1061,8 @@ A context that can simulate any scenarios in which atoms are used from a view or
 |:--|:--|
 |[unwatch(_:)](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/atomtestcontext/unwatch(_:))|Simulates a scenario in which the atom is no longer watched.|
 |[override(_:with:)](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/atomtestcontext/override(_:with:)-1ce4h)|Overwrites the output of a specific atom or all atoms of the given type with the fixed value.|
-|[waitForUpdate(timeout:)](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/atomtestcontext/waitforupdate(timeout:))|Waits until any of atoms watched through this context is updated.|
+|[waitForUpdate(timeout:)](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/atomtestcontext/waitforupdate(timeout:))|Waits until any of the atoms watched through this context have been updated.|
+|[wait(for:timeout:until:)](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/atomtestcontext/wait(for:timeout:until:))|Waits for the given atom until it will be a certain state.|
 |[onUpdate](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/atomtestcontext/onupdate)|Sets a closure that notifies there has been an update to one of the atoms.|
 
 <details><summary><code>📖 Expand to see example</code></summary>

Sources/Atoms/Context/AtomTestContext.swift

@@ -23,8 +23,8 @@ public struct AtomTestContext: AtomWatchableContext {
         nonmutating set { state.onUpdate = newValue }
     }
 
-    /// Waits until any of atoms watched through this context is updated for up to
-    /// the specified timeout, and then return a boolean value indicating whether an update is done.
+    /// Waits until any of the atoms watched through this context have been updated up to the
+    /// specified timeout, and then returns a boolean value indicating whether an update is done.
     ///
     /// ```swift
     /// func testAsyncUpdate() async {
@@ -41,45 +41,86 @@ public struct AtomTestContext: AtomWatchableContext {
     /// }
     /// ```
     ///
-    /// - Parameter interval: The maximum timeout interval that this function can wait until
-    ///                      the next update. The default timeout interval is nil.
+    /// - Parameter duration: The maximum duration that this function can wait until
+    ///                       the next update. The default timeout interval is nil.
     /// - Returns: A boolean value indicating whether an update is done.
     @discardableResult
-    public func waitForUpdate(timeout interval: TimeInterval? = nil) async -> Bool {
-        let updates = AsyncStream<Void> { continuation in
-            let cancellable = state.notifier.sink(
-                receiveCompletion: { completion in
-                    continuation.finish()
-                },
-                receiveValue: {
-                    continuation.yield()
-                }
-            )
-
-            continuation.onTermination = { termination in
-                switch termination {
-                case .cancelled:
-                    cancellable.cancel()
+    public func waitForUpdate(timeout duration: TimeInterval? = nil) async -> Bool {
+        await withTaskGroup(of: Bool.self) { group in
+            let updates = state.makeUpdateStream()
 
-                case .finished:
-                    break
+            group.addTask { @MainActor in
+                var iterator = updates.makeAsyncIterator()
+                await iterator.next()
+                return true
+            }
 
-                @unknown default:
-                    break
+            if let duration {
+                group.addTask {
+                    try? await Task.sleep(seconds: duration)
+                    return false
                 }
             }
+
+            let didUpdate = await group.next() ?? false
+            group.cancelAll()
+
+            return didUpdate
         }
+    }
 
-        return await withTaskGroup(of: Bool.self) { group in
-            group.addTask {
-                var iterator = updates.makeAsyncIterator()
-                await iterator.next()
-                return true
+    /// Waits for the given atom until it will be a certain state up to the specified timeout,
+    /// and then returns a boolean value indicating whether an update is done.
+    ///
+    /// ```swift
+    /// func testAsyncUpdate() async {
+    ///     let context = AtomTestContext()
+    ///
+    ///     let initialPhase = context.watch(AsyncCalculationAtom().phase)
+    ///     XCTAssertEqual(initialPhase, .suspending)
+    ///
+    ///     let didUpdate = await context.wait(for: AsyncCalculationAtom().phase, until: \.isSuccess)
+    ///     let currentPhase = context.watch(AsyncCalculationAtom().phase)
+    ///
+    ///     XCTAssertTure(didUpdate)
+    ///     XCTAssertEqual(currentPhase, .success(123))
+    /// }
+    /// ```
+    ///
+    /// - Parameters:
+    ///   - atom: An atom that this method waits updating to a certain state.
+    ///   - duration: The maximum duration that this function can wait until
+    ///               the next update. The default timeout interval is nil.
+    ///   - predicate: A predicate that determines when to stop waiting.
+    ///
+    /// - Returns: A boolean value indicating whether an update is done.
+    ///
+    @discardableResult
+    public func wait<Node: Atom>(
+        for atom: Node,
+        timeout duration: TimeInterval? = nil,
+        until predicate: @escaping (Node.Loader.Value) -> Bool
+    ) async -> Bool {
+        await withTaskGroup(of: Bool.self) { group in
+            let updates = state.makeUpdateStream()
+
+            group.addTask { @MainActor in
+                guard !predicate(read(atom)) else {
+                    return false
+                }
+
+                for await _ in updates {
+                    if predicate(read(atom)) {
+                        return true
+                    }
+                }
+
+                return false
             }
 
-            if let interval {
+            if let duration {
                 group.addTask {
-                    try? await Task.sleep(nanoseconds: UInt64(interval * 1_000_000_000))
+                    try? await Task.sleep(seconds: duration)
                     return false
                 }
             }
@@ -261,10 +302,37 @@ private extension AtomTestContext {
         let store = AtomStore()
         let token = ScopeKey.Token()
         let container = SubscriptionContainer()
-        let notifier = PassthroughSubject<Void, Never>()
         var overrides = [OverrideKey: any AtomOverrideProtocol]()
         var onUpdate: (() -> Void)?
 
+        private let notifier = PassthroughSubject<Void, Never>()
+
+        func makeUpdateStream() -> AsyncStream<Void> {
+            AsyncStream { continuation in
+                let cancellable = notifier.sink(
+                    receiveCompletion: { _ in
+                        continuation.finish()
+                    },
+                    receiveValue: {
+                        continuation.yield()
+                    }
+                )
+
+                continuation.onTermination = { termination in
+                    switch termination {
+                    case .cancelled:
+                        cancellable.cancel()
+
+                    case .finished:
+                        break
+
+                    @unknown default:
+                        break
+                    }
+                }
+            }
+        }
+
         func notifyUpdate() {
             onUpdate?()
             notifier.send()

Sources/Atoms/Core/TaskExtensions.swift

@@ -0,0 +1,5 @@
+internal extension Task where Success == Never, Failure == Never {
+    static func sleep(seconds duration: Double) async throws {
+        try await sleep(nanoseconds: UInt64(duration * 1_000_000_000))
+    }
+}

Tests/AtomsTests/Atom/AsyncSequenceAtomTests.swift

@@ -16,15 +16,15 @@ final class AsyncSequenceAtomTests: XCTestCase {
         do {
             // Value
             pipe.continuation.yield(0)
-            await context.waitForUpdate()
+            await context.wait(for: atom, until: \.isSuccess)
 
             XCTAssertEqual(context.watch(atom).value, 0)
         }
 
         do {
             // Failure
             pipe.continuation.finish(throwing: URLError(.badURL))
-            await context.waitForUpdate()
+            await context.wait(for: atom, until: \.isFailure)
 
             XCTAssertEqual(context.watch(atom).error as? URLError, URLError(.badURL))
         }

Tests/AtomsTests/Atom/PublisherAtomTests.swift

@@ -20,15 +20,15 @@ final class PublisherAtomTests: XCTestCase {
             // Value
             subject.send(0)
 
-            await context.waitForUpdate()
+            await context.wait(for: atom, until: \.isSuccess)
             XCTAssertEqual(context.watch(atom), .success(0))
         }
 
         do {
             // Error
             subject.send(completion: .failure(URLError(.badURL)))
 
-            await context.waitForUpdate()
+            await context.wait(for: atom, until: \.isFailure)
             XCTAssertEqual(context.watch(atom), .failure(URLError(.badURL)))
         }
 

Tests/AtomsTests/Context/AtomTestContextTests.swift

@@ -50,6 +50,44 @@ final class AtomTestContextTests: XCTestCase {
         XCTAssertFalse(didUpdate1)
     }
 
+    func testWaitFor() async {
+        let atom = TestStateAtom(defaultValue: 0)
+        let context = AtomTestContext()
+
+        context.watch(atom)
+
+        for i in 0..<3 {
+            Task {
+                try? await Task.sleep(seconds: Double(i))
+                context[atom] += 1
+            }
+        }
+
+        let didUpdate0 = await context.wait(for: atom) {
+            $0 == 0
+        }
+
+        XCTAssertFalse(didUpdate0)
+
+        let didUpdate1 = await context.wait(for: atom) {
+            $0 == 3
+        }
+
+        XCTAssertTrue(didUpdate1)
+
+        let didUpdate2 = await context.wait(for: atom, timeout: 1) {
+            $0 == 100
+        }
+
+        XCTAssertFalse(didUpdate2)
+
+        let didUpdate3 = await context.wait(for: atom) {
+            $0 == 3
+        }
+
+        XCTAssertFalse(didUpdate3)
+    }
+
     func testOverride() {
         let atom0 = TestValueAtom(value: 100)
         let atom1 = TestStateAtom(defaultValue: 200)

Tests/AtomsTests/Core/Loader/AtomLoaderContextTests.swift

@@ -72,7 +72,7 @@ final class AtomLoaderContextTests: XCTestCase {
         ) { _, _ in }
 
         await context.transaction { _ in
-            try? await Task.sleep(nanoseconds: 0)
+            try? await Task.sleep(seconds: 0)
         }
 
         XCTAssertTrue(isCommitted)

Tests/AtomsTests/Modifier/TaskPhaseModifierTests.swift

@@ -6,14 +6,14 @@ import XCTest
 @MainActor
 final class TaskPhaseModifierTests: XCTestCase {
     func testPhase() async {
-        let atom = TestTaskAtom(value: 0)
+        let atom = TestTaskAtom(value: 0).phase
         let context = AtomTestContext()
 
-        XCTAssertEqual(context.watch(atom.phase), .suspending)
+        XCTAssertEqual(context.watch(atom), .suspending)
 
-        await context.waitForUpdate(timeout: 1)
+        await context.wait(for: atom, until: \.isSuccess)
 
-        XCTAssertEqual(context.watch(atom.phase), .success(0))
+        XCTAssertEqual(context.watch(atom), .success(0))
     }
 
     func testKey() {