From e49b0369c04b2284408a7dcfb7a19bbe359452bc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Domen=20Ko=C5=BEar?= <domen@dev.si>
Date: Thu, 31 Jan 2019 16:18:22 +0700
Subject: [PATCH] servant-client: add a test case for StreamBody

---
 servant-client/test/Servant/StreamSpec.hs | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/servant-client/test/Servant/StreamSpec.hs b/servant-client/test/Servant/StreamSpec.hs
index e41aa370..3c84fb38 100644
--- a/servant-client/test/Servant/StreamSpec.hs
+++ b/servant-client/test/Servant/StreamSpec.hs
@@ -35,8 +35,9 @@ import qualified Network.HTTP.Client           as C
 import           Prelude ()
 import           Prelude.Compat
 import           Servant.API
-                 ((:<|>) ((:<|>)), (:>), JSON, NetstringFraming,
-                 NewlineFraming, NoFraming, OctetStream, SourceIO, StreamGet)
+                 ((:<|>) ((:<|>)), (:>), JSON, NetstringFraming, StreamBody,
+                 NewlineFraming, NoFraming, OctetStream, SourceIO, StreamGet,
+                 )
 import           Servant.Client.Streaming
 import           Servant.ClientSpec
                  (Person (..))
@@ -72,13 +73,15 @@ type StreamApi =
          "streamGetNewline" :> StreamGet NewlineFraming JSON (SourceIO Person)
     :<|> "streamGetNetstring" :> StreamGet NetstringFraming JSON (SourceIO Person)
     :<|> "streamALot" :> StreamGet NoFraming OctetStream (SourceIO BS.ByteString)
+    :<|> "streamBody" :> StreamBody NoFraming OctetStream (SourceIO BS.ByteString) :> StreamGet NoFraming OctetStream (SourceIO BS.ByteString)
 
 api :: Proxy StreamApi
 api = Proxy
 
 getGetNL, getGetNS :: ClientM (SourceIO Person)
 getGetALot :: ClientM (SourceIO BS.ByteString)
-getGetNL :<|> getGetNS :<|> getGetALot = client api
+getStreamBody :: SourceT IO BS.ByteString -> ClientM (SourceIO BS.ByteString)
+getGetNL :<|> getGetNS :<|> getGetALot :<|> getStreamBody = client api
 
 alice :: Person
 alice = Person "Alice" 42
@@ -90,9 +93,9 @@ server :: Application
 server = serve api
     $    return (source [alice, bob, alice])
     :<|> return (source [alice, bob, alice])
-
     -- 2 ^ (18 + 10) = 256M
     :<|> return (SourceT ($ lots (powerOfTwo 18)))
+    :<|> return
   where
     lots n
         | n < 0     = Stop
@@ -126,6 +129,12 @@ streamSpec = beforeAll (CS.startWaiApp server) $ afterAll CS.endWaiApp $ do
         withClient getGetNS baseUrl $ \(Right res) ->
             testRunSourceIO res `shouldReturn` Right [alice, bob, alice]
 
+    it "works with Servant.API.StreamBody" $ \(_, baseUrl) -> do
+        withClient (getStreamBody (source inout)) baseUrl $ \(Right res) ->
+            testRunSourceIO res `shouldReturn` Right inout
+        where
+          inout = ["foo", "bar"]
+
 {-
     it "streams in constant memory" $ \(_, baseUrl) -> do
         Right rs <- runClient getGetALot baseUrl