How to Stub URLSession Responses… or Die Trying

Posted 2 days, 13 hours ago.

The internet is filled to the brim with the gravestones of attempts to stub URLSession responses. As something of a strange holdover from the Objective-C days, URLSession is notoriously hard to penetrate with stubbing approaches.

You might think “It’s not a final class, I’ll just subclass and override the methods I care about” but the moment you do, you’re already in a world of pain. Many of the methods you’re probably using these days can’t be overridden directly, because they’re defined in extensions. Even if you do manage it, the sneaky static factory methods in disguise as initialisers will beat you down until you find yourself sobbing in the fetal position.

And then you discover URLProtocol. It promises you the world, allowing you to inject your own logic in response to any request, but when using it for unit testing, it comes at a cost.

What is that cost?

I’m glad you asked.

If you spend any amount of time digging into how to use URLProtocol to stub responses for unit testing, you’re likely to come across something like the following code.

class StubProtocol: URLProtocol {

    static var handler: ((URLRequest) throws -> (Data, URLResponse))?

    override class func canInit(with request: URLRequest) -> Bool {
        true
    }

    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
        request
    }

    override func startLoading() {
        guard let client, let handler = Self.handler else { return }

        do {
            let (data, response) = try handler(request)
            client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
            client.urlProtocol(self, didLoad: data)
            client.urlProtocolDidFinishLoading(self)
        } catch {
            client.urlProtocol(self, didFailWithError: error)
        }
    }

    override func stopLoading() {}

}

This is about as little as you can possibly write to get a URLProtocol that will handle any requests you throw at it, but if you try to build with it, you’ll immediately run into structured concurrency issues with the handler.

Because URLProtocol gets registered by type, and initialised for each incoming request, the handler property has to be static. It also has to be a var because we want to be able to set it from the test itself, otherwise we can’t actually provide the stubs.

So the next step on the journey is likely wrapping the property in some kind of locking mechanism to serialise any reads and writes. One way to do this is to use a little actor container.

class StubProtocol: URLProtocol {

    actor HandlerContainer {

        var handler: ((URLRequest) throws -> (Data, URLResponse))?

    }

    static let container: HandlerContainer

    private var activeTask: Task<Void, Never>?

    override class func canInit(with request: URLRequest) -> Bool {
        true
    }

    override class func canonicalRequest(for request: URLRequest) -> URLRequest {
        request
    }

    override func startLoading() {
        guard let client, let handler = Self.container.handler else { return }

        activeTask = Task {
            do {
                let (data, response) = try await handler(request)
                client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
                client.urlProtocol(self, didLoad: data)
                client.urlProtocolDidFinishLoading(self)
            } catch {
                client.urlProtocol(self, didFailWithError: error)
            }
        }
    }

    override func stopLoading() {
        activeTask?.cancel()
    }

}

A simpler way would be to denote the original static var as being @MainActor isolated, but both result in async dispatches that aren’t necessary, slowing down tests (however minutely) and worse, introducing complexity where it shouldn’t be necessary.

And even then, it still isn’t a perfect solution. Now that you’ve achieved a mechanism that plays nice with structured concurrency on paper, you’re stuck with running tests sequentially. No concurrent tests for you. Not as long as that static annotation has anything to do with it.

If you’re trawling the internet for answers, this is probably the point at which you give up. There’s not a lot of answers out there, in part because URLProtocol is a pretty obscure API.

And yet, there is an answer. How do we make a generic URLProtocol subclass which can reference a function in a way that that is both safe to use with structured concurrency and in a way that ensures it doesn’t leak beyond the scope of the test?

Wait.

Generic?

This is where we get a little weird. If we associated a generic type with the URLProtocol we can use Swift’s ability to define types inside function scopes against it. You still wouldn’t be dealing with instances, so there’s no passing data back and forth between the URLProtocol subclass and your tests (except through the obvious means), but you’d get the benefit of being able to use minimal boilerplate to define something like the following.

@Test func example() async {
    struct Response: Stub {

        static func stub(for request: URLRequest) -> (Data, URLResponse) {
            let data = Data(#"{"foo": "bar"}"#.utf8)
            let response = HTTPURLResponse(url: request.url!, statusCode: 200, httpVersion: nil, headerFields: nil)!
            return (data, response)
        }

    }

    let urlSession = URLSession.stubbed(with: Response.self)
    let (data, response) = try await urlSession.data(for: URL(string: "https://example.com") )

    #expect(data == Data(#"{"foo": "bar"}"#.utf8))
    #expect((response as? HTTPURLSession).statusCode == 200)
}

Look, it’s a really contrived example, but that’s not the point. If we could have an interface like this, where we define this Response type inside of the test method, the scope would be captured by the type, not a global property.

It’s a little more verbose than simply defining the handler, but not by much… and it totally does work.

Here’s the magic code:

public protocol Stub {

    static func stub(for request: URLRequest) throws -> (Data, URLResponse)

}

public extension URLSession {

    private class StubProtocol<S: Stub>: URLProtocol {

        override class func canInit(with request: URLRequest) -> Bool {
            true
        }

        override class func canonicalRequest(for request: URLRequest) -> URLRequest {
            request
        }

        override func startLoading() {
            guard let client else { return }

            do {
                let (data, response) = try S.stub(for: request)
                client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
                client.urlProtocol(self, didLoad: data)
                client.urlProtocolDidFinishLoading(self)
            } catch {
                client.urlProtocol(self, didFailWithError: error)
            }
        }

        override func stopLoading() {}

    }

    static func stubbed<S: Stub>(with stub: S.Type) -> URLSession {
        let config = URLSessionConfiguration.ephemeral
        config.protocolClasses = [StubProtocol<S>.self]
        return URLSession(configuration: config)
    }

}

With this code, you can basically use the previous example to set up a test, and the scope of the function is that of the type with the Stub conformance. You can set it up however you like. It would potentially be possible to set these stubs up in such a way that you have one, global Stub per endpoint response you want to test, then just feed them into your URLProtocol to have it replay them, but you can take it a long way even with just one Stub conformance per test.

If you’d like to play around with this code for yourself, here’s a convenient gist that I’ve prepared with the relevant samples. I’m still experimenting with what this approach can do myself, but I’m also planning to use it in a few different places to cover my usage of URLSession.

Here’s hoping this is the death knell of the actor-locked function approach. I’d love to see what you do with it.