diff --git a/DVR/Session.swift b/DVR/Session.swift index a09fa2a..704bbe0 100644 --- a/DVR/Session.swift +++ b/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 } diff --git a/DVR/SessionDataTask.swift b/DVR/SessionDataTask.swift index 9e0c32c..322228d 100644 --- a/DVR/SessionDataTask.swift +++ b/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() } } diff --git a/DVR/SessionUploadTask.swift b/DVR/SessionUploadTask.swift index b6b3b78..8d91883 100644 --- a/DVR/SessionUploadTask.swift +++ b/DVR/SessionUploadTask.swift @@ -9,7 +9,15 @@ 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 @@ -17,7 +25,8 @@ class SessionUploadTask: NSURLSessionUploadTask { 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 diff --git a/DVR/Tests/SessionTests.swift b/DVR/Tests/SessionTests.swift index 3ff4e05..c27c24a 100644 --- a/DVR/Tests/SessionTests.swift +++ b/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() { diff --git a/DVR/Tests/SessionUploadTests.swift b/DVR/Tests/SessionUploadTests.swift index 99b841c..5de9926 100644 --- a/DVR/Tests/SessionUploadTests.swift +++ b/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 {