wai-effectful-1.0.2: test/Main.hs

{-# OPTIONS_GHC -Wno-missing-local-signatures #-}
{-# OPTIONS_GHC -Wno-monomorphism-restriction #-}

module Main where

import Control.Concurrent (forkFinally, newEmptyMVar, putMVar, takeMVar)
import Control.Monad (void)
import Data.ByteString qualified as ByteString
import Data.ByteString.Builder (toLazyByteString, word8)
import Data.ByteString.Lazy qualified as LazyByteString
import Data.Foldable (for_)
import Data.List qualified as List
import Data.Tuple qualified as Tuple
import Data.Word (Word8)
import Effectful
import Effectful.FileSystem (runFileSystem)
import Effectful.FileSystem.IO.ByteString (readFile)
import Effectful.Hspec
import Effectful.Prim.IORef
import Effectful.Wai
import Network.Wai.Internal qualified as Wai
import Prelude hiding (readFile)

main :: IO ()
main = runEff . runPrim . runFileSystem . runHspec . describe "Wai" $ do
    describe "responseToStream" do
        let getBody res = do
                let (_, _, f) = responseToStream res
                f \streamingBody -> do
                    builderRef <- newIORef mempty
                    let add b = atomicModifyIORef builderRef \builder -> (builder <> b, ())
                        flush = pure ()
                    streamingBody add flush
                    LazyByteString.toStrict . toLazyByteString <$> readIORef builderRef
        prop "responseLBS" \bytes -> do
            body <- getBody . responseLBS undefined undefined . LazyByteString.pack $ bytes
            body `shouldBe` ByteString.pack bytes
        prop "responseBuilder" \bytes -> do
            body <- getBody . responseBuilder undefined undefined . foldMap word8 $ bytes
            body `shouldBe` ByteString.pack bytes
        prop "responseStream" \chunks -> do
            body <- getBody $ responseStream undefined undefined \sendChunk _ ->
                for_ chunks $ sendChunk . foldMap word8
            body `shouldBe` ByteString.concat (fmap ByteString.pack chunks)
        it "responseFile total" do
            let fp = "LICENCE"
            body <- getBody $ responseFile undefined undefined fp Nothing
            expected <- readFile fp
            body `shouldBe` expected
        prop "responseFile partial" \offset' count' -> do
            let fp = "LICENCE"
            totalBS <- readFile fp
            let total = ByteString.length totalBS
                offset = abs offset' `mod` total
                count = abs count' `mod` (total - offset)
            body <-
                getBody . responseFile undefined undefined fp . Just $
                    FilePart
                        { filePartOffset = fromIntegral offset
                        , filePartByteCount = fromIntegral count
                        , filePartFileSize = fromIntegral total
                        }
            let expected = ByteString.take count $ ByteString.drop offset totalBS
            body `shouldBe` expected
    describe "liftResponse" do
        it "lifts a raw response that sends from another thread" do
            failure <- liftIO newEmptyMVar
            let raw recv send = do
                    chunk <- recv
                    void . forkFinally (send chunk) $
                        putMVar failure . either (Just . show) (const Nothing)
                fallback = Wai.ResponseBuilder undefined undefined mempty
            withSeqEffToIO \unlift ->
                case unliftResponse unlift . liftResponse $ Wai.ResponseRaw raw fallback of
                    Wai.ResponseRaw run _ -> run (pure "chunk") . const $ pure ()
                    _ -> fail "not a raw response"
            liftIO (takeMVar failure) >>= (`shouldBe` Nothing)
    describe "lazyRequestBody" do
        prop "works" \chunks -> do
            req <- mkRequestFromChunks chunks
            body <- lazyRequestBody req
            body `shouldBe` LazyByteString.fromChunks (ByteString.pack <$> chunks)
        it "is lazy" do
            let req = setRequestBodyChunks (error "requestBody") defaultRequest
            _ <- lazyRequestBody req
            pure ()
    describe "strictRequestBody" do
        prop "works" $ \chunks -> do
            req <- mkRequestFromChunks chunks
            body <- strictRequestBody req
            body `shouldBe` LazyByteString.fromChunks (fmap ByteString.pack chunks)

mkRequestFromChunks :: (IOE :> es, Prim :> es) => [[Word8]] -> Eff es (Request es)
mkRequestFromChunks chunks = do
    ref <- newIORef . fmap ByteString.pack . filter (not . null) $ chunks
    pure . flip setRequestBodyChunks defaultRequest . atomicModifyIORef ref $
        foldMap Tuple.swap . List.uncons