[v2] [3/X] RequestChain + Response Parsing Rewrite

Tests/ApolloInternalTestHelpers/MockHTTPResponse.swift

@@ -2,23 +2,6 @@ import Apollo
 import ApolloAPI
 import Foundation
 
-extension HTTPResponse {
-  public static func mock(
-    statusCode: Int = 200,
-    headerFields: [String : String] = [:],
-    data: Data = Data()
-  ) -> HTTPResponse {
-    return HTTPResponse(
-      response: .mock(
-        statusCode: statusCode,
-        headerFields: headerFields
-      ),
-      rawData: data,
-      parsedResponse: nil
-    )
-  }
-}
-
 extension HTTPURLResponse {
   public static func mock(
     url: URL = TestURL.mockServer.url,

Tests/ApolloInternalTestHelpers/MockResponseProvider.swift

@@ -0,0 +1,119 @@
+import Foundation
+
+/// A protocol that a type can conform to to provide mock responses to a ``MockURLSession``.
+///
+/// Typically an `XCTestCase` will conform to this protocol directly.
+///
+/// - Important: To ensure test isolation, requests handlers must be cleaned up after each
+/// individual test by calling `Self.cleanUpRequestHandlers()`.
+///
+/// This is an example of how an `XCTestCase` should use this protocol.
+///
+/// ```
+/// class MyTests: XCTestCase, MockResponseProvider {
+///
+///   var session: MockURLSession!
+///
+///   override func setUp() async throws {
+///     try await super.setUp()
+///
+///     session = MockURLSession(requestProvider: Self.self)
+///   }
+///
+///   override func tearDown() async throws {
+///     await Self.cleanUpRequestHandlers()
+///     session = nil
+///     sessionConfiguration = nil
+///
+///     try await super.tearDown()
+///   }
+///
+///   func testThatMocksAResponse() async throws {
+///     let url = URL(string: "www.test.com")!
+///     let responseStrings: [String] = [
+///       ... // An array of strings for each multipart response chunk in your mocked response.
+///     ]
+///     await Self.registerRequestHandler(for: url) { request in
+///       let response = HTTPURLResponse(
+///         url: url,
+///         statusCode: 200,
+///         httpVersion: nil,
+///         headerFields: nil
+///       )
+///
+///       let stream = AsyncThrowingStream { continuation in
+///         for string in responseStrings {
+///           continuation.yield(string.data(using: .utf8)!)
+///         }
+///         continuation.finish()
+///       }
+///
+///       return (response. stream)
+///     }
+///
+///     // Send mockRequest
+///     let (dataStream, response) = try await self.session.bytes(for: request, delegate: nil)
+///
+///     for try await chunk in dataStream.chunks {
+///       let chunkString = String(data: chunk, encoding: .utf8)
+///
+///       ... // test assertions on response
+///     }
+///   }
+/// ```
+public protocol MockResponseProvider {
+  typealias MultiResponseHandler = @Sendable (URLRequest) async throws -> (HTTPURLResponse, AsyncThrowingStream<Data, any Error>?)
+  typealias SingleResponseHandler = @Sendable (URLRequest) async throws -> (HTTPURLResponse, Data?)
+}
+
+extension MockResponseProvider {
+
+  public static func registerRequestHandler(for url: URL, handler: @escaping MultiResponseHandler) async {
+    await requestStorage.registerRequestHandler(for: Self.self, url: url, handler: handler)
+  }
+
+  public static func registerRequestHandler(for url: URL, handler: @escaping SingleResponseHandler) async {
+    await requestStorage.registerRequestHandler(for: Self.self, url: url, handler: { request in
+      let (response, data) = try await handler(request)
+      let stream = AsyncThrowingStream { data }
+
+      return (response, stream)
+    })
+  }
+
+  public static func requestHandler(for url: URL) async -> MultiResponseHandler? {
+    await requestStorage.requestHandler(for: Self.self, url: url)
+  }
+
+  /// Removes all request handlers for the provider type.
+  ///
+  /// - Important: To ensure test isolation, requests handlers must be cleaned up after each
+  /// individual test.
+  public static func cleanUpRequestHandlers() async {
+    await requestStorage.removeRequestHandlers(for: Self.self)
+  }
+}
+
+fileprivate let requestStorage = RequestStorage()
+fileprivate actor RequestStorage {
+  private var providers: [ObjectIdentifier: [URL: MockResponseProvider.MultiResponseHandler]] = [:]
+
+  func registerRequestHandler(
+    for type: Any.Type,
+    url: URL,
+    handler: @escaping MockResponseProvider.MultiResponseHandler
+  ) {
+    let typeId = ObjectIdentifier(type)
+    var handlers = providers[typeId, default: [:]]
+    handlers[url] = handler
+    providers[typeId] = handlers
+  }
+
+  func requestHandler(for type: Any.Type, url: URL) -> MockResponseProvider.MultiResponseHandler? {
+    providers[ObjectIdentifier(type)]?[url]
+  }
+
+  func removeRequestHandlers(for type: Any.Type) {
+    providers[ObjectIdentifier(type)] = nil
+  }
+}

Tests/ApolloInternalTestHelpers/MockURLProtocol.swift

@@ -1,56 +1,53 @@
 import Foundation
 
-public class MockURLProtocol<RequestProvider: MockRequestProvider>: URLProtocol {
-  
+public final class MockURLProtocol<RequestProvider: MockResponseProvider>: URLProtocol {
+
   override class public func canInit(with request: URLRequest) -> Bool {
     return true
   }
 
   override class public func canonicalRequest(for request: URLRequest) -> URLRequest {
     return request
   }
-
+
+  private var asyncTask: Task<Void, any Error>?
+
   override public func startLoading() {
-    guard
-      let url = self.request.url,
-      let handler = RequestProvider.requestHandlers[url]
-    else { return }
-
-    DispatchQueue.main.asyncAfter(deadline: .now() + Double.random(in: 0.0...0.25)) {
-      defer {
-        RequestProvider.requestHandlers.removeValue(forKey: url)
-      }
-
+    self.asyncTask = Task {
+      guard
+        let url = self.request.url,
+        let handler = await RequestProvider.requestHandler(for: url)
+      else { return }
+
       do {
-        let result = try handler(self.request)
-
-        switch result {
-        case let .success((response, data)):
-          self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
-
-          if let data = data {
-            self.client?.urlProtocol(self, didLoad: data)
-          }
-
+        defer {
           self.client?.urlProtocolDidFinishLoading(self)
-        case let .failure(error):
-          self.client?.urlProtocol(self, didFailWithError: error)
         }
-
+
+        let (response, dataStream) = try await handler(self.request)
+
+        try Task.checkCancellation()
+
+        self.client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
+
+        guard let dataStream else {
+          return
+        }
+
+        for try await data in dataStream {
+          try Task.checkCancellation()
+
+          self.client?.urlProtocol(self, didLoad: data)
+        }
+
       } catch {
         self.client?.urlProtocol(self, didFailWithError: error)
       }
     }
   }
 
   override public func stopLoading() {
+    self.asyncTask?.cancel()
   }
 
 }
-
-public protocol MockRequestProvider {
-  typealias MockRequestHandler = ((URLRequest) throws -> Result<(HTTPURLResponse, Data?), any Error>)
-
-  // Dictionary of mock request handlers where the `key` is the URL of the request.
-  static var requestHandlers: [URL: MockRequestHandler] { get set }
-}

Tests/ApolloInternalTestHelpers/MockURLSession.swift

@@ -2,75 +2,98 @@ import Foundation
 import Apollo
 import ApolloAPI
 
-public final class MockURLSessionClient: URLSessionClient {
+public struct MockURLSession: ApolloURLSession {
 
-  @Atomic public var lastRequest: URLRequest?
-  @Atomic public var requestCount = 0
+  public let session: URLSession
 
-  public var jsonData: JSONObject?
-  public var data: Data?
-  var responseData: Data? {
-    if let data = data { return data }
-    if let jsonData = jsonData {
-      return try! JSONSerializationFormat.serialize(value: jsonData)
-    }
-    return nil
+  public init<T: MockResponseProvider>(responseProvider: T.Type) {
+    let configuration = URLSessionConfiguration.ephemeral
+    configuration.protocolClasses = [MockURLProtocol<T>.self]
+    session = URLSession(configuration: configuration)
   }
-  public var response: HTTPURLResponse?
-  public var error: (any Error)?
-
-  private let callbackQueue: DispatchQueue
-
-  public init(callbackQueue: DispatchQueue? = nil, response: HTTPURLResponse? = nil, data: Data? = nil) {
-    self.callbackQueue = callbackQueue ?? .main
-    self.response = response
-    self.data = data
-  }
-
-  public override func sendRequest(_ request: URLRequest,
-                                   taskDescription: String? = nil,
-                                   rawTaskCompletionHandler: URLSessionClient.RawCompletion? = nil,
-                                   completion: @escaping URLSessionClient.Completion) -> URLSessionTask {
-    self.$lastRequest.mutate { $0 = request }
-    self.$requestCount.increment()
-
-    // Capture data, response, and error instead of self to ensure we complete with the current state
-    // even if it is changed before the block runs.
-    callbackQueue.async { [responseData, response, error] in
-      rawTaskCompletionHandler?(responseData, response, error)
-
-      if let error = error {
-        completion(.failure(error))
-      } else {
-        guard let data = responseData else {
-          completion(.failure(URLSessionClientError.dataForRequestNotFound(request: request)))
-          return
-        }
-
-        guard let response = response else {
-          completion(.failure(URLSessionClientError.noHTTPResponse(request: request)))
-          return
-        }
-
-        completion(.success((data, response)))
-      }
-    }
 
-    let mockTaskType: any URLSessionDataTaskMockProtocol.Type = URLSessionDataTaskMock.self
-    let mockTask = mockTaskType.init() as! URLSessionDataTaskMock
-    return mockTask
+  public func bytes(
+    for request: URLRequest,
+    delegate: (any URLSessionTaskDelegate)?
+  ) async throws -> (URLSession.AsyncBytes, URLResponse) {
+    return try await session.bytes(for: request, delegate: delegate)
   }
-}
 
-protocol URLSessionDataTaskMockProtocol {
-  init()
-}
-
-private final class URLSessionDataTaskMock: URLSessionDataTask, URLSessionDataTaskMockProtocol {
-
-  override func resume() {
-    // No-op
+  public func invalidateAndCancel() {	
+    session.invalidateAndCancel()
   }
-
-  override func cancel() {}
 }
+
+#warning("TODO: Delete this")
+//public final class MockURLSessionClient: URLSessionClient {
+//
+//  @Atomic public var lastRequest: URLRequest?
+//  @Atomic public var requestCount = 0
+//
+//  public var jsonData: JSONObject?
+//  public var data: Data?
+//  var responseData: Data? {
+//    if let data = data { return data }
+//    if let jsonData = jsonData {
+//      return try! JSONSerializationFormat.serialize(value: jsonData)
+//    }
+//    return nil
+//  }
+//  public var response: HTTPURLResponse?
+//  public var error: (any Error)?
+//  
+//  private let callbackQueue: DispatchQueue
+//  
+//  public init(callbackQueue: DispatchQueue? = nil, response: HTTPURLResponse? = nil, data: Data? = nil) {
+//    self.callbackQueue = callbackQueue ?? .main
+//    self.response = response
+//    self.data = data
+//  }
+//
+//  public override func sendRequest(_ request: URLRequest,
+//                                   taskDescription: String? = nil,
+//                                   rawTaskCompletionHandler: URLSessionClient.RawCompletion? = nil,
+//                                   completion: @escaping URLSessionClient.Completion) -> URLSessionTask {
+//    self.$lastRequest.mutate { $0 = request }
+//    self.$requestCount.increment()
+//
+//    // Capture data, response, and error instead of self to ensure we complete with the current state
+//    // even if it is changed before the block runs.
+//    callbackQueue.async { [responseData, response, error] in
+//      rawTaskCompletionHandler?(responseData, response, error)
+//      
+//      if let error = error {
+//        completion(.failure(error))
+//      } else {
+//        guard let data = responseData else {
+//          completion(.failure(URLSessionClientError.dataForRequestNotFound(request: request)))
+//          return
+//        }
+//        
+//        guard let response = response else {
+//          completion(.failure(URLSessionClientError.noHTTPResponse(request: request)))
+//          return
+//        }
+//        
+//        completion(.success((data, response)))
+//      }
+//    }
+//
+//    let mockTaskType: any URLSessionDataTaskMockProtocol.Type = URLSessionDataTaskMock.self
+//    let mockTask = mockTaskType.init() as! URLSessionDataTaskMock
+//    return mockTask
+//  }
+//}
+//
+//protocol URLSessionDataTaskMockProtocol {
+//  init()
+//}
+//
+//private final class URLSessionDataTaskMock: URLSessionDataTask, URLSessionDataTaskMockProtocol {
+//
+//  override func resume() {
+//    // No-op
+//  }
+//
+//  override func cancel() {}
+//}