Merge pull request #1121 from pushkarnk/uploadtask-delegate

Author: swift-ci

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 requestBodyLength: 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))
+                requestBodyLength = Int64(length)
             case (_, .none):
                 set(requestBodyLength: .unknown)
             }
@@ -495,6 +498,16 @@ extension _HTTPURLProtocol: _EasyHandleDelegate {
         }
     }
 
+    fileprivate func notifyDelegate(aboutUploadedData count: Int64) {
+        let session = self.task?.session as! URLSession
+        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.requestBodyLength)
+        }
+    }
+
     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(aboutUploadedData: 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,14 +339,20 @@ 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.data(using: .utf8)!.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)
+            return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.data(using: .utf8)!.count)", body: text)
         }
 
         if uri == "/requestHeaders" {
             let text = request.getCommaSeparatedHeaders()
-            return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.characters.count)", body: text)
+            return _HTTPResponse(response: .OK, headers: "Content-Length: \(text.data(using: .utf8)!.count)", body: text)
         }
 
 	if uri == "/UnitedStates" {

TestFoundation/TestNSURLSession.swift

@@ -42,6 +42,7 @@ class TestURLSession : LoopbackServerTest {
             ("test_missingContentLengthButStillABody", test_missingContentLengthButStillABody),
             ("test_illegalHTTPServerResponses", test_illegalHTTPServerResponses),
             ("test_dataTaskWithSharedDelegate", test_dataTaskWithSharedDelegate),
+            ("test_simpleUploadWithDelegate", test_simpleUploadWithDelegate),
         ]
     }
 
@@ -443,6 +444,21 @@ class TestURLSession : LoopbackServerTest {
         dataTask.resume()
         waitForExpectations(timeout: 20)
     }
+
+    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 SharedDelegate: NSObject {
@@ -649,3 +665,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()
+    }
+}