Testing Old Asynchronous Code with Swift Testing

Fast tests matter. Artificial delays waste time and obscure failures. In a previous article , we explored this with XCTest. Now there’s a new wrinkle: Swift Testing. Swift Testing works really well alongside modern code taking advantage of Swift Concurrency, but if you are migrating a project to Swift Testing and still need to test some older code that uses completion handlers, you will find yourself between a rock and a hard place. Swift Testing has no equivalent of XCTestExpectation, and an improperly used continuation can silently hang your entire CI pipeline.

Here is the implementation code we were working on in the earlier article:

class UserGetter {
    func getUser(_ completionHandler: (User) -> Void) {
        // do some network stuff here...
        completionHandler(user)
    }
}

With XCTest we are able to use XCTestExpectationto wait for the completion handler to be called, but there is no equivalent in Swift Testing. In Swift Testing, we need to use continuations to bridge completion handler based code into async code so it can be tested properly. To test the function above with a continuation, we can do something like this:

struct UserGetterTests {
    @Test func getUser_callsCompletionHandler() async {
        let subject = UserGetter()
    
        let user = await withCheckedContinuation { continuation in
            subject.getUser { user in
                continuation.resume(returning: user)
            }
        }
      
        #expect(user.name == "John Doe")
    }
}

When calling withCheckedContinuation(_:), we pass in a block that receives a CheckedContinuation. Execution suspends in the calling context until the CheckedContinuation.resume(returning:) function is called on the continuation, at which point withCheckedContinuation(_:) returns the User and control flow returns to the test. Inside the block, we call the UserGetter.getUser(_:) function and then resume the continuation when the completion handler is called. This results in a test that immediately continues on to the assertions once the getUser(_:) function has completed.

With continuations, it’s important to handle resuming with care. You must call the resume(returning:) function (or another resume function) once and only once; the caller will be suspended until you do. So, if there’s a chance that your completion handler will not get called in this test, you need to be prepared to deal with that. Since a test is one tool that allows you to find incorrect code, even if you expect that the completion handler will always be called, you should build your test in a way that it can degrade gracefully if the completion handler does not get called, clearly directing you towards the problem in the case of a failure. This is especially important when your tests run under a continuous integration environment, where the test could hang the entire job, and it might be difficult to identify the offending test.

We want to add a timeout to make sure our test does not hang. There are a few ways to accomplish this, but we’ll use a task group. Setting this up is a bit involved. We create two tasks:

• A task that attempts to run the getUser(_:) method
• A task that just waits for 2 seconds and then completes (to drive the timeout)

Then we wait for the first result from the task group. If the test successfully gets the user, the first task will complete first. If getting the user takes longer than 2 seconds, the second task will finish first. We then cancel any remaining tasks and return the user (if we have one) or indicate a timeout (if we didn’t get a user):

let subject = UserGetter()

enum UserResult {
    case user(User)
    case timeout
}

let result = await withTaskGroup(of: UserResult.self) { group in
    // add the first task to try to get the user
    group.addTask {
        let user = await withCheckedContinuation { continuation in
            subject.getUser { user in
                continuation.resume(returning: user)
            }
        }
        return .user(user)
    }
    // add the second task to time out after 2 seconds
    group.addTask {
        try? await Task.sleep(for: .seconds(2))
        return .timeout
    }
    // get the result of whichever task finishes first, then cancel the other task
    let result = await group.next()!
    group.cancelAll()
    return result
}

switch result {
case .user(let user):
    #expect(user.name == "John Doe")
case .timeout:
    Issue.record("Completion handler was not called in time.")
}

Now that we have a timeout, we know that this test will not hang the testing process. If we have to write several tests using this combination (Swift Testing + completion handlers), it probably makes sense to abstract this timeout behavior into a function:

struct TimeoutError: Error {}

func withContinuationWithTimeout<ResultType>(_ closure: @escaping (CheckedContinuation<ResultType, Never>) -> Void) async throws -> ResultType {

    let result = await withTaskGroup(of: Result<ResultType, Error>.self) { group in
        group.addTask {
            let result = await withCheckedContinuation { continuation in
                closure(continuation)
            }
            return .success(result)
        }
        group.addTask {
            try? await Task.sleep(for: .seconds(2))
            return .failure(TimeoutError())
        }
        let result = await group.next()!
        group.cancelAll()
        return result
    }

    switch result {
    case .success(let resultContent):
        return resultContent
    case .failure(let error):
        throw error
    }
}

We can use it like this in a test:

let subject = UserGetter()

let user = try await withContinuationWithTimeout { continuation in
    subject.getUser { user in
        continuation.resume(returning: user)
    }
}

#expect(user.name == "John Doe")

Using continuations with our helper method gets us pretty close to the level of conciseness we had with with expectations in XCTest, and helps us avoid having to remember to deal with the pitfalls around using continuations. Over time, as asynchronous code gets updated to async/await, we can easily update the tests and enjoy the benefits of using Swift Testing, since we’ve already implemented the tests there.

• Al Wold

  Software Engineer

  More from Al