From df6834a335b124d5aa014b35bd3b5ec00dadd9e2 Mon Sep 17 00:00:00 2001
From: Mark Wotton <mwotton@gmail.com>
Date: Mon, 21 Mar 2016 12:13:23 -0400
Subject: [PATCH] add followRedirect

---
 yesod-test/Yesod/Test.hs    | 20 ++++++++++++++++++++
 yesod-test/test/main.hs     | 29 +++++++++++++++++++++++++++--
 yesod-test/yesod-test.cabal |  3 ++-
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/yesod-test/Yesod/Test.hs b/yesod-test/Yesod/Test.hs
index 33d60364..32b6b0b1 100644
--- a/yesod-test/Yesod/Test.hs
+++ b/yesod-test/Yesod/Test.hs
@@ -50,6 +50,7 @@ module Yesod.Test
     , get
     , post
     , postBody
+    , followRedirect
     , request
     , addRequestHeader
     , setMethod
@@ -693,6 +694,25 @@ get url = request $ do
     setMethod "GET"
     setUrl url
 
+-- | Follow a redirect, if the last response was a redirect.
+-- ==== __Examples__
+-- > get HomeR
+--
+-- > followRedirect
+followRedirect :: Yesod site
+               =>  YesodExample site ()
+followRedirect = do
+  mr <- getResponse
+  case mr of
+   Nothing ->  failure "no response, so no redirect to follow"
+   Just r -> do
+     if not ((H.statusCode $ simpleStatus r) `elem` [301,303])
+       then failure "followRedirect called, but previous request was not a redirect"
+       else do
+         case lookup "Location" (simpleHeaders r) of
+          Nothing -> failure "No location header set"
+          Just h -> get (TE.decodeUtf8 h)
+
 -- | Sets the HTTP method used by the request.
 --
 -- ==== __Examples__
diff --git a/yesod-test/test/main.hs b/yesod-test/test/main.hs
index 5461e8bd..f7808bc0 100644
--- a/yesod-test/test/main.hs
+++ b/yesod-test/test/main.hs
@@ -20,11 +20,13 @@ import Data.Monoid ((<>))
 import Control.Applicative
 import Network.Wai (pathInfo, requestHeaders)
 import Data.Maybe (fromMaybe)
+import Data.Either (isLeft)
+import Control.Monad.Catch (try)
 
 import Data.ByteString.Lazy.Char8 ()
 import qualified Data.Map as Map
 import qualified Text.HTML.DOM as HD
-import Network.HTTP.Types.Status (unsupportedMediaType415)
+import Network.HTTP.Types.Status (status301, status303, unsupportedMediaType415)
 
 parseQuery_ = either error id . parseQuery
 findBySelector_ x = either error id . findBySelector x
@@ -213,8 +215,28 @@ main = hspec $ do
                 setMethod "POST"
                 setUrl ("/" :: Text)
             statusIs 403
+    describe "test redirects" $ yesodSpec app $ do
+        yit "follows 303 redirects when requested" $ do
+            get ("/redirect303" :: Text)
+            statusIs 303
+            followRedirect
+            statusIs 200
+            bodyContains "we have been successfully redirected"
+
+        yit "follows 301 redirects when requested" $ do
+            get ("/redirect301" :: Text)
+            statusIs 301
+            followRedirect
+            statusIs 200
+            bodyContains "we have been successfully redirected"
 
 
+        yit "throws an exception when no redirect was returned" $ do
+            get ("/" :: Text)
+            statusIs 200
+            r <- followRedirect
+            statusIs 200
+            -- assertBool "expected exception" $ isLeft r
 
 instance RenderMessage LiteApp FormMessage where
     renderMessage _ _ = defaultFormMessage
@@ -235,6 +257,9 @@ app = liteApp $ do
         case mfoo of
             Nothing -> error "No foo"
             Just foo -> return foo
+    onStatic "redirect301" $ dispatchTo $ redirectWith status301 ("/redirectTarget" :: Text) >> return ()
+    onStatic "redirect303" $ dispatchTo $ redirectWith status303 ("/redirectTarget" :: Text) >> return ()
+    onStatic "redirectTarget" $ dispatchTo $ return ("we have been successfully redirected" :: Text)
     onStatic "form" $ dispatchTo $ do
         ((mfoo, widget), _) <- runFormPost
                         $ renderDivs
@@ -290,4 +315,4 @@ postHomeR = defaultLayout
     [whamlet|
         <p>
             Welcome to my test application.
-    |]
\ No newline at end of file
+    |]
diff --git a/yesod-test/yesod-test.cabal b/yesod-test/yesod-test.cabal
index 907624e5..f0c607f6 100644
--- a/yesod-test/yesod-test.cabal
+++ b/yesod-test/yesod-test.cabal
@@ -4,7 +4,7 @@ license:            MIT
 license-file:       LICENSE
 author:             Nubis <nubis@woobiz.com.ar>
 maintainer:         Michael Snoyman, Greg Weber, Nubis <nubis@woobiz.com.ar>
-synopsis:           integration testing for WAI/Yesod Applications 
+synopsis:           integration testing for WAI/Yesod Applications
 category:           Web, Yesod, Testing
 stability:          Experimental
 cabal-version:      >= 1.8
@@ -60,6 +60,7 @@ test-suite test
                           , yesod-form
                           , text
                           , wai
+                          , exceptions
                           , http-types
 
 source-repository head