Implement SessionTask.State

venmo · May 2, 2019 · 0cc3eb1 · 0cc3eb1
1 parent 341154c
commit 0cc3eb1
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/Sources/DVR/SessionDataTask.swift b/Sources/DVR/SessionDataTask.swift
@@ -31,13 +31,20 @@ final class SessionDataTask: URLSessionDataTask {
 
     // MARK: - URLSessionTask
 
+    var _state: URLSessionTask.State?
+    override var state: URLSessionTask.State {
+        return _state ?? .suspended
+    }
+
     override func cancel() {
         // Don't do anything
     }
 
     override func resume() {
         let cassette = session.cassette
 
+        _state = .running
+
         // Find interaction
         if let interaction = session.cassette?.interactionForRequest(request) {
             self.interaction = interaction
@@ -47,7 +54,15 @@ final class SessionDataTask: URLSessionDataTask {
                     completion(interaction.responseData, interaction.response, nil)
                 }
             }
-            session.finishTask(self, interaction: interaction, playback: true)
+
+            queue.async { [weak self] in
+                guard let self = self else {
+                    return
+                }
+
+                self._state = .completed
+                self.session.finishTask(self, interaction: interaction, playback: true)
+            }
             return
         }
 
@@ -76,6 +91,8 @@ final class SessionDataTask: URLSessionDataTask {
                 this.completion?(data, response, nil)
             }
 
+            this._state = .completed
+
             // Create interaction
             this.interaction = Interaction(request: this.request, response: response, responseData: data)
             this.session.finishTask(this, interaction: this.interaction!, playback: false)