Unit testing async await

[ahmedk92]

XCTest provides a very convenient way that makes testing methods marked as async a breeze. Just marking the test method as async does the job:

func test_myAsyncMethod() async throws {
  XCTAssertTrue(try await myAsyncMethod())
}

However, sometimes the API of the code we want test is not marked as async although it’s actually async, i.e. it starts a Task under the hood. This makes it hard to test, as we would then resort to expectations which are a bit kludgy. In this article, I’ll explore a way to maintain the access to XCTest’s async.

The idea is inspired by the clever trick of injecting DispatchQueues to be able to wait for them during tests, discussed in this article by John Sundell.

Consider the following code:

import Foundation
import Combine

@MainActor
class ViewModel {
    @Published var email: String = ""
    @Published private(set) var isEmailAvailable: Bool = false

    private var cancellables: Set<AnyCancellable> = []

    typealias CheckEmailAvailability = (String) async -> Bool
    private let checkEmailAvailability: CheckEmailAvailability

    init(checkEmailAvailability: @escaping CheckEmailAvailability) {
        self.checkEmailAvailability = checkEmailAvailability
        observeEmail()
    }

    private func observeEmail() {
        $email.sink { [weak self] email in
            guard let self = self else { return }

            Task {
                self.isEmailAvailable = await self.checkEmailAvailability(email)
            }
        }.store(in: &cancellables)
    }
}

So, what this view model does is:

1. Observe value changes for the email property
2. Delegates the availability checking logic to the async checkEmailAvailability closure
3. Updates the isEmailAvailable property with the result

As mentioned before, we cannot make use of XCTest’s async convenience because we have no control over that implicit Task. If you think about it, what we need from the Task object is just an async context. If we can inject something that provides us with such context, maybe we can have a chance to play well with XCTest’s async. Let’s try this step by step.

First, we can safely assume that any closure we will be passing to Task’s initializer has the following signature:

() async throws -> T

Therefore, we can define a protocol that accepts such closure, and run it:

typealias AsyncClosure<T> = () async throws -> T

protocol AsyncRunner {
    func runAsync<T>(closure: @escaping AsyncClosure<T>)
}

Now, for non-testing purposes, we are interested in a default implementation that runs the given async closure inside a Task. So, we can have this default implementation:

class DefaultAsyncRunner: AsyncRunner {
    func runAsync<T>(closure: @escaping AsyncClosure<T>) {
        Task {
            try await closure()
        }
    }
}

Now, let’s refactor the view model above to inject an instance of this protocol:

import Foundation
import Combine

@MainActor
class ViewModel {
    @Published var email: String = ""
    @Published private(set) var isEmailAvailable: Bool = false

    private var cancellables: Set<AnyCancellable> = []

    typealias CheckEmailAvailability = (String) async -> Bool
    private let checkEmailAvailability: CheckEmailAvailability
    private let asyncRunner: AsyncRunner

    init(
        asyncRunner: AsyncRunner = DefaultAsyncRunner(),
        checkEmailAvailability: @escaping CheckEmailAvailability
    ) {
        self.asyncRunner = asyncRunner
        self.checkEmailAvailability = checkEmailAvailability
        observeEmail()
    }

    private func observeEmail() {
        $email.sink { [weak self] email in
            guard let self = self else { return }

            self.asyncRunner.runAsync {
                self.isEmailAvailable = await self.checkEmailAvailability(email)
            }
        }.store(in: &cancellables)
    }
}

Good so far. Now, let’s see how can we test this view model. You can imagine how unit testing the view model will look like, so I’ll focus first on how the mock implementation for the AsyncRunner we will use in the test will look like.

private class MockAsyncRunner: AsyncRunner {
    private var asyncClosures: [AsyncClosure<Any?>] = []

    func runAsync<T>(closure: @escaping AsyncClosure<T>) {
        asyncClosures.append(closure)
    }

    func awaitAll() async throws {
        while !asyncClosures.isEmpty {
            let closure = asyncClosures.removeFirst()
            _ = try await closure()
        }
    }
}

The idea here is, instead of immediately executing the async closure passed to runAsync like in the default implementation above, we save it to an array. Then before we do our test assertions, we execute all the saved closures in-order. And since we now have access to each async closure, we can explicitly await them at once, maintaining the convenience of asynchronously executing in an async XCTest method. Full unit test code:

import XCTest
@testable import TestingAsyncAwaitExample

final class TestingAsyncAwaitExampleTests: XCTestCase {
    private var viewModel: ViewModel!
    private var mockAsyncRunner: MockAsyncRunner!
    private var isEmailAvailable: [String : Bool]! = [:]

    @MainActor override func setUpWithError() throws {
        mockAsyncRunner = .init()
        viewModel = .init(
            asyncRunner: mockAsyncRunner,
            checkEmailAvailability: { [unowned self] email in
                self.isEmailAvailable[email] ?? false
            }
        )
    }

    override func tearDownWithError() throws {
        viewModel = nil
        mockAsyncRunner = nil
        isEmailAvailable = nil
    }

    @MainActor
    func test_notAvailable() async throws {
        // Given
        let givenEmail = "ahmed@example.com"
        isEmailAvailable[givenEmail] = false

        // When
        viewModel.email = givenEmail

        try await mockAsyncRunner.awaitAll()

        // Then
        XCTAssertFalse(viewModel.isEmailAvailable)
    }

    @MainActor
    func test_available() async throws {
        // Given
        let givenEmail = "ahmed@example.com"
        isEmailAvailable[givenEmail] = true

        // When
        viewModel.email = givenEmail

        try await mockAsyncRunner.awaitAll()

        // Then
        XCTAssertTrue(viewModel.isEmailAvailable)
    }

}

private class MockAsyncRunner: AsyncRunner {
    private var asyncClosures: [AsyncClosure<Any?>] = []

    func runAsync<T>(closure: @escaping AsyncClosure<T>) {
        asyncClosures.append(closure)
    }

    func awaitAll() async throws {
        while !asyncClosures.isEmpty {
            let closure = asyncClosures.removeFirst()
            _ = try await closure()
        }
    }
}

Conclusion

Although this solution looks appealing, I’d first consider converting my APIs to be marked async. This is vital for structured concurrency, and maintaining a healthy Task hierarchy that makes automatic cancellation possible. Anyway, in situations when this is not feasible or can be inconvenient (e.g. the view model above could expose an async setter method for the email instead of a mutable published property, which breaks the FRP style), I think this solution can come handy.