How to unit test AsyncStream functionality

Using Swift
, ,
1

I have a simple view model like this:

import Foundation
import SwiftUI

typealias MessagePollResult = Result<[Message], Error>

@MainActor
final class ChatScreenViewModel: ObservableObject {
    @Published var messages: [Message] = []
    
    private let messageProvider: MessageProviderUseCase
    
    private var streamTask: Task<Void, Never>?
    
    init(messageProvider: MessageProviderUseCase) {
        self.messageProvider = messageProvider
    }
    
    func onAppear() {
        streamTask = Task {
            for await messages in await messageProvider.poll(interval: 0.2) {
                if case let .success(messages) = messages {
                    self.messages = messages
                } else {
                    //Show toast
                }
            }
        }
    }
}

protocol MessageProviderUseCase: Sendable {
    func poll(interval: TimeInterval) async -> AsyncStream<MessagePollResult>
}

enum MessageParticipantType: Equatable {
    case sender
    case recipient
}

struct Message: Equatable, Sendable {
    let id: String
    let participant: User
    let participantType: MessageParticipantType
    let content: String
    let timestamp: Date
}

public enum UserAvatarType: Equatable, Sendable {
    case url(URL?)
    case local(Image) // Entity shouldn't know about UI / Image. Will have to refactor later.
}

public struct User: Equatable, Sendable {
    let id: String
    let name: String
    let description: String
    let avatar: UserAvatarType?
}

And now I am trying to make some unit tests for the polling functionality. But they are red, possibly because I might have set the tests wrong, seeing that the real implementation itself is actually working fine. Here's one of the tests:


@MainActor
struct ChatScreenViewModelTests {
    
    @Test func onAppearShouldReturnInitialMessagesAndStartPolling() async throws {
        let mockMessageProvider = MockMessageProvider()
        let sut = createSUT(messageProvider: mockMessageProvider)
        let cancellable: AnyCancellable?
        var messages: [Message] = []
        
        defer {
            cancellable?.cancel()
            #expect(messages[0].count == 0)
            #expect(messages[1].count > 0)
            #expect(mockMessageProvider.pollCallCount == 1)
        }
        
        cancellable = sut.$messages.sink { incoming in
            messages.append(incoming)
        }
        
        sut.onAppear()
        mockMessageProvider.emit(MessagePollResult.success(fakeMessages)) //Emit initial messages

    }

    //- MARK: Helper funcs
    private func createSUT(
        messageProvider: MessageProviderUseCase = MockMessageProvider()) -> ChatScreenViewModel {
        return ChatScreenViewModel(messageProvider: messageProvider)
    }
}

private class MockMessageProvider: MessageProviderUseCase {
    private var continuation: AsyncStream<MessagePollResult>.Continuation?
    private(set) var pollCallCount: Int = 0
    private(set) var pollIntervalInput: TimeInterval?
    
    func poll(interval: TimeInterval) async -> AsyncStream<MessagePollResult> {
        pollCallCount += 1
        pollIntervalInput = interval
        return AsyncStream { continuation in
            self.continuation = continuation
        }
    }
    
    func emit(_ value: MessagePollResult) {
        continuation?.yield(value)
    }

    func finish() {
        continuation?.finish()
    }
}

let fakeSender = User(id: UUID().uuidString,
                      name: "Fulan Doe",
                      description: "Test sender",
                      avatar: .url(nil))

let fakeRecipient = User(id: UUID().uuidString,
                         name: "Fulanah Doe",
                         description: "Test recipient",
                         avatar: .url(nil))

private let fakeMessages = [
    Message(
        id: UUID().uuidString,
        participant: fakeSender,
        participantType: .sender,
        content: "Hi! How are you? I am Fulan. What's your name?",
        timestamp: Date()
    ),
    Message(
        id: UUID().uuidString,
        participant: fakeRecipient,
        participantType: .recipient,
        content: "Hi Fulan! I'm fine thanks. My name is Fulanah. Nice to meet you! What are you doing right now?",
        timestamp: Date()
    )
]

I've added some console prints and it turns out the mockMessageProvider never emits anything. What did I do wrong here? Thanks.

TLDR: The async stream code is working fine but the unit tests are not.

2

For the most part unit tests on AsyncStream are challenging because of the lack of back-pressure. AsyncChannel can get you most of the way there… but the back pressure is applied after the consumer has already started listening.

The other direction I would recommend here is to build your own AsyncSequence implementation that blocks when you send an element and inject this one only for testing. You can find some more ideas on that in here:

3

The standard problem of asynchronous tests is to figure out how to wait for the condition you want to become true.

In this case, you could consider doing something like this in your test:

for await messages in viewModel.$messages.buffer(...).values {
    // assert something about the published messages
}