Add receive response delegate for response disposition control

Tests/KingfisherTests/ImageDownloaderTests.swift

@@ -547,9 +547,9 @@ class ImageDownloaderTests: XCTestCase {
     
     
     func testSessionDelegate() {
-        class ExtensionDelegate:SessionDelegate {
+        class ExtensionDelegate: SessionDelegate {
             //'exp' only for test
-            public let exp:XCTestExpectation
+            public let exp: XCTestExpectation
             init(_ expectation:XCTestExpectation) {
                 exp = expectation
             }
@@ -566,6 +566,78 @@ class ImageDownloaderTests: XCTestCase {
         }
         waitForExpectations(timeout: 3, handler: nil)
     }
+
+    func testDownloaderReceiveResponsePass() {
+
+        let exp = expectation(description: #function)
+
+        let url = testURLs[0]
+        stub(url, data: testImageData, headers: ["someKey": "someValue"])
+
+        let handler = TaskResponseCompletion()
+        let obj = NSObject()
+        handler.onReceiveResponse.delegate(on: obj) { (obj, response) in
+            guard let httpResponse = response as? HTTPURLResponse else {
+                XCTFail("Should be an HTTP response.")
+                return .cancel
+            }
+            XCTAssertEqual(httpResponse.statusCode, 200)
+            XCTAssertEqual(httpResponse.url, url)
+            XCTAssertEqual(httpResponse.allHeaderFields["someKey"] as? String, "someValue")
+
+            return .allow
+        }
+
+        downloader.delegate = handler
+        downloader.downloadImage(with: url) { result in
+            XCTAssertNotNil(result.value)
+            XCTAssertNil(result.error)
+
+            self.downloader.delegate = nil
+            // hold delegate
+            _ = handler
+            exp.fulfill()
+        }
+        waitForExpectations(timeout: 3, handler: nil)
+    }
+
+    func testDownloaderReceiveResponseFailure() {
+        let exp = expectation(description: #function)
+
+        let url = testURLs[0]
+        stub(url, data: testImageData, headers: ["someKey": "someValue"])
+
+        let handler = TaskResponseCompletion()
+        let obj = NSObject()
+        handler.onReceiveResponse.delegate(on: obj) { (obj, response) in
+            guard let httpResponse = response as? HTTPURLResponse else {
+                XCTFail("Should be an HTTP response.")
+                return .cancel
+            }
+            XCTAssertEqual(httpResponse.statusCode, 200)
+            XCTAssertEqual(httpResponse.url, url)
+            XCTAssertEqual(httpResponse.allHeaderFields["someKey"] as? String, "someValue")
+
+            return .cancel
+        }
+
+        downloader.delegate = handler
+        downloader.downloadImage(with: url) { result in
+            XCTAssertNil(result.value)
+            XCTAssertNotNil(result.error)
+
+            if case .responseError(reason: .cancelledByDelegate) = result.error! {
+            } else {
+                XCTFail()
+            }
+
+            self.downloader.delegate = nil
+            // hold delegate
+            _ = handler
+            exp.fulfill()
+        }
+        waitForExpectations(timeout: 3, handler: nil)
+    }
 }
 
 class URLNilDataModifier: ImageDownloaderDelegate {
@@ -580,6 +652,19 @@ class TaskNilDataModifier: ImageDownloaderDelegate {
     }
 }
 
+class TaskResponseCompletion: ImageDownloaderDelegate {
+
+    let onReceiveResponse = Delegate<URLResponse, URLSession.ResponseDisposition>()
+
+    func imageDownloader(
+        _ downloader: ImageDownloader,
+        didReceive response: URLResponse,
+        completionHandler: @escaping (URLSession.ResponseDisposition) -> Void
+    ) {
+        completionHandler(onReceiveResponse.call(response)!)
+    }
+}
+
 class URLModifier: ImageDownloadRequestModifier {
     var url: URL? = nil
     func modified(for request: URLRequest) -> URLRequest? {

Tests/KingfisherTests/Utils/StubHelpers.swift

@@ -27,16 +27,29 @@
 import Foundation
 
 @discardableResult
-func stub(_ url: URL, data: Data, statusCode: Int = 200, length: Int? = nil) -> LSStubResponseDSL {
-    var stubResult = stubRequest("GET", url.absoluteString as NSString).andReturn(statusCode)?.withBody(data as NSData)
+func stub(_ url: URL,
+          data: Data,
+          statusCode: Int = 200,
+          length: Int? = nil,
+          headers: [String: String] = [:]
+) -> LSStubResponseDSL {
+    var stubResult = stubRequest("GET", url.absoluteString as NSString)
+        .andReturn(statusCode)?
+        .withHeaders(headers)?
+        .withBody(data as NSData)
     if let length = length {
         stubResult = stubResult?.withHeader("Content-Length", "\(length)")
     }
     return stubResult!
 }
 
-func delayedStub(_ url: URL, data: Data, statusCode: Int = 200, length: Int? = nil) -> LSStubResponseDSL {
-    let result = stub(url, data: data, statusCode: statusCode, length: length)
+func delayedStub(_ url: URL,
+                 data: Data,
+                 statusCode: Int = 200,
+                 length: Int? = nil,
+                 headers: [String: String] = [:]
+) -> LSStubResponseDSL {
+    let result = stub(url, data: data, statusCode: statusCode, length: length, headers: headers)
     return result.delay()!
 }