While Apple’s APIs cover most cases, they are not always the ideal solution for specific scenarios. When I was writing tests for WeTransfer’s concurrency code, I frequently encountered flaky tests. At a certain point, I found out about Point-Free’s Swift Concurrency Extras repository and started using it throughout the project as a solution for stable tests.

You won’t need this framework everywhere, but there are cases where you expect a test to succeed while it doesn’t. You’ll likely be able to solve it with the lessons learned today.

A flaky test example

For this lesson, we’re going to write a test to validate the isLoading state in this image fetcher:

import Foundation

#if os(macOS)
import AppKit
typealias PlatformImage = NSImage
#else
import UIKit
typealias PlatformImage = UIImage
#endif

@MainActor
@Observable
final class ImageFetcher {

    var isLoading = false

    // Allow for injecting the download logic explictly for testing purposes.
    let downloadImage: (URL) async throws -> Data

    init(downloadImage: @escaping (URL) async throws -> Data) {
        self.downloadImage = downloadImage
    }

    func fetchImage(imageURL: URL) async throws -> PlatformImage {
        isLoading = true
        defer { isLoading = false }

        let fallbackImage = PlatformImage(named: "fallback_image")!

        let imageData = try await downloadImage(imageURL)

        /// Check for cancellation after the network request.
        /// Potentially, you would perform image operations on `imageData`
        guard !Task.isCancelled else {
            return fallbackImage
        }

        /// ... Perform heavy image operations on `imageData` since the task is not cancelled.

        guard let image = PlatformImage(data: imageData) else {
            /// Converting the data to `UIImage` failed, return our fallback image.
            return fallbackImage
        }

        /// We completed the image download and heavy operations without cancellations, return the image.
        return image
    }
}

You might recognize this code from Module 3 where we worked with it to better understand tasks.

The code is slightly different as it’s optimized for testing with the downloadImage property. It allows you to inject the logic for downloading the image, making it possible to mock the download logic for tests. It also comes with a new isLoading state that is set to true or false depending on whether a download is active.

When you write a test for the isLoading state, it might look as follows:

func testIsLoading() async throws {
    let imageFetcher = ImageFetcher(downloadImage: { url in
        /// Return plain `Data()` for this test as we're just interested in the `isLoading` state.
        return Data()
    })

    /// Call into the `fetchImage` method to ensure the `isLoading` state flips to `true` and `false`.
    let task = Task { _ = try await imageFetcher.fetchImage(imageURL: URL(string: "https://example-image.url")!) }
    XCTAssertEqual(imageFetcher.isLoading, true)

    /// Await the result of the image fetching.
    try await task.value

    /// Validate that `isLoading` restores to `false`.
    XCTAssertEqual(imageFetcher.isLoading, false)
}

This test fails almost always as the Task { ... } returns before we check whether isLoading is set to true. We need a way to pause execution, allowing us to validate the isLoading state in-between. This is where the Swift Concurrency Extras framework comes into place.

Using a Main Serial Executor using Swift Concurrency Extras

The framework offers several solutions, but I’ve been using their main serial executor exclusively. The library includes a static function, withMainSerialExecutor, which attempts to run all tasks spawned in an operation serially and deterministically. This function can be used to make asynchronous tests faster and less flaky.

Do note what they’re saying about this method:

Warning: This API is only intended to be used from tests to make them more reliable. Please do not use it from application code.

We say that it “attempts to run all tasks spawned in an operation serially and deterministically” because under the hood it relies on a global, mutable variable in the Swift runtime to do its job, and there are no scoping guarantees should this mutable variable change during the operation.

To start using the framework, you’ll need to add it as a Swift package using the following URL:

https://github.com/pointfreeco/swift-concurrency-extras.git

You can then import the framework:

import ConcurrencyExtras

And start using the main serial executor inside the test. By running the test using a main serial executor, we can utilize the Task.yield() smartly and read a state in between. The final test looks as follows:

/// By running the test using a main serial executor, we allow ourselves to use `Task.yield()` and read state in between.
func testIsLoadingWithMainSerialExecutor() async throws {
    try await withMainSerialExecutor {
        let imageFetcher = ImageFetcher(downloadImage: { url in

            /// Yield here to allow the test to continue evaluation. This allows us to check the `isLoading` state.
            await Task.yield()

            /// Return plain `Data()` for this test as we're just interested in the `isLoading` state.
            return Data()
        })

        /// Call into the `fetchImage` method to ensure the `isLoading` state flips to `true` and `false`.
        let task = Task { _ = try await imageFetcher.fetchImage(imageURL: URL(string: "https://example-image.url")!) }

        /// Suspends the current test task and allow the above task to start executing up until `downloadImage(...)` is called.
        await Task.yield()

        /// Check the `isLoading` state while we're 'paused' just at `downloadImage`.
        XCTAssertEqual(imageFetcher.isLoading, true)

        /// Await the result of the image fetching.
        try await task.value

        /// Validate that `isLoading` restores to `false`.
        XCTAssertEqual(imageFetcher.isLoading, false)
    }
}

Take your time to read through the test and its inline comments. This explains best what is happening, but I’ll also explain it here.

Since we run the entire test on the main serial executor, we can benefit from using Task.yield() to switch back and forth between the test function and the logic under test.

Better understand the usage of Task.yield()

To fully understand what’s happening, it’s important to take a step back and understand the Task.yield() functionality. The documentation describes it as follows:

If this task is the highest-priority task in the system, the executor immediately resumes execution of the same task. As such, this method isn’t necessarily a way to avoid resource starvation.

In other words, the yield might switch back and forth depending on the executor.

You might try, and find out that your test will succeed most of the time if you use the same yields without a main serial executor:

/// By running the test using a main serial executor, we allow ourselves to use `Task.yield()` and read state in between.
func testIsLoadingWithMainSerialExecutor() async throws {
//        try await withMainSerialExecutor {
        let imageFetcher = ImageFetcher(downloadImage: { url in

            /// Yield here to allow the test to continue evaluation. This allows us to check the `isLoading` state.
            await Task.yield()

            /// Return plain `Data()` for this test as we're just interested in the `isLoading` state.
            return Data()
        })

        /// Call into the `fetchImage` method to ensure the `isLoading` state flips to `true` and `false`.
        let task = Task { _ = try await imageFetcher.fetchImage(imageURL: URL(string: "https://example-image.url")!) }

        /// Suspends the current test task and allow the above task to start executing up until `downloadImage(...)` is called.
        await Task.yield()

        /// Check the `isLoading` state while we're 'paused' just at `downloadImage`.
        XCTAssertEqual(imageFetcher.isLoading, true)

        /// Await the result of the image fetching.
        try await task.value

        /// Validate that `isLoading` restores to `false`.
        XCTAssertEqual(imageFetcher.isLoading, false)
//        }
}

That is correct when running the test in isolation, but if you were to run the test repeatedly, you’ll notice a failure after a while. This is because the executor might or might not be busy, which it probably is when you run a whole suite of tests.

By using the main serialized executor, we know for sure that a yield will be respected and a switch will take place. Since all operations run serially, it will also result in much faster test execution.

Using invokeTest to run all tests on the main serial executor

In case you’re using XCTest, you can decide to run all tests with the main serial executor:

override func invokeTest() {
    withMainSerialExecutor {
        super.invokeTest()
    }
}

Using a main serial executor using Swift Testing

You might start creating a TestTrait for Swift Testing so you can run a test using the main serial executor easily. I thought so too, but wondered why the library didn’t provide a trait out of the box. I found the answer inside this PR, and it turns out that it’s not working nicely together with the parallel execution of Swift Testing. The only way to make this work nicely is by defining your Swift Testing Suite to run serially:

@Suite(.serialized)
@MainActor
final class ImageFetcherSwiftTesting {
    // ...
}

Running tests in parallel

When you run your tests in parallel, you might end up running into issues when using the main serial executor. It’s unfortunate, but it’s caused by the way it works under the hood.

withMainSerialExecutor overrides Swift’s global executor with the main serial executor in an unchecked fashion. When used, all tasks will be enqueued on the main serial executor till its scope ends. If you run tests in parallel, you might affect other tests simultaneously. That’s why it’s essential to use this only when running tests serially.

Summary

In this lesson, we’ve looked into a way to make asynchronous tests more reliable. Flaky tests can occur when validating state when unstructured tasks are being used. By running tests using a main serial executor, we allow ourselves to yield tasks and validate state in between.

Next is the assessment to validate your learning from this module. Good luck!