Fix: Fixed Subscription test to not discard ConcurrentEventStream immediately

Author: d-exclaimation

Sources/GraphQL/Subscription/EventStream.swift

@@ -31,7 +31,7 @@ public class ConcurrentEventStream<Element>: EventStream<Element> {
 extension AsyncThrowingStream {
     func mapStream<To>(_ closure: @escaping (Element) throws -> To) -> AsyncThrowingStream<To, Error> {
         return AsyncThrowingStream<To, Error> { continuation in
-            Task {
+            let task = Task {
                 do {
                     for try await event in self {
                         let newEvent = try closure(event)
@@ -42,12 +42,16 @@ extension AsyncThrowingStream {
                     continuation.finish(throwing: error)
                 }
             }
+
+            continuation.onTermination = { @Sendable reason in 
+                task.cancel()
+            }
         }
     }
 
     func filterStream(_ isIncluded: @escaping (Element) throws -> Bool) -> AsyncThrowingStream<Element, Error> {
         return AsyncThrowingStream<Element, Error> { continuation in
-            Task {
+            let task = Task {
                 do {
                     for try await event in self {
                         if try isIncluded(event) {
@@ -59,6 +63,10 @@ extension AsyncThrowingStream {
                     continuation.finish(throwing: error)
                 }
             }
+
+            continuation.onTermination = { @Sendable _ in 
+                task.cancel()
+            }
         }
     }
 }

Tests/GraphQLTests/SubscriptionTests/SubscriptionTests.swift

@@ -605,12 +605,15 @@ class SubscriptionTests : XCTestCase {
 
         var results = [GraphQLResult]()
         var expectation = XCTestExpectation()
-        _ = stream.map { event in
+        
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        let keepForNow = stream.map { event in
             event.map { result in
                 results.append(result)
                 expectation.fulfill()
             }
         }
+
         var expected = [GraphQLResult]()
 
         db.trigger(email: Email(
@@ -675,6 +678,9 @@ class SubscriptionTests : XCTestCase {
         )
         wait(for: [expectation], timeout: timeoutDuration)
         XCTAssertEqual(results, expected)
+
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        _ = keepForNow
     }
 
     /// 'should not trigger when subscription is already done'
@@ -701,7 +707,8 @@ class SubscriptionTests : XCTestCase {
 
         var results = [GraphQLResult]()
         var expectation = XCTestExpectation()
-        _ = stream.map { event in
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        let keepForNow = stream.map { event in
             event.map { result in
                 results.append(result)
                 expectation.fulfill()
@@ -747,6 +754,9 @@ class SubscriptionTests : XCTestCase {
         // Ensure that the current result was the one before the db was stopped
         wait(for: [expectation], timeout: timeoutDuration)
         XCTAssertEqual(results, expected)
+
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        _ = keepForNow
     }
 
     /// 'should not trigger when subscription is thrown'
@@ -861,7 +871,8 @@ class SubscriptionTests : XCTestCase {
 
         var results = [GraphQLResult]()
         var expectation = XCTestExpectation()
-        _ = stream.map { event in
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        let keepForNow = stream.map { event in
             event.map { result in
                 results.append(result)
                 expectation.fulfill()
@@ -925,6 +936,9 @@ class SubscriptionTests : XCTestCase {
         )
         wait(for: [expectation], timeout: timeoutDuration)
         XCTAssertEqual(results, expected)
+
+        // So that the Task won't immediately be cancelled since the ConcurrentEventStream is discarded
+        _ = keepForNow
     }
 
     /// 'should pass through error thrown in source event stream'