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

@@ -164,7 +164,8 @@ open class NetworkManager: NSObject {
     /// - 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,
+        completion: CompletionHandler? = nil
     ) -> Cancelable? {
         guard let urlRequest = try? buildUrlRequest(request: request) else { return nil }
 
@@ -178,7 +179,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.register(task, progress: progress, completion: completion)
         return task
     }
 

Sources/YNetwork/NetworkManager/Progress/FileProgress.swift

@@ -15,6 +15,8 @@
 ///
 /// Guaranteed to be called back on the main thread
 public typealias ProgressHandler = (Percentage) -> Void
+public typealias CompletionHandler = (Data) -> Void
+public typealias CancellationHandler = () -> Void
 
 /// Asynchronous completion handler that reports the status of request.
 ///
@@ -29,6 +31,7 @@
 internal class FileProgress: NSObject {
     /// Stores the progress handlers to be called, keyed by unique task identifier
     private var progressHandlersByTaskID: [Int: ProgressHandler] = [:]
+    private var completionHandlerByTaskID: [Int: CompletionHandler] = [:]
     private var downloadHandlersByTaskID: [Int: FileDownloadHandler] = [:]
 
     /// Updates the progress handler for the specified task with the percentage value
@@ -41,6 +44,17 @@
             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 = completionHandlerByTaskID[taskIdentifier] else { return }
+        DispatchQueue.main.async {
+            completionHandler(data)
+        }
+    }
 
     /// Updates the request status for the specified task with the file URL
     /// - Parameters:
@@ -57,10 +71,17 @@
     /// - Parameters:
     ///   - cancelable: optional cancelable task
     ///   - progress: optional progress handler
-    func register(_ cancelable: Cancelable?, progress: ProgressHandler?) {
+    func register(
+        _ cancelable: Cancelable?,
+        progress: ProgressHandler?,
+        completion: CompletionHandler? = nil
+    ) {
         guard let task = cancelable as? URLSessionTask,
               let progress = progress else { return }
         progressHandlersByTaskID[task.taskIdentifier] = progress
+        if let completion = completion {
+            completionHandlerByTaskID[task.taskIdentifier] = completion
+        }
     }
 
     /// Registers a data task for file progress.
@@ -84,6 +105,12 @@
         progressHandlersByTaskID.removeValue(forKey: taskIdentifier)
         downloadHandlersByTaskID.removeValue(forKey: taskIdentifier)
     }
+
+    /// Unregisters a completion handler once final response received
+    /// - Parameter taskIdentifier: unique task identifier
+    func unregisterCompletion(forKey taskIdentifier: Int) {
+        completionHandlerByTaskID.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)
+        unregisterCompletion(forKey: dataTask.taskIdentifier)
+    }
 }