Added redirect support.

OHHTTPStubs/OHHTTPStubs.m

@@ -260,17 +260,38 @@ - (void)startLoading
             [[NSHTTPCookieStorage sharedHTTPCookieStorage] setCookies:cookies forURL:request.URL mainDocumentURL:request.mainDocumentURL];
         }
 
-        execute_after(requestTime,^{
-            [client URLProtocol:self didReceiveResponse:urlResponse cacheStoragePolicy:NSURLCacheStorageNotAllowed];
-
-            execute_after(responseTime,^{
-                [client URLProtocol:self didLoadData:responseStub.responseData];
-                [client URLProtocolDidFinishLoading:self];
+
+        NSString* redirectLocation = [responseStub.httpHeaders objectForKey:@"Location"];
+        NSURL* redirectLocationURL;
+        if (redirectLocation)
+        {
+            redirectLocationURL = [NSURL URLWithString:redirectLocation];
+        }
+        else
+        {
+            redirectLocationURL = nil;
+        }
+        if (((responseStub.statusCode/100)==3) && redirectLocationURL)
+        {
+            NSURLRequest* redirectRequest = [NSURLRequest requestWithURL:redirectLocationURL];
+            execute_after(responseTime, ^{
+                [client URLProtocol:self wasRedirectedToRequest:redirectRequest redirectResponse:urlResponse];
+            });
+        }
+        else
+        {
+            execute_after(requestTime,^{
+                [client URLProtocol:self didReceiveResponse:urlResponse cacheStoragePolicy:NSURLCacheStorageNotAllowed];
+
+                execute_after(responseTime,^{
+                    [client URLProtocol:self didLoadData:responseStub.responseData];
+                    [client URLProtocolDidFinishLoading:self];
+                });
             });
-        });
 #if ! __has_feature(objc_arc)
-        [urlResponse autorelease];
+            [urlResponse autorelease];
 #endif
+        }
     } else {
         // Send the canned error
         execute_after(responseStub.responseTime, ^{

OHHTTPStubs/UnitTests/Test Suites/NSURLConnectionDelegateTests.m

@@ -34,6 +34,9 @@ @implementation NSURLConnectionDelegateTests
 {
     NSMutableData* _data;
     NSError* _error;
+
+    NSURL* _redirectRequestURL;
+    NSInteger _redirectResponseStatusCode;
 }
 
 ///////////////////////////////////////////////////////////////////////////////////
@@ -58,6 +61,27 @@ -(void)tearDown
     [super tearDown];
 }
 
+- (NSURLRequest *)connection:(NSURLConnection *)connection willSendRequest:(NSURLRequest *)request redirectResponse:(NSURLResponse *)response
+{
+    _redirectRequestURL = [request URL];
+    if (response)
+    {
+        if ([response isKindOfClass:[NSHTTPURLResponse class]])
+        {
+            _redirectResponseStatusCode = [((NSHTTPURLResponse *) response) statusCode];
+        }
+        else
+        {
+            _redirectResponseStatusCode = 0;
+        }
+    }
+    else
+    {
+        // we get a nil response when NSURLConnection canonicalizes the URL, we don't care about that.
+    }
+    return request;
+}
+
 -(void)connection:(NSURLConnection *)connection didReceiveResponse:(NSURLResponse *)response
 {
     [_data setLength:0U];
@@ -225,4 +249,99 @@ -(void)test_NSURLConnection_cookies
 
 }
 
+
+///////////////////////////////////////////////////////////////////////////////////
+#pragma mark Redirected requests
+///////////////////////////////////////////////////////////////////////////////////
+
+- (void)test_NSURLConnection_redirected
+{
+    static const NSTimeInterval kResponseTime = 1.0;
+    NSData* redirectData = [[NSString stringWithFormat:@"%@ - redirect", NSStringFromSelector(_cmd)] dataUsingEncoding:NSUTF8StringEncoding];
+    NSData* testData = [NSStringFromSelector(_cmd) dataUsingEncoding:NSUTF8StringEncoding];
+    NSURL* redirectURL = [NSURL URLWithString:@"http://www.yahoo.com/"];
+    NSString* redirectCookieName = @"yahooCookie";
+    NSString* redirectCookieValue = [[NSProcessInfo processInfo] globallyUniqueString];
+
+    // Set the cookie accept policy to accept all cookies from the main document domain
+    // (especially in case the previous policy was "NSHTTPCookieAcceptPolicyNever")
+    NSHTTPCookieStorage* cookieStorage = [NSHTTPCookieStorage sharedHTTPCookieStorage];
+    NSHTTPCookieAcceptPolicy previousAcceptPolicy = [cookieStorage cookieAcceptPolicy]; // keep it to restore later
+    [cookieStorage setCookieAcceptPolicy:NSHTTPCookieAcceptPolicyOnlyFromMainDocumentDomain];
+
+    NSString* endCookieName = @"googleCookie";
+    NSString* endCookieValue = [[NSProcessInfo processInfo] globallyUniqueString];
+    NSURL *endURL = [NSURL URLWithString:@"http://www.google.com/"];
+
+    [OHHTTPStubs shouldStubRequestsPassingTest:^BOOL(NSURLRequest *request) {
+        return YES;
+    } withStubResponse:^OHHTTPStubsResponse *(NSURLRequest *request) {
+        if ([[request URL] isEqual:redirectURL]) {
+            NSString* redirectCookie = [NSString stringWithFormat:@"%@=%@;", redirectCookieName, redirectCookieValue];
+            NSDictionary* headers = [NSDictionary dictionaryWithObjectsAndKeys:
+                                     [endURL absoluteString], @"Location",
+                                     redirectCookie, @"Set-Cookie",
+                                     nil];
+            return [OHHTTPStubsResponse responseWithData:redirectData
+                                              statusCode:311 // any 300-level request will do
+                                            responseTime:kResponseTime
+                                                 headers:headers];
+        } else {
+            NSString* endCookie = [NSString stringWithFormat:@"%@=%@;", endCookieName, endCookieValue];
+            NSDictionary* headers = [NSDictionary dictionaryWithObject:endCookie forKey:@"Set-Cookie"];
+            return [OHHTTPStubsResponse responseWithData:testData
+                                              statusCode:200
+                                            responseTime:kResponseTime
+                                                 headers:headers];
+        }
+    }];
+
+    NSURLRequest* req = [NSURLRequest requestWithURL:redirectURL];
+    NSDate* startDate = [NSDate date];
+
+    NSURLConnection* cxn = [NSURLConnection connectionWithRequest:req delegate:self];
+
+//    [self waitForAsyncOperationWithTimeout:kResponseTime+kResponseTimeTolerence];
+    [self waitForAsyncOperationWithTimeout:500000];
+
+    STAssertEqualObjects(_redirectRequestURL, endURL, @"Invalid redirect request URL");
+    STAssertEquals(_redirectResponseStatusCode, 311, @"Invalid redirect response status code");
+    STAssertEqualObjects(_data, testData, @"Invalid data response");
+    STAssertNil(_error, @"Received unexpected network error %@", _error);
+    STAssertEqualsWithAccuracy(-[startDate timeIntervalSinceNow], kResponseTime, kResponseTimeTolerence, @"Invalid response time");
+
+    /* Check that the redirect cookie has been properly stored */
+    NSArray* redirectCookies = [cookieStorage cookiesForURL:req.URL];
+    BOOL redirectCookieFound = NO;
+    for (NSHTTPCookie* cookie in redirectCookies)
+    {
+        if ([cookie.name isEqualToString:redirectCookieName])
+        {
+            redirectCookieFound = YES;
+            STAssertEqualObjects(cookie.value, redirectCookieValue, @"The redirect cookie does not have the expected value");
+        }
+    }
+    STAssertTrue(redirectCookieFound, @"The redirect cookie was not stored as expected");
+
+    /* Check that the end cookie has been properly stored */
+    NSArray* endCookies = [cookieStorage cookiesForURL:endURL];
+    BOOL endCookieFound = NO;
+    for (NSHTTPCookie* cookie in endCookies)
+    {
+        if ([cookie.name isEqualToString:endCookieName])
+        {
+            endCookieFound = YES;
+            STAssertEqualObjects(cookie.value, endCookieValue, @"The end cookie does not have the expected value");
+        }
+    }
+    STAssertTrue(endCookieFound, @"The end cookie was not stored as expected");
+
+    // in case we timed out before the end of the request (test failed), cancel the request to avoid further delegate method calls
+    [cxn cancel];
+
+
+    // As a courtesy, restore previous policy before leaving
+    [cookieStorage setCookieAcceptPolicy:previousAcceptPolicy];
+}
+
 @end