XCTestCase+URLSessionDataTask

XCTestCase+URLSessionDataTask.swift

import XCTest

extension XCTestCase {
    
    func expectationDataTask(dataTask: NSURLSessionDataTask) -> XCTestExpectation {
        dataTask.resume()
        
        let expectation = keyValueObservingExpectationForObject(dataTask, keyPath: "state") {  _, _ in
            return dataTask.state == .Completed || dataTask.state == .Canceling
        }
        return expectation
    }
    
    func waitForExpectationDataTask(dataTask: NSURLSessionDataTask) -> XCTestExpectation {
        let expectation = expectationDataTask(dataTask)
        
        let timeout = dataTask.currentRequest?.timeoutInterval ?? 65
        waitForExpectationsWithTimeout(timeout, handler: nil)
        
        return expectation
    }
}