From e9466b775241e0cc74554100c9886f6ef100d7ea Mon Sep 17 00:00:00 2001
From: Oleg Grenrus <oleg.grenrus@iki.fi>
Date: Sat, 2 Feb 2019 12:44:35 +0200
Subject: [PATCH] Implement HasClient (StreamBody ... :> api)

---
 .../Servant/Client/Core/Internal/HasClient.hs | 47 ++++++++++++++++---
 servant-client/test/Servant/StreamSpec.hs     |  7 +--
 2 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/servant-client-core/src/Servant/Client/Core/Internal/HasClient.hs b/servant-client-core/src/Servant/Client/Core/Internal/HasClient.hs
index 71200700..6be92ec6 100644
--- a/servant-client-core/src/Servant/Client/Core/Internal/HasClient.hs
+++ b/servant-client-core/src/Servant/Client/Core/Internal/HasClient.hs
@@ -16,6 +16,8 @@ module Servant.Client.Core.Internal.HasClient where
 import           Prelude ()
 import           Prelude.Compat
 
+import           Control.Concurrent.MVar
+                 (modifyMVar, newMVar)
 import qualified Data.ByteString                        as BS
 import qualified Data.ByteString.Lazy                   as BL
 import           Data.Foldable
@@ -36,13 +38,14 @@ import qualified Network.HTTP.Types                     as H
 import           Servant.API
                  ((:<|>) ((:<|>)), (:>), AuthProtect, BasicAuth, BasicAuthData,
                  BuildHeadersTo (..), Capture', CaptureAll, Description,
-                 EmptyAPI, FramingUnrender (..), FromSourceIO (..), Header',
-                 Headers (..), HttpVersion, IsSecure, MimeRender (mimeRender),
+                 EmptyAPI, FramingRender (..), FramingUnrender (..),
+                 FromSourceIO (..), Header', Headers (..), HttpVersion,
+                 IsSecure, MimeRender (mimeRender),
                  MimeUnrender (mimeUnrender), NoContent (NoContent), QueryFlag,
                  QueryParam', QueryParams, Raw, ReflectMethod (..), RemoteHost,
                  ReqBody', SBoolI, Stream, StreamBody', Summary, ToHttpApiData,
-                 Vault, Verb, WithNamedContext, contentType, getHeadersHList,
-                 getResponse, toQueryParam, toUrlPiece)
+                 ToSourceIO (..), Vault, Verb, WithNamedContext, contentType,
+                 getHeadersHList, getResponse, toQueryParam, toUrlPiece)
 import           Servant.API.ContentTypes
                  (contentTypes)
 import           Servant.API.Modifiers
@@ -538,7 +541,7 @@ instance (MimeRender ct a, HasClient m api)
     hoistClientMonad pm (Proxy :: Proxy api) f (cl a)
 
 instance
-    ( HasClient m api
+    ( HasClient m api, MimeRender ctype chunk, FramingRender framing, ToSourceIO chunk a
     ) => HasClient m (StreamBody' mods framing ctype a :> api)
   where
 
@@ -547,7 +550,39 @@ instance
     hoistClientMonad pm _ f cl = \a ->
       hoistClientMonad pm (Proxy :: Proxy api) f (cl a)
 
-    clientWithRoute _pm Proxy _req _body = error "HasClient @StreamBody"
+    clientWithRoute pm Proxy req body
+        = clientWithRoute pm (Proxy :: Proxy api)
+        $ setRequestBody (RequestBodyStreamChunked givesPopper) (contentType ctypeP) req
+      where
+        ctypeP   = Proxy :: Proxy ctype
+        framingP = Proxy :: Proxy framing
+
+        sourceIO = framingRender
+            framingP
+            (mimeRender ctypeP :: chunk -> BL.ByteString)
+            (toSourceIO body)
+
+        -- not pretty.
+        givesPopper :: (IO BS.ByteString -> IO ()) -> IO ()
+        givesPopper needsPopper = S.unSourceT sourceIO $ \step0 -> do
+            ref <- newMVar step0
+
+            -- Note sure we need locking, but it's feels safer.
+            let popper :: IO BS.ByteString
+                popper = modifyMVar ref nextBs
+
+            needsPopper popper
+
+        nextBs S.Stop          = return (S.Stop, BS.empty)
+        nextBs (S.Error err)   = fail err
+        nextBs (S.Skip s)      = nextBs s
+        nextBs (S.Effect ms)   = ms >>= nextBs
+        nextBs (S.Yield lbs s) = case BL.toChunks lbs of
+            []     -> nextBs s
+            (x:xs) | BS.null x -> nextBs step'
+                   | otherwise -> return (step', x)
+                where
+                  step' = S.Yield (BL.fromChunks xs) s
 
 -- | Make the querying function append @path@ to the request path.
 instance (KnownSymbol path, HasClient m api) => HasClient m (path :> api) where
diff --git a/servant-client/test/Servant/StreamSpec.hs b/servant-client/test/Servant/StreamSpec.hs
index 3c84fb38..0e9c557a 100644
--- a/servant-client/test/Servant/StreamSpec.hs
+++ b/servant-client/test/Servant/StreamSpec.hs
@@ -130,10 +130,11 @@ streamSpec = beforeAll (CS.startWaiApp server) $ afterAll CS.endWaiApp $ do
             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
+        withClient (getStreamBody (source input)) baseUrl $ \(Right res) ->
+            testRunSourceIO res `shouldReturn` Right output
         where
-          inout = ["foo", "bar"]
+          input = ["foo", "", "bar"]
+          output = ["foo", "bar"]
 
 {-
     it "streams in constant memory" $ \(_, baseUrl) -> do