Implement cancellabe button image task

Sources/UIButton+Kingfisher.swift

@@ -209,6 +209,9 @@ extension UIButton {
                 
                 dispatch_async_safely_main_queue {
                     if let sSelf = self {
+                        
+                        sSelf.kf_setImageTask(nil)
+                        
                         if imageURL == sSelf.kf_webURLForState(state) && image != nil {
                             sSelf.setImage(image, forState: state)
                         }
@@ -216,6 +219,8 @@ extension UIButton {
                     }
                 }
             })
+        
+        kf_setImageTask(task)
         return task
     }
     
@@ -248,6 +253,9 @@ extension UIButton {
 }
 
 private var lastURLKey: Void?
+private var imageTaskKey: Void?
+
+// MARK: - Runtime for UIButton image
 extension UIButton {
     /**
     Get the image URL binded to this button for a specified state. 
@@ -276,6 +284,14 @@ extension UIButton {
     private func kf_setWebURLs(URLs: NSMutableDictionary) {
         objc_setAssociatedObject(self, &lastURLKey, URLs, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
     }
+    
+    private var kf_imageTask: RetrieveImageTask? {
+        return objc_getAssociatedObject(self, &imageTaskKey) as? RetrieveImageTask
+    }
+    
+    private func kf_setImageTask(task: RetrieveImageTask?) {
+        objc_setAssociatedObject(self, &imageTaskKey, task, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
+    }
 }
 
 /**
@@ -456,7 +472,11 @@ extension UIButton {
             },
             completionHandler: { [weak self] image, error, cacheType, imageURL in
                 dispatch_async_safely_main_queue {
+                    
                     if let sSelf = self {
+                        
+                        sSelf.kf_setBackgroundImageTask(nil)
+                        
                         if imageURL == sSelf.kf_backgroundWebURLForState(state) && image != nil {
                             sSelf.setBackgroundImage(image, forState: state)
                         }
@@ -464,6 +484,8 @@ extension UIButton {
                     }
                 }
             })
+        
+        kf_setBackgroundImageTask(task)
         return task
     }
     
@@ -497,6 +519,9 @@ extension UIButton {
 }
 
 private var lastBackgroundURLKey: Void?
+private var backgroundImageTaskKey: Void?
+    
+// MARK: - Runtime for UIButton background image
 extension UIButton {
     /**
     Get the background image URL binded to this button for a specified state.
@@ -525,8 +550,35 @@ extension UIButton {
     private func kf_setBackgroundWebURLs(URLs: NSMutableDictionary) {
         objc_setAssociatedObject(self, &lastBackgroundURLKey, URLs, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
     }
+    
+    private var kf_backgroundImageTask: RetrieveImageTask? {
+        return objc_getAssociatedObject(self, &backgroundImageTaskKey) as? RetrieveImageTask
+    }
+    
+    private func kf_setBackgroundImageTask(task: RetrieveImageTask?) {
+        objc_setAssociatedObject(self, &backgroundImageTaskKey, task, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
+    }
 }
 
+// MARK: - Cancel image download tasks.
+extension UIButton {
+    /**
+     Cancel the image download task bounded to the image view if it is running.
+     Nothing will happen if the downloading has already finished.
+     */
+    public func kf_cancelImageDownloadTask() {
+        kf_imageTask?.downloadTask?.cancel()
+    }
+    
+    /**
+     Cancel the background image download task bounded to the image view if it is running.
+     Nothing will happen if the downloading has already finished.
+     */
+    public func kf_cancelBackgroundImageDownloadTask() {
+        kf_backgroundImageTask?.downloadTask?.cancel()
+    }
+}
+    
 #elseif os(OSX)
 
 import AppKit

Tests/KingfisherTests/UIButtonExtensionTests.swift

@@ -109,4 +109,50 @@ class UIButtonExtensionTests: XCTestCase {
         }
         waitForExpectationsWithTimeout(5, handler: nil)
     }
+    
+    func testCacnelImageTask() {
+        let expectation = expectationWithDescription("wait for downloading image")
+        
+        let URLString = testKeys[0]
+        let stub = stubRequest("GET", URLString).andReturn(200).withBody(testImageData).delay()
+        let URL = NSURL(string: URLString)!
+
+        button.kf_setImageWithURL(URL, forState: UIControlState.Highlighted, placeholderImage: nil, optionsInfo: nil, progressBlock: { (receivedSize, totalSize) -> () in
+                XCTFail("Progress block should not be called.")
+            }) { (image, error, cacheType, imageURL) -> () in
+                XCTAssertNotNil(error)
+                XCTAssertEqual(error?.code, NSURLErrorCancelled)
+
+                expectation.fulfill()
+        }
+        dispatch_after(dispatch_time(DISPATCH_TIME_NOW, Int64(Double(NSEC_PER_SEC) * 0.1)), dispatch_get_main_queue()) { () -> Void in
+            self.button.kf_cancelImageDownloadTask()
+            stub.go()
+        }
+
+        waitForExpectationsWithTimeout(5, handler: nil)
+    }
+    
+    func testCacnelBackgroundImageTask() {
+        let expectation = expectationWithDescription("wait for downloading image")
+        
+        let URLString = testKeys[0]
+        let stub = stubRequest("GET", URLString).andReturn(200).withBody(testImageData).delay()
+        let URL = NSURL(string: URLString)!
+        
+        button.kf_setBackgroundImageWithURL(URL, forState: UIControlState.Normal, placeholderImage: nil, optionsInfo: nil, progressBlock: { (receivedSize, totalSize) -> () in
+            XCTFail("Progress block should not be called.")
+            }) { (image, error, cacheType, imageURL) -> () in
+                XCTAssertNotNil(error)
+                XCTAssertEqual(error?.code, NSURLErrorCancelled)
+                
+                expectation.fulfill()
+        }
+        dispatch_after(dispatch_time(DISPATCH_TIME_NOW, Int64(Double(NSEC_PER_SEC) * 0.1)), dispatch_get_main_queue()) { () -> Void in
+            self.button.kf_cancelBackgroundImageDownloadTask()
+            stub.go()
+        }
+        
+        waitForExpectationsWithTimeout(5, handler: nil)
+    }
 }