wordpress-mobile · crazytonyli · Jan 11, 2024 · Dec 20, 2023 · Dec 20, 2023 · Dec 20, 2023
diff --git a/WordPressKit/HTTPClient.swift b/WordPressKit/HTTPClient.swift
@@ -9,37 +9,92 @@ struct HTTPAPIResponse<Body> {
 
 extension URLSession {
 
+    /// Send a HTTP request and return its response as a `WordPressAPIResult` instance.
+    ///
+    /// ## Progress Tracking and Cancellation
+    ///
+    /// You can track the HTTP request's overall progress by passing a `Progress` instance to the `fulfillingProgress`
+    /// parameter, which must satisify following requirements:
+    /// - `totalUnitCount` must not be zero.
+    /// - `completedUnitCount` must be zero.
+    /// - It's used exclusivity for tracking the HTTP request overal progress: No children in its progress tree.
-    /// - It's used exclusivity for tracking the HTTP request overal progress: No children in its progress tree.
+    /// - It's used exclusively for tracking the HTTP request overall progress: No children in its progress tree.
-    /// - It's used exclusivity for tracking the HTTP request overal progress: No children in its progress tree.
+    /// - It's used exclusively for tracking the HTTP request overall progress: No children in its progress tree.
+    /// - `cancellationHandler` must be nil. You can call `fulfillingProgress.cancel()` to cancel the ongoing HTTP request.
+    ///
+    ///  Upon completion, the HTTP request's progress fulfills the `fulfillingProgress`.
+    ///
+    /// - Parameters:
+    ///   - builder: A `HTTPRequestBuilder` instance that represents an HTTP request to be sent.
+    ///   - acceptableStatusCodes: HTTP status code ranges that are considered a successful response. Responses with
+    ///         a status code outside of these ranges are returned as a `WordPressAPIResult.unacceptableStatusCode` instance.
+    ///   - parentProgress: A `Progress` instance that will be used as the parent progress of the HTTP request's overall
+    ///         progress. See the function documentation regarding requirements on this argument.
+    ///   - errorType: The concret endpoint error type.
     func perform<E: LocalizedError>(
         request builder: HTTPRequestBuilder,
         acceptableStatusCodes: [ClosedRange<Int>] = [200...299],
+        fulfillingProgress parentProgress: Progress? = nil,
         errorType: E.Type = E.self
     ) async -> WordPressAPIResult<HTTPAPIResponse<Data>, E> {
+        if let parentProgress {
+            assert(parentProgress.completedUnitCount == 0 && parentProgress.totalUnitCount > 0, "Invalid parent progress")
+            assert(parentProgress.cancellationHandler == nil, "The progress instance's cancellationHandler property must be nil")
+        }
+
         guard let request = try? builder.build() else {
             return .failure(.requestEncodingFailure)
         }
 
-        let result: (Data, URLResponse)
-        do {
-            result = try await data(for: request)
-        } catch {
-            if let urlError = error as? URLError {
-                return .failure(.connection(urlError))
-            } else {
-                return .failure(.unknown(underlyingError: error))
+        return await withCheckedContinuation { continuation in
+            let task = dataTask(with: request) { data, response, error in
+                let result: WordPressAPIResult<HTTPAPIResponse<Data>, E> = Self.parseResponse(
+                    data: data,
+                    response: response,
+                    error: error,
+                    acceptableStatusCodes: acceptableStatusCodes
+                )
+
+                continuation.resume(returning: result)
             }
-        }
+            task.resume()
 
-        let (body, response) = result
+            if let parentProgress, parentProgress.totalUnitCount > parentProgress.completedUnitCount {
+                let pending = parentProgress.totalUnitCount - parentProgress.completedUnitCount
+                parentProgress.addChild(task.progress, withPendingUnitCount: pending)
 
-        guard let response = response as? HTTPURLResponse else {
-            return .failure(.unparsableResponse(response: nil, body: body))
+                parentProgress.cancellationHandler = { [weak task] in
+                    task?.cancel()
+                }
+            }
         }
+    }
 
-        guard acceptableStatusCodes.contains(where: { $0 ~= response.statusCode }) else {
-            return .failure(.unacceptableStatusCode(response: response, body: body))
+    private static func parseResponse<E: LocalizedError>(
+        data: Data?,
+        response: URLResponse?,
+        error: Error?,
+        acceptableStatusCodes: [ClosedRange<Int>]
+    ) -> WordPressAPIResult<HTTPAPIResponse<Data>, E> {
+        let result: WordPressAPIResult<HTTPAPIResponse<Data>, E>
+
+        if let error {
+            if let urlError = error as? URLError {
+                result = .failure(.connection(urlError))
+            } else {
+                result = .failure(.unknown(underlyingError: error))
+            }
+        } else {
+            if let httpResponse = response as? HTTPURLResponse {
+                if acceptableStatusCodes.contains(where: { $0 ~= httpResponse.statusCode }) {
+                    result = .success(HTTPAPIResponse(response: httpResponse, body: data ?? Data()))
+                } else {
+                    result = .failure(.unacceptableStatusCode(response: httpResponse, body: data ?? Data()))
+                }
+            } else {
+                result = .failure(.unparsableResponse(response: nil, body: data))
+            }
         }
 
-        return .success(.init(response: response, body: body))
+        return result
     }
 
 }

diff --git a/WordPressKitTests/Utilities/URLSessionHelperTests.swift b/WordPressKitTests/Utilities/URLSessionHelperTests.swift
@@ -106,6 +106,43 @@ class URLSessionHelperTests: XCTestCase {
 
         try XCTAssertEqual(result.get().title, "Hello Post")
     }
+
+    func testProgressTracking() async throws {
+        stub(condition: isPath("/hello")) { _ in
+            HTTPStubsResponse(data: "success".data(using: .utf8)!, statusCode: 200, headers: nil)
+        }
+
+        let progress = Progress.discreteProgress(totalUnitCount: 20)
+        XCTAssertEqual(progress.completedUnitCount, 0)
+        XCTAssertEqual(progress.fractionCompleted, 0)
+
+        let _ = await URLSession.shared.perform(request: .init(url: URL(string: "https://wordpress.org/hello")!), fulfillingProgress: progress, errorType: TestError.self)
+        XCTAssertEqual(progress.completedUnitCount, 20)
+        XCTAssertEqual(progress.fractionCompleted, 1)
+    }
+
+    func testCancellation() async throws {
+        // Give a slow HTTP request that takes 0.5 second to complete
+        stub(condition: isPath("/hello")) { _ in
+            let response = HTTPStubsResponse(data: "success".data(using: .utf8)!, statusCode: 200, headers: nil)
+            response.responseTime = 0.5
+            return response
+        }
+
+        // and cancelling it (in 0.1 second) before it completes
+        let progress = Progress.discreteProgress(totalUnitCount: 20)
+        DispatchQueue.main.asyncAfter(deadline: .now() + .milliseconds(100)) {
+            progress.cancel()
+        }
+
+        // The result should be an cancellation result
+        let result = await URLSession.shared.perform(request: .init(url: URL(string: "https://wordpress.org/hello")!), fulfillingProgress: progress, errorType: TestError.self)
+        if case let .failure(.connection(urlError)) = result, urlError.code == .cancelled {
+            // Do nothing
+        } else {
+            XCTFail("Unexpected result: \(result)")
+        }
+    }
 }
 
 private enum TestError: LocalizedError, Equatable {