From c0d6380e852fa68d1552f81277dff7382a694f95 Mon Sep 17 00:00:00 2001
From: Pushkar Kulkarni <pushkar.nk@in.ibm.com>
Date: Mon, 19 Jun 2017 18:51:09 +0530
Subject: [PATCH] Fix for test_finishTaskAndInvalidate failure

---
 Foundation/NSURLProtocol.swift                | 27 ++++++++++---------
 .../NSURLSession/http/HTTPURLProtocol.swift   |  1 -
 TestFoundation/TestNSURLSession.swift         |  2 +-
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/Foundation/NSURLProtocol.swift b/Foundation/NSURLProtocol.swift
index 11a23408fb..face732a9c 100644
--- a/Foundation/NSURLProtocol.swift
+++ b/Foundation/NSURLProtocol.swift
@@ -156,23 +156,24 @@ internal class _ProtocolClient : NSObject, URLProtocolClient {
         guard let session = task.session as? URLSession else { fatalError() }
         switch session.behaviour(for: task) {
         case .taskDelegate(let delegate):
-            guard let s = session as? URLSession else { fatalError() }
-            s.delegateQueue.addOperation {
-                delegate.urlSession(s, task: task, didCompleteWithError: nil)
+            session.delegateQueue.addOperation {
+                delegate.urlSession(session, task: task, didCompleteWithError: nil)
                 task.state = .completed
+                session.taskRegistry.remove(task)
             }
         case .noDelegate:
             task.state = .completed
+            session.taskRegistry.remove(task)
         case .dataCompletionHandler(let completion):
             let data = Data()
             guard let client = `protocol`.client else { fatalError() }
             client.urlProtocol(`protocol`, didLoad: data)
             return
         case .downloadCompletionHandler(let completion):
-            guard let s = session as? URLSession else { fatalError() }
-            s.delegateQueue.addOperation {
+            session.delegateQueue.addOperation {
                 completion(task.currentRequest?.url, task.response, nil)
                 task.state = .completed
+                session.taskRegistry.remove(task)
             }
         }
     }
@@ -194,6 +195,7 @@ internal class _ProtocolClient : NSObject, URLProtocolClient {
             s.delegateQueue.addOperation {
                 completion(data, task.response, nil)
                 task.state = .completed
+                s.taskRegistry.remove(task)
             }
         default: return
         }
@@ -204,24 +206,25 @@ internal class _ProtocolClient : NSObject, URLProtocolClient {
         guard let session = task.session as? URLSession else { fatalError() }
         switch session.behaviour(for: task) {
         case .taskDelegate(let delegate):
-            guard let s = session as? URLSession else { fatalError() }
-            s.delegateQueue.addOperation {
-                delegate.urlSession(s, task: task, didCompleteWithError: error as Error)
+            session.delegateQueue.addOperation {
+                delegate.urlSession(session, task: task, didCompleteWithError: error as Error)
                 task.state = .completed
+                session.taskRegistry.remove(task)
             }
         case .noDelegate:
             task.state = .completed
+            session.taskRegistry.remove(task)
         case .dataCompletionHandler(let completion):
-            guard let s = session as? URLSession else { fatalError() }
-            s.delegateQueue.addOperation {
+            session.delegateQueue.addOperation {
                 completion(nil, nil, error)
                 task.state = .completed
+                session.taskRegistry.remove(task)
             }
         case .downloadCompletionHandler(let completion):
-            guard let s = session as? URLSession else { fatalError() }
-            s.delegateQueue.addOperation {
+            session.delegateQueue.addOperation {
                 completion(nil, nil, error)
                 task.state = .completed
+                session.taskRegistry.remove(task)
             }
         }
     }
diff --git a/Foundation/NSURLSession/http/HTTPURLProtocol.swift b/Foundation/NSURLSession/http/HTTPURLProtocol.swift
index 017eaae7bf..bf20a2303f 100644
--- a/Foundation/NSURLSession/http/HTTPURLProtocol.swift
+++ b/Foundation/NSURLSession/http/HTTPURLProtocol.swift
@@ -55,7 +55,6 @@ internal class _HTTPURLProtocol: URLProtocol {
             self.internalState = .transferFailed
             guard let error = self.task?.error else { fatalError() }
             completeTask(withError: error)
-            return
         }
     }
 
diff --git a/TestFoundation/TestNSURLSession.swift b/TestFoundation/TestNSURLSession.swift
index 3bb16f60f3..e4ffbdea3b 100644
--- a/TestFoundation/TestNSURLSession.swift
+++ b/TestFoundation/TestNSURLSession.swift
@@ -30,7 +30,7 @@ class TestURLSession : XCTestCase {
 //            ("test_downloadTaskWithURLRequest", test_downloadTaskWithURLRequest),
             ("test_downloadTaskWithRequestAndHandler", test_downloadTaskWithRequestAndHandler),
             ("test_downloadTaskWithURLAndHandler", test_downloadTaskWithURLAndHandler),
-//            ("test_finishTaskAndInvalidate", test_finishTasksAndInvalidate),
+            ("test_finishTaskAndInvalidate", test_finishTasksAndInvalidate),
 //            ("test_taskError", test_taskError),
             ("test_taskCopy", test_taskCopy),
 //            ("test_cancelTask", test_cancelTask),