Add missing delegate method invocations with upload tasks(SR-5516)

Author: Pushkar Kulkarni

Foundation/NSURLSession/http/HTTPURLProtocol.swift

@@ -14,6 +14,8 @@ internal class _HTTPURLProtocol: URLProtocol {
 
     fileprivate var easyHandle: _EasyHandle!
     fileprivate var totalDownloaded = 0
+    fileprivate var totalUploaded: Int64 = 0
+    fileprivate var bodyLength: Int64 = 0
     fileprivate var tempFileURL: URL
 
     public required init(task: URLSessionTask, cachedResponse: CachedURLResponse?, client: URLProtocolClient?) {
@@ -127,6 +129,7 @@ fileprivate extension _HTTPURLProtocol {
                 set(requestBodyLength: .noBody)
             case (_, .some(let length)):
                 set(requestBodyLength: .length(length))
+                bodyLength = Int64(length)
             case (_, .none):
                 set(requestBodyLength: .unknown)
             }
@@ -495,6 +498,16 @@ extension _HTTPURLProtocol: _EasyHandleDelegate {
         }
     }
 
+    fileprivate func notifyDelegate(aboutSentData count: Int64) {
+        guard let session = self.task?.session as? URLSession else { fatalError() }
+        guard case .taskDelegate(let delegate) = session.behaviour(for: self.task!), self.task is URLSessionUploadTask else { return }
+        totalUploaded += count
+        session.delegateQueue.addOperation {
+            delegate.urlSession(session, task: self.task!, didSendBodyData: count,
+                totalBytesSent: self.totalUploaded, totalBytesExpectedToSend: self.bodyLength)
+        }
+    }
+
     func fill(writeBuffer buffer: UnsafeMutableBufferPointer<Int8>) -> _EasyHandle._WriteBufferResult {
         guard case .transferInProgress(let ts) = internalState else { fatalError("Requested to fill write buffer, but transfer isn't in progress.") }
         guard let source = ts.requestBodySource else { fatalError("Requested to fill write buffer, but transfer state has no body source.") }
@@ -503,6 +516,7 @@ extension _HTTPURLProtocol: _EasyHandleDelegate {
             copyDispatchData(data, infoBuffer: buffer)
             let count = data.count
             assert(count > 0)
+            notifyDelegate(aboutSentData: Int64(count))
             return .bytes(count)
         case .done:
             return .bytes(0)

TestFoundation/HTTPServer.swift

@@ -330,7 +330,7 @@ public class TestURLSessionServer {
     }
 
     func process(request: _HTTPRequest) -> _HTTPResponse {
-        if request.method == .GET || request.method == .POST {
+        if request.method == .GET || request.method == .POST || request.method == .PUT {
             return getResponse(request: request)
         } else {
             fatalError("Unsupported method!")
@@ -339,6 +339,12 @@ public class TestURLSessionServer {
 
     func getResponse(request: _HTTPRequest) -> _HTTPResponse {
         let uri = request.uri
+
+        if uri == "/upload" {
+            let text = "Upload completed!"
+            return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.characters.count)", body: text)
+        }
+
         if uri == "/country.txt" {
             let text = capitals[String(uri.characters.dropFirst())]!
             return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.characters.count)", body: text)

TestFoundation/TestNSURLSession.swift

@@ -41,6 +41,7 @@ class TestURLSession : LoopbackServerTest {
             ("test_outOfRangeButCorrectlyFormattedHTTPCode", test_outOfRangeButCorrectlyFormattedHTTPCode),
             ("test_missingContentLengthButStillABody", test_missingContentLengthButStillABody),
             ("test_illegalHTTPServerResponses", test_illegalHTTPServerResponses),
+            ("test_simpleUploadWithDelegate", test_simpleUploadWithDelegate),
         ]
     }
 
@@ -429,6 +430,21 @@ class TestURLSession : LoopbackServerTest {
             waitForExpectations(timeout: 12)
         }
     }
+
+    func test_simpleUploadWithDelegate() {
+        let delegate = HTTPUploadDelegate()
+        let session = URLSession(configuration: .default, delegate: delegate, delegateQueue: nil)
+        let urlString = "http://127.0.0.1:\(TestURLSession.serverPort)/upload"
+        var request = URLRequest(url: URL(string: urlString)!)
+        request.httpMethod = "PUT"
+
+        delegate.uploadCompletedExpectation = expectation(description: "PUT \(urlString): Upload data")
+
+        let fileData = Data(count: 16*1024)
+        let task = session.uploadTask(with: request, from: fileData)
+        task.resume()
+        waitForExpectations(timeout: 20)
+    }
 }
 
 class SessionDelegate: NSObject, URLSessionDelegate {
@@ -619,3 +635,21 @@ extension HTTPRedirectionDataTask : URLSessionTaskDelegate {
         completionHandler(request)
     }
 }
+
+class HTTPUploadDelegate: NSObject {
+    var uploadCompletedExpectation: XCTestExpectation!
+    var totalBytesSent: Int64 = 0
+}
+
+extension HTTPUploadDelegate: URLSessionTaskDelegate {
+    func urlSession(_ session: URLSession, task: URLSessionTask, didSendBodyData bytesSent: Int64, totalBytesSent: Int64, totalBytesExpectedToSend: Int64) {
+        self.totalBytesSent = totalBytesSent
+    }
+}
+
+extension HTTPUploadDelegate: URLSessionDataDelegate {
+    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
+        XCTAssertEqual(self.totalBytesSent, 16*1024)
+        uploadCompletedExpectation.fulfill()
+    }
+}