Notify delegate when task finishes

DVR/Session.swift

@@ -17,6 +17,9 @@ public class Session: NSURLSession {
     private var completedInteractions = [Interaction]()
     private var completionBlock: (Void -> Void)?
 
+    override public var delegate: NSURLSessionDelegate? {
+        return backingSession.delegate
+    }
 
     // MARK: - Initializers
 
@@ -110,6 +113,14 @@ public class Session: NSURLSession {
         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 {
+            delegate.URLSession?(self, task: task, didCompleteWithError: nil)
+        }
     }
 
 

DVR/SessionDataTask.swift

@@ -13,6 +13,11 @@ class SessionDataTask: NSURLSessionDataTask {
     let request: NSURLRequest
     let completion: Completion?
     private let queue = dispatch_queue_create("com.venmo.DVR.sessionDataTaskQueue", nil)
+    private var interaction: Interaction?
+
+    override var response: NSURLResponse? {
+        return interaction?.response
+    }
 
 
     // MARK: - Initializers
@@ -35,6 +40,7 @@ class SessionDataTask: NSURLSessionDataTask {
 
         // Find interaction
         if let interaction = session.cassette?.interactionForRequest(request) {
+            self.interaction = interaction
             // Forward completion
             if let completion = completion {
                 dispatch_async(queue) {
@@ -75,8 +81,8 @@ class SessionDataTask: NSURLSessionDataTask {
             }
 
             // Create interaction
-            let interaction = Interaction(request: this.request, response: response, responseData: data)
-            this.session.finishTask(this, interaction: interaction, playback: false)
+            this.interaction = Interaction(request: this.request, response: response, responseData: data)
+            this.session.finishTask(this, interaction: this.interaction!, playback: false)
         }
         task.resume()
     }

DVR/Tests/SessionTests.swift

@@ -133,4 +133,58 @@ class SessionTests: XCTestCase {
 
         waitForExpectationsWithTimeout(1, handler: nil)
     }
+
+    func testTaskDelegate() {
+        class Delegate: NSObject, NSURLSessionTaskDelegate {
+            let expectation: XCTestExpectation
+            var response: NSURLResponse?
+
+            init(expectation: XCTestExpectation) {
+                self.expectation = expectation
+            }
+
+            @objc private func URLSession(session: NSURLSession, task: NSURLSessionTask, didCompleteWithError error: NSError?) {
+                response = task.response
+                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: "example", backingSession: backingSession)
+        session.recordingEnabled = false
+
+        let task = session.dataTaskWithRequest(request)
+        task.resume()
+
+        waitForExpectationsWithTimeout(1, handler: nil)
+    }
+
+    func testDataDelegate() {
+        class Delegate: NSObject, NSURLSessionDataDelegate {
+            let expectation: XCTestExpectation
+
+            init(expectation: XCTestExpectation) {
+                self.expectation = expectation
+            }
+
+            @objc func URLSession(session: NSURLSession, dataTask: NSURLSessionDataTask, didReceiveData data: NSData) {
+                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: "example", backingSession: backingSession)
+        session.recordingEnabled = false
+
+        let task = session.dataTaskWithRequest(request)
+        task.resume()
+
+        waitForExpectationsWithTimeout(1, handler: nil)
+    }
 }