From b8bc44a4306df060214463f8478c730ec7527c93 Mon Sep 17 00:00:00 2001
From: Andy Durdin <andy.durdin@bitwink.com>
Date: Tue, 10 May 2016 15:37:33 +0100
Subject: [PATCH] Bugfix: task completion block never called when not following
 redirects.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

When `-[NSURLSessionTaskDelegate URLSession:task:willPerformHTTPRedirection:newRequest:completionHandler:]` calls its completion handler with `nil` to indicate that the redirect should not be followed, the completion handler from `-[NSURLSession dataTaskWithRequest:completionHandler:]` is normally called with the actual redirect response itself. But when using OHTTPStubs such that a stub returns a redirect response, the task’s completion handler block is simply never called.

- Ensure `OHHTTPStubsProtocol` finishes redirect responses fully when not following redirects.
- Update `NSURLSessionTests` to exercise this behaviour.
- Fix `AFNetworkingTests` test that had been looking for the incorrect behaviour.
---
 CHANGELOG.md                                  |   3 +
 OHHTTPStubs/Sources/OHHTTPStubs.m             |  61 +++++----
 .../UnitTests/Test Suites/AFNetworkingTests.m |  65 ++++++++--
 .../UnitTests/Test Suites/NSURLSessionTests.m | 118 +++++++++++++++++-
 4 files changed, 207 insertions(+), 40 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 289b381d..7222c913 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,9 @@
 
 ## Master
 
+* Bugfix: task completion block never called when not following redirects.
+  [@adurdin](https://github.com/adurdin), [#175](https://github.com/AliSoftware/OHHTTPStubs/pull/175)
+
 * Declare in the project settings that the library contains swift code.  
   [@rodericj](https://github.com/rodericj), [#173](https://github.com/AliSoftware/OHHTTPStubs/pull/173)
 
diff --git a/OHHTTPStubs/Sources/OHHTTPStubs.m b/OHHTTPStubs/Sources/OHHTTPStubs.m
index dde3a52a..0b575c62 100644
--- a/OHHTTPStubs/Sources/OHHTTPStubs.m
+++ b/OHHTTPStubs/Sources/OHHTTPStubs.m
@@ -418,6 +418,7 @@ - (void)startLoading
         {
             redirectLocationURL = nil;
         }
+        // Notify if a redirection occurred
         if (((responseStub.statusCode > 300) && (responseStub.statusCode < 400)) && redirectLocationURL)
         {
             NSURLRequest* redirectRequest = [NSURLRequest requestWithURL:redirectLocationURL];
@@ -432,39 +433,37 @@ - (void)startLoading
                 }
             }];
         }
-        else
-        {
-            [self executeOnClientRunLoopAfterDelay:responseStub.requestTime block:^{
-                if (!self.stopped)
+        // Send the response (even for redirections)
+        [self executeOnClientRunLoopAfterDelay:responseStub.requestTime block:^{
+            if (!self.stopped)
+            {
+                [client URLProtocol:self didReceiveResponse:urlResponse cacheStoragePolicy:NSURLCacheStorageNotAllowed];
+                if(responseStub.inputStream.streamStatus == NSStreamStatusNotOpen)
                 {
-                    [client URLProtocol:self didReceiveResponse:urlResponse cacheStoragePolicy:NSURLCacheStorageNotAllowed];
-                    if(responseStub.inputStream.streamStatus == NSStreamStatusNotOpen)
-                    {
-                        [responseStub.inputStream open];
-                    }
-                    [self streamDataForClient:client
-                             withStubResponse:responseStub
-                                   completion:^(NSError * error)
-                     {
-                         [responseStub.inputStream close];
-                         NSError *blockError = nil;
-                         if (error==nil)
-                         {
-                             [client URLProtocolDidFinishLoading:self];
-                         }
-                         else
-                         {
-                             [client URLProtocol:self didFailWithError:responseStub.error];
-                             blockError = responseStub.error;
-                         }
-                         if (OHHTTPStubs.sharedInstance.afterStubFinishBlock)
-                         {
-                             OHHTTPStubs.sharedInstance.afterStubFinishBlock(request, self.stub, responseStub, blockError);
-                         }
-                     }];
+                    [responseStub.inputStream open];
                 }
-            }];
-        }
+                [self streamDataForClient:client
+                         withStubResponse:responseStub
+                               completion:^(NSError * error)
+                 {
+                     [responseStub.inputStream close];
+                     NSError *blockError = nil;
+                     if (error==nil)
+                     {
+                         [client URLProtocolDidFinishLoading:self];
+                     }
+                     else
+                     {
+                         [client URLProtocol:self didFailWithError:responseStub.error];
+                         blockError = responseStub.error;
+                     }
+                     if (OHHTTPStubs.sharedInstance.afterStubFinishBlock)
+                     {
+                         OHHTTPStubs.sharedInstance.afterStubFinishBlock(request, self.stub, responseStub, blockError);
+                     }
+                 }];
+            }
+        }];
     } else {
         // Send the canned error
         [self executeOnClientRunLoopAfterDelay:responseStub.responseTime block:^{
diff --git a/OHHTTPStubs/UnitTests/Test Suites/AFNetworkingTests.m b/OHHTTPStubs/UnitTests/Test Suites/AFNetworkingTests.m
index 378e339b..866c8229 100644
--- a/OHHTTPStubs/UnitTests/Test Suites/AFNetworkingTests.m	
+++ b/OHHTTPStubs/UnitTests/Test Suites/AFNetworkingTests.m	
@@ -125,32 +125,35 @@ -(void)test_AFHTTPRequestOperation_redirect
     static const NSTimeInterval kRequestTime = 0.05;
     static const NSTimeInterval kResponseTime = 0.1;
     
-    NSURL* redirectURL = [NSURL URLWithString:@"http://www.iana.org/domains/another/example"];
+    NSURL* redirectURL = [NSURL URLWithString:@"https://httpbin.org/get"];
     [OHHTTPStubs stubRequestsPassingTest:^BOOL(NSURLRequest *request) {
         return YES;
     } withStubResponse:^OHHTTPStubsResponse *(NSURLRequest *request) {
-        return [[OHHTTPStubsResponse responseWithData:[NSData data] statusCode:311 headers:@{@"Location":redirectURL.absoluteString}]
+        return [[OHHTTPStubsResponse responseWithData:[NSData data] statusCode:302 headers:@{@"Location":redirectURL.absoluteString}]
                 requestTime:kRequestTime responseTime:kResponseTime];
     }];
     
+    XCTestExpectation* redirectExpectation = [self expectationWithDescription:@"AFHTTPRequestOperation request was redirected"];
     XCTestExpectation* expectation = [self expectationWithDescription:@"AFHTTPRequestOperation request finished"];
-    
-    NSURLRequest* req = [NSURLRequest requestWithURL:[NSURL URLWithString:@"http://www.iana.org/domains/example/"]];
+
+    NSURLRequest* req = [NSURLRequest requestWithURL:[NSURL URLWithString:@"https://httpbin.org/redirect/1"]];
     AFHTTPSessionManager *manager = [AFHTTPSessionManager manager];
-    [manager setResponseSerializer:[AFHTTPResponseSerializer serializer]];
-    
+    AFHTTPResponseSerializer *serializer = [AFHTTPResponseSerializer serializer];
+    serializer.acceptableStatusCodes = [NSIndexSet indexSetWithIndex:302];
+    [manager setResponseSerializer:serializer];
+
     __block __strong NSURL* url = nil;
     [manager setTaskWillPerformHTTPRedirectionBlock:^NSURLRequest * (NSURLSession * session, NSURLSessionTask * task, NSURLResponse * response, NSURLRequest * request) {
         if (response == nil) {
             return request;
         }
         url = request.URL;
-        [expectation fulfill];
+        [redirectExpectation fulfill];
         return nil;
     }];
     
     [manager GET:req.URL.absoluteString parameters:nil progress:nil success:^(NSURLSessionTask *task, id responseObject) {
-        XCTFail(@"Unexpected response");
+        // Expect the 302 response when the redirection block returns nil (don't follow redirects)
         [expectation fulfill];
     } failure:^(NSURLSessionTask *operation, NSError *error) {
         XCTFail(@"Unexpected network failure");
@@ -162,6 +165,52 @@ -(void)test_AFHTTPRequestOperation_redirect
     XCTAssertEqualObjects(url, redirectURL, @"Unexpected data received");
 }
 
+/*
+ * In order to establish that test_AFHTTPRequestOperation_redirect was incorrect and needed fixing, I needed
+ * to demonstrate identical behaviour--that is, returning the redirect response itself to the success block--
+ * when running without the NSURLProtocol stubbing the request. The test below, if enabled, establishes this,
+ * as it is identical to test_AFHTTPRequestOperation_redirect except that it does not stub the requests.
+ */
+#if 0
+-(void)test_AFHTTPRequestOperation_redirect_baseline
+{
+    NSURL* redirectURL = [NSURL URLWithString:@"https://httpbin.org/get"];
+
+    XCTestExpectation* redirectExpectation = [self expectationWithDescription:@"AFHTTPRequestOperation request was redirected"];
+    XCTestExpectation* expectation = [self expectationWithDescription:@"AFHTTPRequestOperation request finished"];
+
+    NSURLRequest* req = [NSURLRequest requestWithURL:[NSURL URLWithString:@"https://httpbin.org/redirect/1"]];
+    AFHTTPSessionManager *manager = [AFHTTPSessionManager manager];
+    AFHTTPResponseSerializer *serializer = [AFHTTPResponseSerializer serializer];
+    serializer.acceptableStatusCodes = [NSIndexSet indexSetWithIndex:302];
+    [manager setResponseSerializer:serializer];
+
+    __block __strong NSURL* url = nil;
+    [manager setTaskWillPerformHTTPRedirectionBlock:^NSURLRequest * (NSURLSession * session, NSURLSessionTask * task, NSURLResponse * response, NSURLRequest * request) {
+        if (response == nil) {
+            return request;
+        }
+        url = request.URL;
+        [redirectExpectation fulfill];
+        return nil;
+    }];
+
+    [manager GET:req.URL.absoluteString parameters:nil progress:nil success:^(NSURLSessionTask *task, id responseObject) {
+        // Expect the 302 response when the redirection block returns nil (don't follow redirects)
+        [expectation fulfill];
+    } failure:^(NSURLSessionTask *operation, NSError *error) {
+        XCTFail(@"Unexpected network failure");
+        [expectation fulfill];
+    }];
+
+    // Allow a longer timeout as this test actually hits the network
+    [self waitForExpectationsWithTimeout:10 handler:nil];
+
+    XCTAssertEqualObjects(url, redirectURL, @"Unexpected data received");
+}
+#endif
+
+
 @end
 
 
diff --git a/OHHTTPStubs/UnitTests/Test Suites/NSURLSessionTests.m b/OHHTTPStubs/UnitTests/Test Suites/NSURLSessionTests.m
index a4cb91f4..c3646e23 100644
--- a/OHHTTPStubs/UnitTests/Test Suites/NSURLSessionTests.m	
+++ b/OHHTTPStubs/UnitTests/Test Suites/NSURLSessionTests.m	
@@ -38,12 +38,13 @@
 @import OHHTTPStubs;
 #endif
 
-@interface NSURLSessionTests : XCTestCase <NSURLSessionDataDelegate> @end
+@interface NSURLSessionTests : XCTestCase <NSURLSessionDataDelegate, NSURLSessionTaskDelegate> @end
 
 @implementation NSURLSessionTests
 {
     NSMutableData* _receivedData;
     XCTestExpectation* _taskDidCompleteExpectation;
+    BOOL _shouldFollowRedirects;
 }
 
 - (void)setUp
@@ -51,6 +52,7 @@ - (void)setUp
     [super setUp];
     [OHHTTPStubs removeAllStubs];
     _receivedData = nil;
+    _shouldFollowRedirects = YES;
 }
 
 - (void)_test_NSURLSession:(NSURLSession*)session
@@ -100,6 +102,67 @@ - (void)_test_NSURLSession:(NSURLSession*)session
     }
 }
 
+- (void)_test_redirect_NSURLSession:(NSURLSession*)session
+                        jsonForStub:(id)json
+                         completion:(void(^)(NSError* errorResponse, NSHTTPURLResponse *response, id jsonResponse))completion
+{
+    if ([NSURLSession class])
+    {
+        static const NSTimeInterval kRequestTime = 0.0;
+        static const NSTimeInterval kResponseTime = 0.2;
+
+        [OHHTTPStubs stubRequestsPassingTest:^BOOL(NSURLRequest *request) {
+            return [[[request URL] path] isEqualToString:@""];
+        } withStubResponse:^OHHTTPStubsResponse *(NSURLRequest *request) {
+            NSDictionary *headers = @{ @"Location": @"foo://unknownhost:666/elsewhere" };
+            return [[OHHTTPStubsResponse responseWithData:[[NSData alloc] init] statusCode:301 headers:headers]
+                    requestTime:kRequestTime responseTime:kResponseTime];
+        }];
+
+        [OHHTTPStubs stubRequestsPassingTest:^BOOL(NSURLRequest *request) {
+            return [[[request URL] path] isEqualToString:@"/elsewhere"];
+        } withStubResponse:^OHHTTPStubsResponse *(NSURLRequest *request) {
+            return [[OHHTTPStubsResponse responseWithJSONObject:json statusCode:200 headers:nil]
+                    requestTime:kRequestTime responseTime:kResponseTime];
+        }];
+
+        XCTestExpectation* expectation = [self expectationWithDescription:@"NSURLSessionDataTask completed"];
+
+        __block __strong NSHTTPURLResponse *redirectResponse = nil;
+        __block __strong id dataResponse = nil;
+        __block __strong NSError* errorResponse = nil;
+        NSMutableURLRequest *request = [NSMutableURLRequest requestWithURL:[NSURL URLWithString:@"foo://unknownhost:666"]];
+        request.HTTPMethod = @"GET";
+        [request setValue:@"application/json" forHTTPHeaderField:@"Content-Type"];
+        [request setValue:@"application/json" forHTTPHeaderField:@"Accept"];
+
+        NSURLSessionDataTask *task =  [session dataTaskWithRequest:request
+                                                 completionHandler:^(NSData *data, NSURLResponse *response, NSError *error)
+                                       {
+                                           errorResponse = error;
+                                           if (!error)
+                                           {
+                                               NSHTTPURLResponse *HTTPResponse = (NSHTTPURLResponse *)response;
+                                               if ([HTTPResponse statusCode] >= 300 && [HTTPResponse statusCode] < 400) {
+                                                   redirectResponse = HTTPResponse;
+                                               } else {
+                                                   NSError *jsonError = nil;
+                                                   NSDictionary *jsonObject = [NSJSONSerialization JSONObjectWithData:data options:0 error:&jsonError];
+                                                   XCTAssertNil(jsonError, @"Unexpected error deserializing JSON response");
+                                                   dataResponse = jsonObject;
+                                               }
+                                           }
+                                           [expectation fulfill];
+                                       }];
+
+        [task resume];
+
+        [self waitForExpectationsWithTimeout:kRequestTime+kResponseTime+0.5 handler:nil];
+
+        completion(errorResponse, redirectResponse, dataResponse);
+    }
+}
+
 // The shared session use the same mechanism as NSURLConnection
 // (based on protocols registered via +[NSURLProtocol registerClass:] and all)
 // and no NSURLSessionConfiguration
@@ -114,6 +177,12 @@ - (void)test_SharedNSURLSession
             XCTAssertNil(errorResponse, @"Unexpected error");
             XCTAssertEqualObjects(jsonResponse, json, @"Unexpected data received");
         }];
+
+        [self _test_redirect_NSURLSession:session jsonForStub:json completion:^(NSError *errorResponse, NSHTTPURLResponse *redirectResponse, id jsonResponse) {
+            XCTAssertNil(errorResponse, @"Unexpected error");
+            XCTAssertNil(redirectResponse, @"Unexpected redirect response received");
+            XCTAssertEqualObjects(jsonResponse, json, @"Unexpected data received");
+        }];
     }
     else
     {
@@ -133,6 +202,35 @@ - (void)test_NSURLSessionDefaultConfig
             XCTAssertNil(errorResponse, @"Unexpected error");
             XCTAssertEqualObjects(jsonResponse, json, @"Unexpected data received");
         }];
+
+        [self _test_redirect_NSURLSession:session jsonForStub:json completion:^(NSError *errorResponse, NSHTTPURLResponse *redirectResponse, id jsonResponse) {
+            XCTAssertNil(errorResponse, @"Unexpected error");
+            XCTAssertNil(redirectResponse, @"Unexpected redirect response received");
+            XCTAssertEqualObjects(jsonResponse, json, @"Unexpected data received");
+        }];
+    }
+    else
+    {
+        NSLog(@"/!\\ Test skipped because the NSURLSession class is not available on this OS version. Run the tests a target with a more recent OS.\n");
+    }
+}
+
+- (void)test_NSURLSessionDefaultConfig_notFollowingRedirects
+{
+    _shouldFollowRedirects = NO;
+
+    if ([NSURLSessionConfiguration class] && [NSURLSession class])
+    {
+        NSURLSessionConfiguration* config = [NSURLSessionConfiguration defaultSessionConfiguration];
+        NSURLSession *session = [NSURLSession sessionWithConfiguration:config delegate:self delegateQueue:nil];
+
+        NSDictionary* json = @{@"Success": @"Yes"};
+        [self _test_redirect_NSURLSession:session jsonForStub:json completion:^(NSError *errorResponse, NSHTTPURLResponse *redirectResponse, id jsonResponse) {
+            XCTAssertNil(errorResponse, @"Unexpected error");
+            XCTAssertNotNil(redirectResponse, @"Redirect response should have been received");
+            XCTAssertEqual(301, [redirectResponse statusCode], @"Expected 301 redirect");
+            XCTAssertNil(jsonResponse, @"Unexpected data received");
+        }];
     }
     else
     {
@@ -152,6 +250,12 @@ - (void)test_NSURLSessionEphemeralConfig
             XCTAssertNil(errorResponse, @"Unexpected error");
             XCTAssertEqualObjects(jsonResponse, json, @"Unexpected data received");
         }];
+
+        [self _test_redirect_NSURLSession:session jsonForStub:json completion:^(NSError *errorResponse, NSHTTPURLResponse *redirectResponse, id jsonResponse) {
+            XCTAssertNil(errorResponse, @"Unexpected error");
+            XCTAssertNil(redirectResponse, @"Unexpected redirect response received");
+            XCTAssertEqualObjects(jsonResponse, json, @"Unexpected data received");
+        }];
     }
     else
     {
@@ -173,6 +277,13 @@ - (void)test_NSURLSessionDefaultConfig_Disabled
             XCTAssertNotNil(errorResponse, @"Expected error but none found");
             XCTAssertNil(jsonResponse, @"Data should not have been received as stubs should be disabled");
         }];
+
+        [self _test_redirect_NSURLSession:session jsonForStub:json completion:^(NSError *errorResponse, NSHTTPURLResponse *redirectResponse, id jsonResponse) {
+            // Stubs were disable for this session, so we should get an error instead of the stubs data
+            XCTAssertNotNil(errorResponse, @"Expected error but none found");
+            XCTAssertNil(redirectResponse, @"Redirect response should not have been received as stubs should be disabled");
+            XCTAssertNil(jsonResponse, @"Data should not have been received as stubs should be disabled");
+        }];
     }
     else
     {
@@ -226,6 +337,11 @@ - (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didComp
     [_taskDidCompleteExpectation fulfill];
 }
 
+- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task willPerformHTTPRedirection:(NSHTTPURLResponse *)response newRequest:(NSURLRequest *)request completionHandler:(void (^)(NSURLRequest * _Nullable))completionHandler
+{
+    completionHandler(_shouldFollowRedirects ? request : nil);
+}
+
 @end
 
 #else