Issue 18 - Adding Upload Completion Handler for Handling Custom/Error…

Sources/YNetwork/NetworkManager/NetworkManager+URLSessionTaskDelegate.swift

@@ -10,7 +10,7 @@ import Foundation
 
 // MARK: - URLSessionTaskDelegate
 
-extension NetworkManager: URLSessionTaskDelegate {
+extension NetworkManager: URLSessionTaskDelegate, URLSessionDataDelegate {
     /// :nodoc:
     public func urlSession(
         _ session: URLSession,
@@ -40,4 +40,15 @@ extension NetworkManager: URLSessionTaskDelegate {
             fileUpload.urlSession(session, task: task, didCompleteWithError: error)
         }
     }
+
+    /// :nodoc:
+    public func urlSession(
+        _ session: URLSession,
+        dataTask: URLSessionDataTask,
+        didReceive data: Data
+    ) {
+        if dataTask is URLSessionUploadTask {
+            fileUpload.urlSession(session, dataTask: dataTask, didReceive: data)
+        }
+    }
 }

Sources/YNetwork/NetworkManager/NetworkManager.swift

@@ -161,10 +161,12 @@ open class NetworkManager: NSObject {
     /// - Parameters:
     ///   - request: the upload network request to submit
     ///   - progress: progress handler (will be called back on main thread)
+    ///   - completionHandler: file upload handler (will be called back on URLSession background thread).
     /// - Returns: a cancelable upload task if one was able to be created, otherwise nil if no task was issued
     @discardableResult open func submitBackgroundUpload(
         _ request: NetworkRequest,
-        progress: ProgressHandler? = nil
+        progress: ProgressHandler? = nil,
+        completionHandler: FileUploadHandler? = nil
     ) -> Cancelable? {
         guard let urlRequest = try? buildUrlRequest(request: request) else { return nil }
 
@@ -178,7 +180,7 @@ open class NetworkManager: NSObject {
 
         // creating the upload task copies the file
         let task = try? configuration?.networkEngine.submitBackgroundUpload(urlRequest, fileUrl: localURL)
-        fileUpload.register(task, progress: progress)
+        fileUpload.registerUpload(task, progress: progress, completion: completionHandler)
         return task
     }
 

Sources/YNetwork/NetworkManager/Progress/FileProgress.swift

@@ -16,6 +16,13 @@ public typealias Percentage = Double
 /// Guaranteed to be called back on the main thread
 public typealias ProgressHandler = (Percentage) -> Void
 
+/// Asynchronous completion handler that is called when a response is received for an upload
+public typealias FileUploadHandler = (Data) -> Void
+
+/// Asynchronous cancellation handler that is called when an upload request is cancelled
+/// This can be used with the cancellationHandler attribute of the Progress object associated with the upload task
+public typealias CancellationHandler = () -> Void
+
 /// Asynchronous completion handler that reports the status of request.
 ///
 /// Guaranteed to be called back on a background thread (because the system erases the temporary file)
@@ -29,6 +36,7 @@ public typealias FileDownloadHandler = (Result<URL, Error>) -> Void
 internal class FileProgress: NSObject {
     /// Stores the progress handlers to be called, keyed by unique task identifier
     private var progressHandlersByTaskID: [Int: ProgressHandler] = [:]
+    private var uploadHandlerByTaskID: [Int: FileUploadHandler] = [:]
     private var downloadHandlersByTaskID: [Int: FileDownloadHandler] = [:]
 
     /// Updates the progress handler for the specified task with the percentage value
@@ -41,6 +49,17 @@ internal class FileProgress: NSObject {
             progressHandler(percent)
         }
     }
+
+    /// Invokes the completion handler for the specified task with the response data
+    /// - Parameters:
+    ///   - data: the response data that can be decoded for custom responses such as error messages
+    ///   - taskIdentifier: unique task identifier
+    func receive(data: Data, forKey taskIdentifier: Int) {
+        guard let completionHandler = uploadHandlerByTaskID[taskIdentifier] else { return }
+        DispatchQueue.main.async {
+            completionHandler(data)
+        }
+    }
 
     /// Updates the request status for the specified task with the file URL
     /// - Parameters:
@@ -53,16 +72,32 @@ internal class FileProgress: NSObject {
         completionhandler(result)
     }
 
-    /// Registers a data task for file progress.
+    /// Registers a data task for file progress of either an upload or download.
     /// - Parameters:
     ///   - cancelable: optional cancelable task
     ///   - progress: optional progress handler
-    func register(_ cancelable: Cancelable?, progress: ProgressHandler?) {
+    func registerProgress(
+        _ cancelable: Cancelable?,
+        progress: ProgressHandler?
+    ) {
         guard let task = cancelable as? URLSessionTask,
               let progress = progress else { return }
         progressHandlersByTaskID[task.taskIdentifier] = progress
     }
 
+    /// Registers the data task with a completion handler to be called when the response to the upload is received.
+    /// - Parameters:
+    ///   - cancelable: optional cancelable task
+    ///   - completion: optional completion handler
+    func registerCompletion(
+        _ cancelable: Cancelable?,
+        completion: FileUploadHandler?
+    ) {
+        guard let task = cancelable as? URLSessionTask,
+              let completion = completion else { return }
+        uploadHandlerByTaskID[task.taskIdentifier] = completion
+    }
+
     /// Registers a data task for file progress.
     /// - Parameters:
     ///   - cancelable: optional cancelable task
@@ -73,17 +108,37 @@ internal class FileProgress: NSObject {
         progress: ProgressHandler?,
         handler: @escaping FileDownloadHandler
     ) {
-        register(cancelable, progress: progress)
+        registerProgress(cancelable, progress: progress)
         guard let task = cancelable as? URLSessionTask else { return }
         downloadHandlersByTaskID[task.taskIdentifier] = handler
     }
+
+    /// Registers a data task for file upload progress and completion.
+    /// - Parameters:
+    ///   - cancelable: optional cancelable task
+    ///   - progress: optional progress handler
+    ///   - completion: optional completion handler
+    func registerUpload(
+        _ cancelable: Cancelable?,
+        progress: ProgressHandler?,
+        completion: FileUploadHandler?
+    ) {
+        registerProgress(cancelable, progress: progress)
+        registerCompletion(cancelable, completion: completion)
+    }
 
     /// Unregisters a data task for file progress
     /// - Parameter taskIdentifier: unique task identifier
     func unregister(forKey taskIdentifier: Int) {
         progressHandlersByTaskID.removeValue(forKey: taskIdentifier)
         downloadHandlersByTaskID.removeValue(forKey: taskIdentifier)
     }
+
+    /// Unregisters a completion handler, should be called once the final response is received
+    /// - Parameter taskIdentifier: unique task identifier
+    func unregisterUploadCompletion(forKey taskIdentifier: Int) {
+        uploadHandlerByTaskID.removeValue(forKey: taskIdentifier)
+    }
 
     func checkResponseForError(task: URLSessionTask) -> Error? {
         guard let httpResponse = task.response as? HTTPURLResponse else {

Sources/YNetwork/NetworkManager/Progress/FileUploadProgress.swift

@@ -15,7 +15,7 @@ import Foundation
 /// to optionally track progress for large file upload tasks.
 internal class FileUploadProgress: FileProgress { }
 
-extension FileUploadProgress: URLSessionTaskDelegate {
+extension FileUploadProgress: URLSessionTaskDelegate, URLSessionDataDelegate {
     func urlSession(
         _ session: URLSession,
         task: URLSessionTask,
@@ -31,4 +31,14 @@ extension FileUploadProgress: URLSessionTaskDelegate {
         // clean up the task now that we're finished with it
         unregister(forKey: task.taskIdentifier)
     }
+
+    public func urlSession(
+        _ session: URLSession,
+        dataTask: URLSessionDataTask,
+        didReceive data: Data
+    ) {
+        // clean up the task now that the final response for the upload was received
+        receive(data: data, forKey: dataTask.taskIdentifier)
+        unregisterUploadCompletion(forKey: dataTask.taskIdentifier)
+    }
 }

Tests/YNetworkTests/NetworkManager/NetworkManagerUploadTests.swift

@@ -74,7 +74,7 @@
 
         XCTAssertNil(sut.receivedError)
 
-        let task = try XCTUnwrap(sut.submitBackgroundUpload(request) { _ in } as? URLSessionTask)
+        let task = try XCTUnwrap(sut.submitBackgroundUpload(request, completionHandler:  { _ in }) as? URLSessionTask)
         task.cancel() // this will make it fail
         task.resume() // resume it
 
@@ -90,7 +90,7 @@
         XCTAssertNil(sut.receivedError)
 
         URLProtocolStub.appendStub(.failure(NetworkError.invalidResponse), type: .upload)
-        let task = sut.submitBackgroundUpload(request) { _ in }
+        let task = sut.submitBackgroundUpload(request, completionHandler:  { _ in })
 
         XCTAssertNotNil(task)
 
@@ -154,7 +154,7 @@
         let sut = NetworkManager()
 
         // Given we submit a request without first configuring the network manager
-        let task = sut.submitBackgroundUpload(request) { _ in }
+        let task = sut.submitBackgroundUpload(request, completionHandler:  { _ in })
 
         // We don't expect a task to be returned
         XCTAssertNil(task)
@@ -209,19 +209,17 @@
             self.fulfill()
         }
     }
+
+    override func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+        self.receivedData = data
+    }
 
     func fulfill() {
         expectation?.fulfill()
         expectation = nil
     }
 }
 
-extension NetworkManagerSpy: URLSessionDataDelegate {
-    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
-        self.receivedData = data
-    }
-}
-
 public enum NetworkSpyError: Error {
     case cancelled
 }