Finish observation when all group tasks are finished or observed sequence is completed

Given observed sequence is finite,
When sequence completes,
Then observation finishes

Given observed sequence is infinite,
When all tasks complete,
Then observation finishes

Author: erikbasargin

Package.swift

@@ -21,6 +21,8 @@ let package = Package(
         .target(
             name: "AsyncSwiftly",
             dependencies: [
+                "AsyncTrigger",
+                "AsyncMaterializedSequence",
                 .product(name: "AsyncAlgorithms", package: "swift-async-algorithms"),
             ],
         ),

Sources/AsyncSwiftly/AsyncTakeUntilSequence.swift

@@ -0,0 +1,60 @@
+//
+//  AsyncTakeUntilSequence.swift
+//  async-swiftly
+//
+//  Created by Erik Basargin on 05/07/2025.
+//
+
+import AsyncAlgorithms
+import AsyncMaterializedSequence
+
+extension AsyncSequence {
+    
+    func takeUntil<TriggerSequence: AsyncSequence>(
+        _ trigger: TriggerSequence
+    ) -> AsyncTakeUntilSequence<Self, TriggerSequence> where Self: Sendable, Self.Element: Sendable, TriggerSequence: Sendable, TriggerSequence.Element: Sendable {
+        .init(self, trigger)
+    }
+}
+
+struct AsyncTakeUntilSequence<
+    Base1: AsyncSequence,
+    Base2: AsyncSequence
+>: AsyncSequence, Sendable where Base1: Sendable, Base1.Element: Sendable, Base2: Sendable, Base2.Element: Sendable {
+    
+    typealias Base = AsyncCombineLatest2Sequence<
+        AsyncMaterializedSequence<Base1>,
+        AsyncMerge2Sequence<AsyncThrowingMapSequence<AsyncSyncSequence<[Int]>, Bool>, AsyncThrowingMapSequence<Base2, Bool>>
+    >
+    
+    let base: Base
+    
+    init(_ base1: Base1, _ base2: Base2) {
+        let startWith = [1].async.map { _ throws in false }
+        let triggerBase = base2.map { _ throws in true }
+        let trigger = merge(startWith, triggerBase)
+        base = combineLatest(base1.materialize(), trigger)
+    }
+    
+    func makeAsyncIterator() -> Iterator {
+        Iterator(base: base.makeAsyncIterator())
+    }
+    
+    struct Iterator: AsyncIteratorProtocol {
+        
+        var base: Base.AsyncIterator
+        
+        mutating func next() async throws -> Base1.Element? {
+            guard let value = try await base.next() else {
+                return nil
+            }
+            
+            switch value {
+            case (.value(let element), false):
+                return element
+            case (.completed, false), (_, true):
+                return nil
+            }
+        }
+    }
+}

Sources/AsyncSwiftly/TestingTaskGroup.swift

@@ -5,6 +5,8 @@
 //  Created by Erik Basargin on 25/06/2025.
 //
 
+import AsyncTrigger
+import AsyncAlgorithms
 import os
 import Foundation
 
@@ -71,6 +73,7 @@ public struct TestingTaskGroup<ObservationElement: Sendable>: ~Copyable {
     let clock: Clock
     var group: ThrowingDiscardingTaskGroup<any Error>
     let events = EventsStorage<ObservationElement>()
+    let finishObservationTrigger = AsyncTrigger()
 
     public init(group: ThrowingDiscardingTaskGroup<any Error>) {
         self.queue = WorkQueue()
@@ -80,11 +83,11 @@ public struct TestingTaskGroup<ObservationElement: Sendable>: ~Copyable {
 
     public mutating func addObserver<Failure: Error>(
         at rawStep: Int,
-        observer: @escaping @Sendable () -> some AsyncSequence<ObservationElement, Failure>
+        observer: @escaping @Sendable () -> some AsyncSequence<ObservationElement, Failure> & Sendable
     ) {
-        addTask(at: rawStep) { [clock, events] in
+        addTask(at: rawStep) { [clock, events, finishObservationTrigger] in
             do {
-                for try await element in observer() {
+                for try await element in observer().takeUntil(finishObservationTrigger) {
                     let tick = clock.now.when.rawValue
                     events.append(.value(tick, element), for: rawStep)
                 }
@@ -98,8 +101,7 @@ public struct TestingTaskGroup<ObservationElement: Sendable>: ~Copyable {
     }
 
     public mutating func addTask(at rawStep: Int, operation: sending @escaping @isolated(any) () async -> Void) {
-        let duration = Clock.Step.step(rawStep)
-        let instant = Clock.Instant(when: duration)
+        let instant = Clock.Instant(when: .step(rawStep))
         let executor = OperationExecutor(instant: instant, queue: queue)
 
         group.addTask { [queue] in
@@ -113,9 +115,11 @@ public struct TestingTaskGroup<ObservationElement: Sendable>: ~Copyable {
             work()
         }
 
+        finishObservationTrigger.fire()
+        
         if queue.isAnyWorkLeft {
             // At this point if work remains unfinished, we've got some tasks that cannot be resolved in provided range of time.
-            // Wait until all work finishes
+            // Wait until all work finishes or times out.
 
             await queue.waitForAll()
         }
@@ -282,7 +286,9 @@ extension TestingTaskGroup.WorkQueue: AsyncSequence {
             }
 
             repeat {
-                await Task.yield()
+                for _ in 0..<100 {
+                    await Task.yield()
+                }
 
                 switch nextAction() {
                 case let .awaitNextWork(queue):

Tests/AsyncSwiftlyTests/TestingTaskGroupTests.swift

@@ -118,8 +118,8 @@ struct TestingTaskGroupTests {
         }
     }
 
-    @Test("Given system under test produces values, When stream is observed, Then values are received in correct moment")
-    func valuesAreObserved() async throws {
+    @Test("Given observed sequence is finite, When sequence completes, Then observation finishes")
+    func observationFinishes() async throws {
         let (stream, continuation) = AsyncStream.makeStream(of: Int.self)
 
         let result = try await withTestingTaskGroup { group in
@@ -132,6 +132,9 @@ struct TestingTaskGroupTests {
             group.addTask(at: 2) {
                 continuation.yield(2)
             }
+            group.addTask(at: 3) {
+                continuation.yield(3)
+            }
         }
 
         #expect(result[0] == [
@@ -140,4 +143,27 @@ struct TestingTaskGroupTests {
             .finished(2),
         ])
     }
+    
+    @Test("Given observed sequence is infinite, When all tasks complete, Then observation finishes")
+    func infiniteObservationFinishes() async throws {
+        let (stream, continuation) = AsyncStream.makeStream(of: Int.self)
+        
+        let result = try await withTestingTaskGroup { group in
+            group.addObserver(at: 0) {
+                stream
+            }
+            group.addTask(at: 1) {
+                continuation.yield(1)
+            }
+            group.addTask(at: 2) {
+                continuation.yield(2)
+            }
+        }
+        
+        #expect(result[0] == [
+            .value(1, 1),
+            .value(2, 2),
+            .finished(3),
+        ])
+    }
 }