Add tests for "Effect cancellation retains publisher" issue

Tests/ComposableArchitectureTests/EffectCancellationTests.swift

@@ -286,4 +286,194 @@ final class EffectCancellationTests: XCTestCase {
     scheduler.advance(by: 1)
     XCTAssertEqual(expectedOutput, [])
   }
+
+  /// Custom Combine Publisher, used in the following tests
+  class CustomPublisher: Publisher {
+    typealias Output = Void
+    typealias Failure = Never
+
+    func receive<S>(subscriber: S) where S: Subscriber, S.Failure == Failure, S.Input == Output {}
+  }
+
+  /// Check if CustomPublisher is correctly released from the memory after cancelling the subscription
+  func testCustomPublisherRelease() {
+    var publisher: CustomPublisher!
+    weak var weakPublisher: CustomPublisher?
+    var receivedRequests = 0
+    var receivedCancels = 0
+    var cancellables = Set<AnyCancellable>()
+
+    publisher = CustomPublisher()
+    weakPublisher = publisher
+    publisher.handleEvents(
+      receiveCancel: { receivedCancels += 1 },
+      receiveRequest: { _ in receivedRequests += 1 })
+      .sink(receiveValue: { _ in })
+      .store(in: &cancellables)
+    publisher = nil
+
+    XCTAssertNotNil(weakPublisher, "should retain publisher after subscribing")
+    XCTAssertEqual(receivedRequests, 1, "should receive request after subscribing")
+    XCTAssertEqual(receivedCancels, 0, "should not receive cancel after subscribing")
+
+    cancellables.removeAll()
+
+    XCTAssertNil(weakPublisher, "should not retain publisher after cancellation")
+    XCTAssertEqual(receivedRequests, 1, "should not receive request after cancellation")
+    XCTAssertEqual(receivedCancels, 1, "should receive cancel after cancellation")
+  }
+
+  /// Check if CustomPublisher is correctly released from the memory after cancelling an Effect originating from it
+  func testEffectCancellationPublisherRelease() {
+    struct State: Equatable {}
+
+    enum Action: Equatable {
+      case subscribe
+      case cancel
+      case event
+    }
+
+    var store: Store<State, Action>!
+    weak var weakStore: Store<State, Action>?
+    var reducer: Reducer<State, Action, Void>!
+    weak var weakPublisher: CustomPublisher?
+    var createdPublishers = 0
+    var receivedRequests = 0
+    var receivedCancels = 0
+
+    reducer = Reducer { state, action, _ in
+      struct EffectId: Hashable {}
+
+      switch action {
+      case .subscribe:
+        let publisher = CustomPublisher()
+        weakPublisher = publisher
+        createdPublishers += 1
+        return publisher
+          .handleEvents(receiveCancel: {
+            receivedCancels += 1
+          }, receiveRequest: { _ in
+            receivedRequests += 1
+          })
+          .map { Action.event }
+          .eraseToEffect()
+          .cancellable(id: EffectId(), cancelInFlight: true)
+
+      case .cancel:
+        return .cancel(id: EffectId())
+
+      case .event:
+        return .none
+      }
+    }
+
+    store = Store(
+      initialState: State(),
+      reducer: reducer,
+      environment: ()
+    )
+    weakStore = store
+
+    XCTAssertEqual(createdPublishers, 0, "should not create publisher before subscribing")
+    XCTAssertEqual(receivedRequests, 0, "should not receive requests before subscribing")
+    XCTAssertEqual(receivedCancels, 0, "should not receive cancels before subscribing")
+
+    ViewStore(store).send(.subscribe)
+
+    XCTAssertEqual(createdPublishers, 1, "should create publisher when subscribing")
+    XCTAssertNotNil(weakPublisher, "should retain publisher when subscribing")
+    XCTAssertEqual(receivedRequests, 1, "should receive request when subscribing")
+    XCTAssertEqual(receivedCancels, 0, "should not receive cancel when subscribing")
+
+    ViewStore(store).send(.cancel)
+
+    XCTAssertEqual(createdPublishers, 1, "should not create another publisher when canceled")
+    XCTAssertNil(weakPublisher, "should not retain publisher when canceled")
+    XCTAssertEqual(receivedRequests, 1, "should not receive requests when canceled")
+    XCTAssertEqual(receivedCancels, 1, "should receive cancel when canceled")
+
+    store = nil
+
+    XCTAssertNil(weakStore, "should not retain store when it's no longer referenced")
+    XCTAssertNil(weakPublisher, "should not retain publisher when store is no longer referenced")
+    XCTAssertEqual(receivedRequests, 1, "should not receive requests when store is no longer referenced")
+    XCTAssertEqual(receivedCancels, 1, "should not receive cancels when store is no longer referenced")
+  }
+
+  /// Check if CustomPublisher is correctly released from the memory after cancelling an Effect originating from it inside a combined Reducer
+  func testCombinedReducerEffectCancellationPublisherRelease() {
+    struct State: Equatable {}
+
+    enum Action: Equatable {
+      case subscribe
+      case cancel
+      case event
+    }
+
+    var store: Store<State, Action>!
+    weak var weakStore: Store<State, Action>?
+    var reducer: Reducer<State, Action, Void>!
+    weak var weakPublisher: CustomPublisher?
+    var createdPublishers = 0
+    var receivedRequests = 0
+    var receivedCancels = 0
+
+    reducer = Reducer { state, action, _ in
+      struct EffectId: Hashable {}
+
+      switch action {
+      case .subscribe:
+        let publisher = CustomPublisher()
+        weakPublisher = publisher
+        createdPublishers += 1
+        return publisher
+          .handleEvents(receiveCancel: {
+            receivedCancels += 1
+          }, receiveRequest: { _ in
+            receivedRequests += 1
+          })
+          .map { Action.event }
+          .eraseToEffect()
+          .cancellable(id: EffectId(), cancelInFlight: true)
+
+      case .cancel:
+        return .cancel(id: EffectId())
+
+      case .event:
+        return .none
+      }
+    }
+
+    store = Store(
+      initialState: State(),
+      reducer: .combine(reducer),
+      environment: ()
+    )
+    weakStore = store
+
+    XCTAssertEqual(createdPublishers, 0, "should not create publisher before subscribing")
+    XCTAssertEqual(receivedRequests, 0, "should not receive requests before subscribing")
+    XCTAssertEqual(receivedCancels, 0, "should not receive cancels before subscribing")
+
+    ViewStore(store).send(.subscribe)
+
+    XCTAssertEqual(createdPublishers, 1, "should create publisher when subscribing")
+    XCTAssertNotNil(weakPublisher, "should retain publisher when subscribing")
+    XCTAssertEqual(receivedRequests, 1, "should receive request when subscribing")
+    XCTAssertEqual(receivedCancels, 0, "should not receive cancel when subscribing")
+
+    ViewStore(store).send(.cancel)
+
+    XCTAssertEqual(createdPublishers, 1, "should not create another publisher when canceled")
+    XCTAssertNil(weakPublisher, "should not retain publisher when canceled")
+    XCTAssertEqual(receivedRequests, 1, "should not receive requests when canceled")
+    XCTAssertEqual(receivedCancels, 1, "should receive cancel when canceled")
+
+    store = nil
+
+    XCTAssertNil(weakStore, "should not retain store when it's no longer referenced")
+    XCTAssertNil(weakPublisher, "should not retain publisher when store is no longer referenced")
+    XCTAssertEqual(receivedRequests, 1, "should not receive requests when store is no longer referenced")
+    XCTAssertEqual(receivedCancels, 1, "should not receive cancels when store is no longer referenced")
+  }
 }