From 7b12f61a914f2c81709c7d232f6b433c376f6c99 Mon Sep 17 00:00:00 2001
From: Jason Whittle <jason.whittle@gmail.com>
Date: Wed, 30 Nov 2016 18:05:48 -0500
Subject: [PATCH] yesod-test: add getLocation test helper.

---
 yesod-test/Yesod/Test.hs | 23 +++++++++++++++++
 yesod-test/test/main.hs  | 55 ++++++++++++++++++++++++++++++++++------
 2 files changed, 70 insertions(+), 8 deletions(-)

diff --git a/yesod-test/Yesod/Test.hs b/yesod-test/Yesod/Test.hs
index b3cf4a294..f95a48e9a 100644
--- a/yesod-test/Yesod/Test.hs
+++ b/yesod-test/Yesod/Test.hs
@@ -52,6 +52,7 @@ module Yesod.Test
     , post
     , postBody
     , followRedirect
+    , getLocation
     , request
     , addRequestHeader
     , setMethod
@@ -749,6 +750,28 @@ followRedirect = do
           Just h -> let url = TE.decodeUtf8 h in
                      get url  >> return (Right url)
 
+-- | Parse the Location header of the last response.
+--
+-- ==== __Examples__
+--
+-- > post ResourcesR
+-- > (Right (ResourceR resourceId)) <- getLocation
+getLocation :: (Yesod site, ParseRoute site)
+            => YesodExample site (Either T.Text (Route site))
+getLocation = do
+  mr <- getResponse
+  case mr of
+    Nothing -> return $ Left "getLocation called, but there was no previous response, so no Location header"
+    Just r -> case lookup "Location" (simpleHeaders r) of
+      Nothing -> return $ Left "getLocation called, but the previous response has no Location header"
+      Just h -> case parseRoute $ decodePath h of
+        Nothing -> return $ Left "getLocation called, but couldn’t parse it into a route"
+        Just l -> return $ Right l
+  where decodePath b = let (x, y) = BS8.break (=='?') b
+                       in (H.decodePathSegments x, unJust <$> H.parseQueryText y)
+        unJust (a, Just b) = (a, b)
+        unJust (a, Nothing) = (a, mempty)
+
 -- | Sets the HTTP method used by the request.
 --
 -- ==== __Examples__
diff --git a/yesod-test/test/main.hs b/yesod-test/test/main.hs
index d97d53398..5b4c1a4e5 100644
--- a/yesod-test/test/main.hs
+++ b/yesod-test/test/main.hs
@@ -5,6 +5,7 @@
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE ViewPatterns #-}
 
 import Test.HUnit hiding (Test)
 import Test.Hspec
@@ -32,6 +33,14 @@ parseQuery_ = either error id . parseQuery
 findBySelector_ x = either error id . findBySelector x
 parseHtml_ = HD.parseLBS
 
+data RoutedApp = RoutedApp
+
+mkYesod "RoutedApp" [parseRoutes|
+/                HomeR      GET POST
+/resources       ResourcesR POST
+/resources/#Text ResourceR  GET
+|]
+
 main :: IO ()
 main = hspec $ do
     describe "CSS selector parsing" $ do
@@ -209,7 +218,7 @@ main = hspec $ do
             statusIs 200
             printBody
             bodyContains "Foo"
-    describe "CSRF with cookies/headers" $ yesodSpec CsrfApp $ do
+    describe "CSRF with cookies/headers" $ yesodSpec RoutedApp $ do
         yit "Should receive a CSRF cookie and add its value to the headers" $ do
             get ("/" :: Text)
             statusIs 200
@@ -251,6 +260,30 @@ main = hspec $ do
             r <- followRedirect
             liftIO $ assertBool "expected a Left when not a redirect" $ isLeft r
 
+    describe "route parsing in tests" $ yesodSpec RoutedApp $ do
+        yit "parses location header into a route" $ do
+            -- get CSRF token
+            get HomeR
+            statusIs 200
+
+            request $ do
+                setMethod "POST"
+                setUrl $ ResourcesR
+                addPostParam "foo" "bar"
+                addTokenFromCookie
+            statusIs 201
+
+            loc <- getLocation
+            liftIO $ assertBool "expected location to be available" $ isRight loc
+            let (Right (ResourceR t)) = loc
+            liftIO $ assertBool "expected location header to contain post param" $ t == "bar"
+
+        yit "returns a Left when no redirect was returned" $ do
+            get HomeR
+            statusIs 200
+            loc <- getLocation
+            liftIO $ assertBool "expected a Left when not a redirect" $ isLeft loc
+
 instance RenderMessage LiteApp FormMessage where
     renderMessage _ _ = defaultFormMessage
 
@@ -307,13 +340,7 @@ cookieApp = liteApp $ do
             redirect ("/cookie/home" :: Text)
             return ()
 
-data CsrfApp = CsrfApp
-
-mkYesod "CsrfApp" [parseRoutes|
-/ HomeR GET POST
-|]
-
-instance Yesod CsrfApp where
+instance Yesod RoutedApp where
     yesodMiddleware = defaultYesodMiddleware . defaultCsrfMiddleware
 
 getHomeR :: Handler Html
@@ -329,3 +356,15 @@ postHomeR = defaultLayout
         <p>
             Welcome to my test application.
     |]
+
+postResourcesR :: Handler ()
+postResourcesR = do
+  ([("foo", t)], _) <- runRequestBody
+  sendResponseCreated $ ResourceR t
+
+getResourceR :: Text -> Handler Html
+getResourceR i = defaultLayout
+    [whamlet|
+        <p>
+            Read item #{i}.
+    |]