Add support for upload tasks

DVR/Session.swift

@@ -128,17 +128,19 @@ public class Session: NSURLSession {
 
         completedInteractions.append(interaction)
 
-        if !recording && outstandingTasks.count == 0 {
-            finishRecording()
-        }
-
         if let delegate = delegate as? NSURLSessionDataDelegate, task = task as? NSURLSessionDataTask, data = interaction.responseData {
             delegate.URLSession?(self, dataTask: task, didReceiveData: data)
         }
 
         if let delegate = delegate as? NSURLSessionTaskDelegate {
+            let bytes = Int64(interaction.responseData?.length ?? 0)
+            delegate.URLSession?(self, task: task, didSendBodyData: bytes, totalBytesSent: bytes, totalBytesExpectedToSend: bytes)
             delegate.URLSession?(self, task: task, didCompleteWithError: nil)
         }
+
+        if !recording && outstandingTasks.count == 0 {
+            finishRecording()
+        }
     }
 
 
@@ -162,7 +164,7 @@ public class Session: NSURLSession {
         var modifiedRequest = backingSession.configuration.HTTPAdditionalHeaders.map(request.requestByAppendingHeaders) ?? request
         modifiedRequest = data.map(modifiedRequest.requestWithBody) ?? modifiedRequest
         let task = SessionUploadTask(session: self, request: modifiedRequest, completion: completionHandler)
-        addTask(task.dataTask)
+        addTask(task)
         return task
     }
 

DVR/SessionDataTask.swift

@@ -13,7 +13,23 @@ class SessionDataTask: NSURLSessionDataTask {
     let request: NSURLRequest
     let completion: Completion?
     private let queue = dispatch_queue_create("com.venmo.DVR.sessionDataTaskQueue", nil)
-    private var interaction: Interaction?
+    internal var interaction: Interaction?
+    private var backingTask: NSURLSessionTask?
+
+    private var _taskDescription: String?
+    override var taskDescription: String? {
+        get {
+            return _taskDescription
+        }
+        set {
+            _taskDescription = newValue
+        }
+    }
+
+    private var _taskIdentifier: Int?
+    override var taskIdentifier: Int {
+        return _taskIdentifier ?? 0
+    }
 
     override var response: NSURLResponse? {
         return interaction?.response
@@ -22,9 +38,10 @@ class SessionDataTask: NSURLSessionDataTask {
 
     // MARK: - Initializers
 
-    init(session: Session, request: NSURLRequest, completion: (Completion)? = nil) {
+    init(session: Session, request: NSURLRequest, backingTask: NSURLSessionTask? = nil, completion: (Completion)? = nil) {
         self.session = session
         self.request = request
+        self.backingTask = backingTask
         self.completion = completion
     }
 
@@ -47,7 +64,7 @@ class SessionDataTask: NSURLSessionDataTask {
                     completion(interaction.responseData, interaction.response, nil)
                 }
             }
-            session.finishTask(self, interaction: interaction, playback: true)
+            session.finishTask(self.backingTask ?? self, interaction: interaction, playback: true)
             return
         }
 
@@ -81,9 +98,12 @@ class SessionDataTask: NSURLSessionDataTask {
             }
 
             // Create interaction
-            this.interaction = Interaction(request: this.request, response: response, responseData: data)
-            this.session.finishTask(this, interaction: this.interaction!, playback: false)
+            let interaction = Interaction(request: this.request, response: response, responseData: data)
+            this.interaction = interaction
+            this.session.finishTask(this.backingTask ?? this, interaction: interaction, playback: false)
         }
+
+        _taskIdentifier = task.taskIdentifier
         task.resume()
     }
 }

DVR/SessionUploadTask.swift

@@ -9,15 +9,24 @@ class SessionUploadTask: NSURLSessionUploadTask {
     weak var session: Session!
     let request: NSURLRequest
     let completion: Completion?
-    let dataTask: SessionDataTask
+    var dataTask: SessionDataTask!
+
+    override var response: NSURLResponse? {
+        return dataTask.interaction?.response
+    }
+
+    override var taskIdentifier: Int {
+        return dataTask.taskIdentifier
+    }
 
     // MARK: - Initializers
 
     init(session: Session, request: NSURLRequest, completion: Completion? = nil) {
         self.session = session
         self.request = request
         self.completion = completion
-        dataTask = SessionDataTask(session: session, request: request, completion: completion)
+        super.init()
+        dataTask = SessionDataTask(session: session, request: request, backingTask: self, completion: completion)
     }
 
     // MARK: - NSURLSessionTask

DVR/Tests/SessionTests.swift

@@ -18,6 +18,7 @@ class SessionTests: XCTestCase {
     func testDataTask() {
         let request = NSURLRequest(URL: NSURL(string: "http://example.com")!)
         let dataTask = session.dataTaskWithRequest(request)
+        dataTask.taskDescription = "description"
 
         XCTAssert(dataTask is SessionDataTask)
 
@@ -138,6 +139,7 @@ class SessionTests: XCTestCase {
         class Delegate: NSObject, NSURLSessionTaskDelegate {
             let expectation: XCTestExpectation
             var response: NSURLResponse?
+            var bytes: Int64?
 
             init(expectation: XCTestExpectation) {
                 self.expectation = expectation
@@ -147,6 +149,10 @@ class SessionTests: XCTestCase {
                 response = task.response
                 expectation.fulfill()
             }
+
+            @objc private func URLSession(session: NSURLSession, task: NSURLSessionTask, didSendBodyData bytesSent: Int64, totalBytesSent: Int64, totalBytesExpectedToSend: Int64) {
+                bytes = bytesSent
+            }
         }
 
         let expectation = expectationWithDescription("didCompleteWithError")
@@ -160,6 +166,7 @@ class SessionTests: XCTestCase {
         task.resume()
 
         waitForExpectationsWithTimeout(1, handler: nil)
+        XCTAssertNotNil(delegate.bytes)
     }
 
     func testDataDelegate() {

DVR/Tests/SessionUploadTests.swift

@@ -65,6 +65,37 @@ class SessionUploadTests: XCTestCase {
         waitForExpectationsWithTimeout(4, handler: nil)
     }
 
+    func testUploadDelegate() {
+        class Delegate: NSObject, NSURLSessionDataDelegate {
+            var task: NSURLSessionTask?
+            let expectation: XCTestExpectation
+
+            init(expectation: XCTestExpectation) {
+                self.expectation = expectation
+            }
+
+            @objc func URLSession(session: NSURLSession, dataTask: NSURLSessionDataTask, didReceiveData data: NSData) {
+                task = dataTask
+                expectation.fulfill()
+            }
+        }
+
+        let expectation = expectationWithDescription("didCompleteWithError")
+        let delegate = Delegate(expectation: expectation)
+        let config = NSURLSessionConfiguration.defaultSessionConfiguration()
+        let backingSession = NSURLSession(configuration: config, delegate: delegate, delegateQueue: nil)
+        let session = Session(cassetteName: "upload-data", backingSession: backingSession)
+        session.recordingEnabled = false
+
+        let data = encodeMultipartBody(NSData(contentsOfURL: testFile)!, parameters: [:])
+
+        let task = session.uploadTaskWithRequest(request, fromData: data)
+        task.resume()
+
+        waitForExpectationsWithTimeout(1, handler: nil)
+        XCTAssertEqual(task, delegate.task)
+    }
+
     // MARK: Helpers
 
     func encodeMultipartBody(data: NSData, parameters: [String: AnyObject]) -> NSData {