From 50b50fa0e6f7798d6b766e59edda64b12d6179e8 Mon Sep 17 00:00:00 2001
From: Viacheslav Lotsmanov <lotsmanov89@gmail.com>
Date: Sun, 4 Oct 2020 05:12:32 +0300
Subject: [PATCH] Write tests for Optional ReqBody' and fix some cases

---
 servant-server/src/Servant/Server/Internal.hs | 22 ++++++------
 servant-server/test/Servant/ServerSpec.hs     | 35 ++++++++++++++++---
 2 files changed, 41 insertions(+), 16 deletions(-)

diff --git a/servant-server/src/Servant/Server/Internal.hs b/servant-server/src/Servant/Server/Internal.hs
index f23a794d1..eaeb3f1f8 100644
--- a/servant-server/src/Servant/Server/Internal.hs
+++ b/servant-server/src/Servant/Server/Internal.hs
@@ -48,8 +48,6 @@ import qualified Data.ByteString.Char8                      as BC8
 import qualified Data.ByteString.Lazy                       as BL
 import           Data.Either
                  (partitionEithers)
-import           Data.Function
-                 ((&))
 import           Data.Maybe
                  (fromMaybe, isNothing, mapMaybe, maybeToList)
 import           Data.Semigroup
@@ -674,9 +672,7 @@ instance ( AllCTUnrender list a, HasServer api context
             Just f  -> return f
 
       -- Body check, we get a body parsing functions as the first argument.
-      bodyCheck f = withRequest $ \ request -> do
-        mrqbody <- f <$> liftIO (lazyRequestBody request)
-
+      bodyCheck f = withRequest $ \ request ->
         let
           hasReqBody =
             case requestBodyLength request of
@@ -685,13 +681,15 @@ instance ( AllCTUnrender list a, HasServer api context
 
           serverErr :: String -> ServerError
           serverErr = formatError rep request . cs
-
-        mrqbody & case (sbool :: SBool (FoldRequired mods), sbool :: SBool (FoldLenient mods), hasReqBody) of
-          (STrue,  STrue,  _)     -> return . bimap cs id
-          (STrue,  SFalse, _)     -> either (delayedFailFatal . serverErr) return
-          (SFalse, _,      False) -> return . const Nothing
-          (SFalse, STrue,  True)  -> return . Just . bimap cs id
-          (SFalse, SFalse, True)  -> either (delayedFailFatal . serverErr) (return . Just)
+        in
+          fmap f (liftIO $ lazyRequestBody request) >>=
+            case (sbool :: SBool (FoldRequired mods), sbool :: SBool (FoldLenient mods), hasReqBody) of
+              (STrue,  STrue,  _)     -> return . bimap cs id
+              (STrue,  SFalse, _)     -> either (delayedFailFatal . serverErr) return
+              (SFalse, STrue,  False) -> return . either (const Nothing) (Just . Right)
+              (SFalse, SFalse, False) -> return . either (const Nothing) Just
+              (SFalse, STrue,  True)  -> return . Just . bimap cs id
+              (SFalse, SFalse, True)  -> either (delayedFailFatal . serverErr) (return . Just)
 
 instance
     ( FramingUnrender framing, FromSourceIO chunk a, MimeUnrender ctype chunk
diff --git a/servant-server/test/Servant/ServerSpec.hs b/servant-server/test/Servant/ServerSpec.hs
index 90e72667e..49c1b8268 100644
--- a/servant-server/test/Servant/ServerSpec.hs
+++ b/servant-server/test/Servant/ServerSpec.hs
@@ -51,9 +51,10 @@ import           Servant.API
                  BasicAuthData (BasicAuthData), Capture, Capture', CaptureAll,
                  Delete, EmptyAPI, Get, Header, Headers, HttpVersion,
                  IsSecure (..), JSON, Lenient, NoContent (..), NoContentVerb,
-                 NoFraming, OctetStream, Patch, PlainText, Post, Put,
-                 QueryFlag, QueryParam, QueryParams, Raw, RemoteHost, ReqBody,
-                 SourceIO, StdMethod (..), Stream, Strict, Verb, addHeader)
+                 NoFraming, OctetStream, Optional, Patch, PlainText, Post, Put,
+                 QueryFlag, QueryParam, QueryParams, Raw, RemoteHost,
+                 ReqBody, ReqBody', SourceIO, StdMethod (..), Stream, Strict,
+                 Verb, addHeader)
 import           Servant.Server
                  (Context ((:.), EmptyContext), Handler, Server, Tagged (..),
                  emptyServer, err401, err403, err404, serve, serveWithContext)
@@ -465,6 +466,7 @@ queryParamSpec = do
 ------------------------------------------------------------------------------
 type ReqBodyApi = ReqBody '[JSON] Person :> Post '[JSON] Person
            :<|> "blah" :> ReqBody '[JSON] Person :> Put '[JSON] Integer
+           :<|> "meh" :> ReqBody' '[Optional, Strict] '[JSON] Person :> Put '[JSON] Integer
 
 reqBodyApi :: Proxy ReqBodyApi
 reqBodyApi = Proxy
@@ -473,7 +475,7 @@ reqBodySpec :: Spec
 reqBodySpec = describe "Servant.API.ReqBody" $ do
 
   let server :: Server ReqBodyApi
-      server = return :<|> return . age
+      server = return :<|> return . age :<|> return . maybe 0 age
       mkReq method x = THW.request method x
          [(hContentType, "application/json;charset=utf-8")]
 
@@ -490,6 +492,31 @@ reqBodySpec = describe "Servant.API.ReqBody" $ do
       THW.request methodPost "/"
         [(hContentType, "application/nonsense")] "" `shouldRespondWith` 415
 
+    describe "optional request body" $ do
+      it "request without body succeeds" $ do
+        THW.request methodPut "/meh" [] mempty `shouldRespondWith` 200
+
+      it "request without body responds with proper default value" $ do
+        response <- THW.request methodPut "/meh" [] mempty
+        liftIO $ simpleBody response `shouldBe` encode (0 :: Integer)
+
+      it "responds with 415 if the request body media type is unsupported" $ do
+        THW.request methodPut "/meh" [(hContentType, "application/nonsense")]
+          (encode alice) `shouldRespondWith` 415
+        THW.request methodPut "/meh" [(hContentType, "application/octet-stream")]
+          (encode alice) `shouldRespondWith` 415
+
+      it "request without body and with content-type header succeeds" $ do
+        mkReq methodPut "/meh" mempty `shouldRespondWith` 200
+
+      it "request without body and with content-type header returns default value" $ do
+        response <- mkReq methodPut "/meh" mempty
+        liftIO $ simpleBody response `shouldBe` encode (0 :: Integer)
+
+      it "optional request body can be provided" $ do
+        response <- mkReq methodPut "/meh" (encode alice)
+        liftIO $ simpleBody response `shouldBe` encode (age alice)
+
 -- }}}
 ------------------------------------------------------------------------------
 -- * headerSpec {{{