Merge pull request #578 from k-bx/master

Add postBody and setRequestBody functions to POST some data in a test.

yesod-test/Yesod/Test.hs

@@ -42,12 +42,14 @@ module Yesod.Test
     --
     , get
     , post
+    , postBody
     , request
     , addRequestHeader
     , setMethod
     , addPostParam
     , addGetParam
     , addFile
+    , setRequestBody
     , RequestBuilder
     , setUrl
 
@@ -165,7 +167,7 @@ getResponse :: YesodExample site (Maybe SResponse)
 getResponse = fmap yedResponse ST.get
 
 data RequestBuilderData site = RequestBuilderData
-    { rbdPosts :: [RequestPart]
+    { rbdPostData :: RBDPostData
     , rbdResponse :: (Maybe SResponse)
     , rbdMethod :: H.Method
     , rbdSite :: site
@@ -174,9 +176,12 @@ data RequestBuilderData site = RequestBuilderData
     , rbdHeaders :: H.RequestHeaders
     }
 
+data RBDPostData = MultipleItemsPostData [RequestPart]
+                 | BinaryPostData BSL8.ByteString
+
 -- | Request parts let us discern regular key/values from files sent in the request.
 data RequestPart
-  = ReqPlainPart T.Text T.Text
+  = ReqKvPart T.Text T.Text
   | ReqFilePart T.Text FilePath BSL8.ByteString T.Text
 
 -- | The RequestBuilder state monad constructs an url encoded string of arguments
@@ -370,9 +375,10 @@ printMatches query = do
 -- | Add a parameter with the given name and value.
 addPostParam :: T.Text -> T.Text -> RequestBuilder site ()
 addPostParam name value =
-    ST.modify $ \rbd -> rbd
+  ST.modify $ \rbd -> rbd { rbdPostData = (addPostData (rbdPostData rbd)) }
-        { rbdPosts = ReqPlainPart name value : rbdPosts rbd
+  where addPostData (BinaryPostData _) = error "Trying to add post param to binary content."
-        }
+        addPostData (MultipleItemsPostData posts) =
+          MultipleItemsPostData $ ReqKvPart name value : posts
 
 addGetParam :: T.Text -> T.Text -> RequestBuilder site ()
 addGetParam name value = ST.modify $ \rbd -> rbd
@@ -386,9 +392,10 @@ addGetParam name value = ST.modify $ \rbd -> rbd
 addFile :: T.Text -> FilePath -> T.Text -> RequestBuilder site ()
 addFile name path mimetype = do
   contents <- liftIO $ BSL8.readFile path
-  ST.modify $ \rbd -> rbd
+  ST.modify $ \rbd -> rbd { rbdPostData = (addPostData (rbdPostData rbd) contents) }
-    { rbdPosts = ReqFilePart name path contents mimetype : rbdPosts rbd
+    where addPostData (BinaryPostData _) _ = error "Trying to add file after setting binary content."
-    }
+          addPostData (MultipleItemsPostData posts) contents =
+            MultipleItemsPostData $ ReqFilePart name path contents mimetype : posts
 
 -- This looks up the name of a field based on the contents of the label pointing to it.
 nameFromLabel :: T.Text -> RequestBuilder site T.Text
@@ -459,8 +466,18 @@ post :: (Yesod site, RedirectUrl site url)
      => url
      -> YesodExample site ()
 post url = request $ do
-    setMethod "POST"
+  setMethod "POST"
-    setUrl url
+  setUrl url
+
+-- | Perform a POST request to url with sending a body into it.
+postBody :: (Yesod site, RedirectUrl site url)
+         => url
+         -> BSL8.ByteString
+         -> YesodExample site ()
+postBody url body = request $ do
+  setMethod "POST"
+  setUrl url
+  setRequestBody body
 
 -- | Perform a GET request to url, using params
 get :: (Yesod site, RedirectUrl site url)
@@ -494,6 +511,12 @@ setUrl url' = do
         , rbdGets = rbdGets rbd ++ H.parseQuery (TE.encodeUtf8 urlQuery)
         }
 
+-- | Simple way to set HTTP request body
+setRequestBody :: (Yesod site)
+               => BSL8.ByteString
+               -> RequestBuilder site ()
+setRequestBody body = ST.modify $ \rbd -> rbd { rbdPostData = BinaryPostData body }
+
 addRequestHeader :: H.Header -> RequestBuilder site ()
 addRequestHeader header = ST.modify $ \rbd -> rbd
     { rbdHeaders = header : rbdHeaders rbd
@@ -508,7 +531,7 @@ request reqBuilder = do
     YesodExampleData app site oldCookies mRes <- ST.get
 
     RequestBuilderData {..} <- liftIO $ ST.execStateT reqBuilder RequestBuilderData
-      { rbdPosts = []
+      { rbdPostData = MultipleItemsPostData []
       , rbdResponse = mRes
       , rbdMethod = "GET"
       , rbdSite = site
@@ -523,11 +546,21 @@ request reqBuilder = do
     let cookies = M.filter (checkCookieTime currentUtc) oldCookies
         cookiesForPath = M.filter (checkCookiePath path) cookies
 
-    let maker
+    let req = case rbdPostData of
-          | DL.any isFile rbdPosts = makeMultipart
+          MultipleItemsPostData x ->
-          | otherwise = makeSinglepart
+            if DL.any isFile x
-        req = maker cookiesForPath rbdPosts rbdMethod rbdHeaders path rbdGets
+            then (multipart x)
-
+            else singlepart
+          BinaryPostData _ -> singlepart
+          where singlepart = makeSinglepart cookiesForPath rbdPostData rbdMethod rbdHeaders path rbdGets
+                multipart x = makeMultipart cookiesForPath x rbdMethod rbdHeaders path rbdGets
+    -- let maker = case rbdPostData of
+    --       MultipleItemsPostData x ->
+    --         if DL.any isFile x
+    --         then makeMultipart
+    --         else makeSinglepart
+    --       BinaryPostData _ -> makeSinglepart
+    -- let req = maker cookiesForPath rbdPostData rbdMethod rbdHeaders path rbdGets
     response <- liftIO $ runSession (srequest req) app
     let newCookies = map (Cookie.parseSetCookie . snd) $ DL.filter (("Set-Cookie"==) . fst) $ simpleHeaders response
         cookies' = M.fromList [(Cookie.setCookieName c, c) | c <- newCookies] `M.union` cookies
@@ -548,15 +581,28 @@ request reqBuilder = do
     boundary :: String
     boundary = "*******noneedtomakethisrandom"
     separator = BS8.concat ["--", BS8.pack boundary, "\r\n"]
+    makeMultipart :: M.Map a0 Cookie.SetCookie
+                  -> [RequestPart]
+                  -> H.Method
+                  -> [H.Header]
+                  -> T.Text
+                  -> H.Query
+                  -> SRequest
     makeMultipart cookies parts method extraHeaders urlPath urlQuery =
-      flip SRequest (BSL8.fromChunks [multiPartBody parts]) $ mkRequest
+      SRequest simpleRequest' (simpleRequestBody' parts)
-        [ ("Cookie", Builder.toByteString $ Cookie.renderCookies
+      where simpleRequestBody' x =
-              [(Cookie.setCookieName c, Cookie.setCookieValue c) | c <- map snd $ M.toList cookies])
+              BSL8.fromChunks [multiPartBody x]
-        , ("Content-Type", BS8.pack $ "multipart/form-data; boundary=" ++ boundary)
+            simpleRequest' = mkRequest
-        ] method extraHeaders urlPath urlQuery
+                             [ ("Cookie", cookieValue)
+                             , ("Content-Type", contentTypeValue)]
+                             method extraHeaders urlPath urlQuery
+            cookieValue = Builder.toByteString $ Cookie.renderCookies cookiePairs
+            cookiePairs = [ (Cookie.setCookieName c, Cookie.setCookieValue c)
+                          | c <- map snd $ M.toList cookies ]
+            contentTypeValue = BS8.pack $ "multipart/form-data; boundary=" ++ boundary
     multiPartBody parts =
       BS8.concat $ separator : [BS8.concat [multipartPart p, separator] | p <- parts]
-    multipartPart (ReqPlainPart k v) = BS8.concat
+    multipartPart (ReqKvPart k v) = BS8.concat
       [ "Content-Disposition: form-data; "
       , "name=\"", TE.encodeUtf8 k, "\"\r\n\r\n"
       , TE.encodeUtf8 v, "\r\n"]
@@ -568,15 +614,29 @@ request reqBuilder = do
       , BS8.concat $ BSL8.toChunks bytes, "\r\n"]
 
     -- For building the regular non-multipart requests
-    makeSinglepart cookies parts method extraHeaders urlPath urlQuery = SRequest (mkRequest
+    makeSinglepart :: M.Map a0 Cookie.SetCookie
-      [ ("Cookie", Builder.toByteString $ Cookie.renderCookies
+                   -> RBDPostData
-              [(Cookie.setCookieName c, Cookie.setCookieValue c) | c <- map snd $ M.toList cookies])
+                   -> H.Method
-      , ("Content-Type", "application/x-www-form-urlencoded")
+                   -> [H.Header]
-      ] method extraHeaders urlPath urlQuery) $
+                   -> T.Text
-      BSL8.fromChunks $ return $ TE.encodeUtf8 $ T.intercalate "&" $ map singlepartPart parts
+                   -> H.Query
-
+                   -> SRequest
-    singlepartPart (ReqFilePart _ _ _ _) = ""
+    makeSinglepart cookies rbdPostData method extraHeaders urlPath urlQuery =
-    singlepartPart (ReqPlainPart k v) = T.concat [k,"=",v]
+      SRequest simpleRequest' (simpleRequestBody' rbdPostData)
+      where
+        simpleRequest' = (mkRequest
+                          [ ("Cookie", cookieValue)
+                          , ("Content-Type", "application/x-www-form-urlencoded")]
+                          method extraHeaders urlPath urlQuery)
+        simpleRequestBody' (MultipleItemsPostData x) =
+          BSL8.fromChunks $ return $ TE.encodeUtf8 $ T.intercalate "&"
+          $ map singlepartPart x
+        simpleRequestBody' (BinaryPostData x) = x
+        cookieValue = Builder.toByteString $ Cookie.renderCookies cookiePairs
+        cookiePairs = [ (Cookie.setCookieName c, Cookie.setCookieValue c)
+                      | c <- map snd $ M.toList cookies ]
+        singlepartPart (ReqFilePart _ _ _ _) = ""
+        singlepartPart (ReqKvPart k v) = T.concat [k,"=",v]
 
     -- General request making
     mkRequest headers method extraHeaders urlPath urlQuery = defaultRequest