Reliably testing code that adopts Swift Concurrency?

A good story often has an epilogue.

As things seemed to point to NotificationCentre.Notifications as the culprit in the hard-to-test code here, I wrote a minimal test to demonstrate the bug, intending to submit a Feedback report to the team. Not difficult to test in a nicely encapsulated test:

func test_notificationCentreNotifications_doesNotDisplayRaceCondition() async throws {
    // Setup test expectation
    let expectation = self.expectation(description: "Notification received.")
    let fulfillExpectation = { expectation.fulfill() }
    
    // Initiate the NotificationCentre.Notifications AsyncSequence
    let asyncSequenceTask = Task {
        let screenshotSequence = await NotificationCenter.default.notifications(
            named: UIApplication.userDidTakeScreenshotNotification
        )
        for await _ in screenshotSequence {
            fulfillExpectation()
        }
    }
    
    // Post the notification created when a screenshot occurs.
    _ = await MainActor.run {
        Task.detached(priority:.userInitiated, operation: {
            await NotificationCenter.default.post(name: UIApplication.userDidTakeScreenshotNotification, object: nil)
        })
    }

    // Test expectations
    await self.waitForExpectations(timeout: 1)
    asyncSequenceTask.cancel()
}

Just one "problem":

Executed 100000 tests, with 0 failures (0 unexpected) in 138.589 (203.695) seconds

OK. So, our diagnosis above, pointing the finger at the internal initialization of NotificationCentre.Notifications seemed plausible, but clearly isn't correct.

So we’re back answering the original topic of the thread, a focus on the testability of async-await, with the original question “So how are we supposed to test code like this?”. Effective problem solving requires clear problem definition. And while it might not be initialization of NotificationCentre.Notifications that was the issue, it still seems like an initialization issue was the problem that was making the code untestable.

This is the insight. While the ViewModel was initialized synchronously, it originally only created its AsyncSequence to monitor for screenshots when the async method task() was called. That could be the cause of a race condition. What if we moved the creation of that AsyncSequence to occur as part of the initialization of the ViewModel, and also kick off monitoring during init()?

@MainActor
class ViewModel: ObservableObject {
    @Published var count = 0
    let screenshots = NotificationCenter.default.notifications(named: UIApplication.userDidTakeScreenshotNotification)
    var screenshotMonitor: Task<(), Never>?
    
    init() {
        Task { await monitorForScreenshots() }
    }

    func monitorForScreenshots() async {
        screenshotMonitor = Task {
            for await _ in screenshots {
                self.count += 1
            }
        }
    }
    
    func cancelMonitoring() {
        screenshotMonitor?.cancel()
    }
}

The test for this is then pretty straightfoward, directly examines the count property on the ViewModel, and does not use dependency injection:

func test_selfInitializedViewModel_doesNotDisplayRaceCondition() async throws {
    // Initialize the ViewModel
    let viewModel = await ViewModel()
    
    // Baseline expectations
    await MainActor.run { // MainActor.run required because ViewModel is @MainActor given @Published count
        XCTAssertEqual(viewModel.count, 0, "Expected 0 count, found \(viewModel.count)")
    }
    
    // Post the notification created when a screenshot occurs, then wait till it has been issued.
    _ = await expectation(forNotification: UIApplication.userDidTakeScreenshotNotification, object: nil)
    await NotificationCenter.default.post(name: UIApplication.userDidTakeScreenshotNotification, object: nil)
    await waitForExpectations(timeout: 1)
    
    // Test expectations
    await MainActor.run { // MainActor.run required because ViewModel is @MainActor given @Published count
        XCTAssertEqual(viewModel.count, 1, "Expected 1 count, found \(viewModel.count)")
    }
    
    // Cleanup
    await viewModel.cancelMonitoring()
}

Let’s be really sure. Result of running this test one million times:

Executed 1000000 tests, with 0 failures (0 unexpected) in 2274.146 (3139.190) seconds

What I’ve learned is: to ensure AsyncSequence code is testable, your production code should initialize your AsyncSequence synchronously, as part of its owner’s initiailzation. This also makes the production code slightly more reliable (eliminates the possibility of missing an early screenshot during that async startup), though functionally this probably would be a non-issue.

8 Likes