servant/servant-http-streams/test/Servant/StreamSpec.hs

111 lines
3.9 KiB
Haskell

{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -freduction-depth=100 #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
module Servant.StreamSpec (spec) where
import Control.Monad.Trans.Except
import qualified Data.ByteString as BS
import Data.Proxy
import Prelude ()
import Prelude.Compat
import Servant.API
((:<|>) ((:<|>)), (:>), JSON, NetstringFraming,
NewlineFraming, NoFraming, OctetStream, SourceIO, StreamBody,
StreamGet)
import Servant.ClientSpec
(Person (..))
import qualified Servant.ClientSpec as CS
import Servant.HttpStreams
import Servant.Server
import Servant.Types.SourceT
import System.Entropy
(getEntropy, getHardwareEntropy)
import Test.Hspec
spec :: Spec
spec = describe "Servant.HttpStreams streaming" $ do
streamSpec
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)
getStreamBody :: SourceT IO BS.ByteString -> ClientM (SourceIO BS.ByteString)
getGetNL :<|> getGetNS :<|> _getGetALot :<|> getStreamBody = client api
alice :: Person
alice = Person "Alice" 42
bob :: Person
bob = Person "Bob" 25
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
| otherwise = Effect $ do
let size = powerOfTwo 10
mbs <- getHardwareEntropy size
bs <- maybe (getEntropy size) pure mbs
return (Yield bs (lots (n - 1)))
powerOfTwo :: Int -> Int
powerOfTwo = (2 ^)
withClient :: ClientM a -> BaseUrl -> (Either ServantError a -> IO r) -> IO r
withClient x burl k = do
withClientEnvIO burl $ \env -> withClientM x env k
testRunSourceIO :: SourceIO a
-> IO (Either String [a])
testRunSourceIO = runExceptT . runSourceT
streamSpec :: Spec
streamSpec = beforeAll (CS.startWaiApp server) $ afterAll CS.endWaiApp $ do
it "works with Servant.API.StreamGet.Newline" $ \(_, baseUrl) -> do
withClient getGetNL baseUrl $ \(Right res) ->
testRunSourceIO res `shouldReturn` Right [alice, bob, alice]
it "works with Servant.API.StreamGet.Netstring" $ \(_, baseUrl) -> do
withClient getGetNS baseUrl $ \(Right res) ->
testRunSourceIO res `shouldReturn` Right [alice, bob, alice]
it "works with Servant.API.StreamBody" $ \(_, baseUrl) -> do
withClient (getStreamBody (source input)) baseUrl $ \(Right res) ->
testRunSourceIO res `shouldReturn` Right output
where
input = ["foo", "", "bar"]
output = ["foo", "bar"]