Mocking URLSession to return a Mocked URLSessionDataTask

Question

I want to mock URLSession, and return a mocked URLSessionDataTask.

To Mock URLSession I create a protocol

protocol URLSessionProtocol {
    func dataTask(with url: URL, completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTask
}

that URLSession can then conform to in an extension

extension URLSession: URLSessionProtocol {}

Now I want to do the same for URLSessionDataTask, and so implement a similar protocol and extension for it. I need to do this, since the way I call URLSession requires use of func dataTask(with url: URL, completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void) -> URLSessionDataTask

protocol URLSessionDataTaskProtocol {
    func resume()
}

extension URLSessionDataTask: URLSessionDataTaskProtocol {}

So then my URLSessionDataTask mock is set up as follows:

class URLSessionMock: URLSessionProtocol {
    typealias CompletionHandler = (Data?, URLResponse?, Error?) -> Void
    // data and error can be set to provide data or an error
    var data: Data?
    var error: Error?
    func dataTask(
        with url: URL,
        completionHandler: @escaping CompletionHandler
        ) -> URLSessionDataTask {
        let data = self.data
        let error = self.error
        return URLSessionDataTaskMock {
            completionHandler(data, nil, error)
        }
    }
}

With my URLSessionDataTaskMock presented with:

class URLSessionDataTaskMock: URLSessionDataTaskProtocol {
    private let closure: () -> Void
    init(closure: @escaping () -> Void) {
        self.closure = closure
    }
    // override resume and call the closure

    func resume() {
        closure()
    }
}

Doesn't work since URLSessionDataTaskMock within the URLSessionProtocol isn't the correct return type - I need to return a URLSessionDataTask.

I can't cast my URLSessionDataTaskMock to URLSessionDataTask as the types are not related.

How can I return my URLSessionDataTaskMock from my URLSessionProtocol?

Answer 1

You can probably get away with something like this. The key would be your associated type in your URLSessionProtocol

protocol URLSessionProtocol {
    associatedtype DataTaskType
    func dataTask(with url: URL, completionHandler: @escaping (Data?, URLResponse?, Error?) -> Void) -> DataTaskType
}

extension URLSession: URLSessionProtocol {}

protocol URLSessionDataTaskProtocol {
    func resume()
}

extension URLSessionDataTask: URLSessionDataTaskProtocol {}

class URLSessionDataTaskMock: URLSessionDataTaskProtocol {
    typealias CompletionHandler = URLSessionMock.CompletionHandler

    private let completion: CompletionHandler

    init(completion: @escaping CompletionHandler) {
        self.completion = completion
    }

    func resume() {
        // create some data
        completion(nil, nil, nil)
    }
}

class URLSessionMock: URLSessionProtocol {
    typealias CompletionHandler = (Data?, URLResponse?, Error?) -> Void
    // data and error can be set to provide data or an error
    var data: Data?
    var error: Error?
    func dataTask(
        with url: URL,
        completionHandler: @escaping CompletionHandler
    ) -> URLSessionDataTaskMock {
        return URLSessionDataTaskMock(completion: completionHandler)
    }
}