Testing the cancellation branch

I'm working on a websocket client reducer. It has a few long-running tasks: a heartbeat, a listener, etc. When new messages are received, they're passed to a continuation held in state, and a parent reducer subscribes to the corresponding AsyncThrowingStream, also held in state. As long as the parent task is not cancelled, I expect it to run indefinitely, reconnecting on disconnects and errors.

Here's the heartbeat action from the reducer:

case let .beginHeartbeat(socket, topic):
  let heartbeatStream = heartbeatClient(socket, clock)

  return .run(operation: { send in
    for try await _ in heartbeatStream {
      await send(.setConnectionStatus(.connected))
    if Task.isCancelled {
      await send(.onCancel)
    } else {
      await send(.didDisconnectInError(topic: topic))
  }, catch: { error, send in
    await send(.didDisconnectInError(topic: topic))
  }).cancellable(id: Cancellables.heartbeat(id: state.requestId))

When the parent task is cancelled, for example when a SwiftUI .task { ...finish() } goes out of scope, I want to do some cleanup and end the stream. If the heartbeat completes and the parent task hasn't been cancelled, I assume that something has gone wrong, maybe the server has closed its connection, and I'd like to attempt to reconnect.

I've tried to write a test that covers the if Task.isCancelled branch here.

func testHeartbeatHappyPath() async throws {
  let state = NetworkingClientReducer.State()
  let testClock = ImmediateClock()
  let store = TestStore(
    initialState: state,
    reducer: NetworkingClientReducer()
      .dependency(\.websocketHeartbeatClient, { _, _ in [(), (), (), ()].asyncThrowingStream })
      .dependency(\.continuousClock, testClock)

  let unimplementedSocket = UnimplementedSocket()

  let task = await store.send(.beginHeartbeat(socket: unimplementedSocket, topic: "fake-topic"))

  await store.receive(.setConnectionStatus(.connected)) { state in
    state.connectionStatus = .connected

  await store.receive(.setConnectionStatus(.connected))

  await store.receive(.setConnectionStatus(.connected))

  await task.cancel()

  await store.receive(.onCancel)
  await store.skipReceivedActions()
  await store.skipInFlightEffects()

Here I try to create a heartbeat client which will send 4 heartbeats, but I cancel the task after the 3rd one is received. I'd expect to hit the if Task.isCancelled branch, but instead I receive a 4th heartbeat, and then .didDisconnectInError(topic: topic).

Is there a way to do the test I'm trying? Maybe Task.isCancelled ought to be an injected dependency, but that still won't cause the for try await in loop to end.

Maybe this is not a long running effect? If this .asyncThrowingStream helper just publish all values one after another, then it has published all 4 values and finish before you cancel the task.

You hit the nail on the head. I was misunderstanding how receive worked. There might be a few problems with my test. Once .send is run, it'll just keep executing and only later check to see if the store received the correct actions.

I was thinking it would run until it received the action, check the state, and then run until it received the next action I assert against, which obviously doesn't make sense now that I think about it.

Additionally, once the task is cancelled, the store stops processing actions anyway. I tried replacing the heartbeat client with an async stream that would produce heartbeats forever like so

.dependency(\.websocketHeartbeatClient, { _, _ in AsyncThrowingStream { () } })

And then set the task to cancel. It hit the if Task.isCancelled branch this time, but the action wasn't registered at all. Found out this is the same behavior in TestStore as it is in Store, so the cleanup I thought I was running on cancellation, wasn't actually being run at all!

Thanks a bunch for the tip!