{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_HADDOCK prune not-home #-}

{- |
Module      : Data.Attoparsec.Framer.Testing
Copyright   : (c) 2022 Tim Emiola
Maintainer  : Tim Emiola <adetokunbo@emio.la>
SPDX-License-Identifier: BSD3

This module provides combinators that simplify unit tests of code that
use @'Framer's@.
-}
module Data.Attoparsec.Framer.Testing (
  -- * testing combinators
  parsesFromFramerOk,
  chunksOfN,
  linkedSrcAndSink,
  linkedSrcAndSink',
) where

import Control.Exception (catch)
import Control.Monad (when)
import qualified Data.Attoparsec.ByteString as A
import Data.Attoparsec.Framer
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.ByteString.Builder (byteStringHex, toLazyByteString)
import qualified Data.ByteString.Lazy.Char8 as C8
import Data.IORef (
  IORef,
  modifyIORef',
  newIORef,
  readIORef,
  writeIORef,
 )
import Data.List (unfoldr)
import Data.Word (Word32)


{- | Creates a 'Framer' and uses 'runFramer to confirm that the expect frames
  are received '
-}
parsesFromFramerOk :: Eq a => (a -> ByteString) -> A.Parser a -> Word32 -> [a] -> IO Bool
parsesFromFramerOk :: forall a.
Eq a =>
(a -> ByteString) -> Parser a -> Word32 -> [a] -> IO Bool
parsesFromFramerOk a -> ByteString
asBytes Parser a
parser Word32
chunkSize' [a]
wanted = do
  IORef (Maybe [ByteString])
chunkStore <- forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing
  IORef [a]
dst <- forall a. a -> IO (IORef a)
newIORef []
  let updateDst :: a -> IO ()
updateDst a
x = forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef [a]
dst ((:) a
x)
      mkChunks :: Int -> [ByteString]
mkChunks Int
n = forall a. Monoid a => [a] -> a
mconcat forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (Int -> ByteString -> [ByteString]
chunksOfN Int
n forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
asBytes) [a]
wanted
      src :: Word32 -> IO ByteString
src = (Int -> [ByteString])
-> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' Int -> [ByteString]
mkChunks IORef (Maybe [ByteString])
chunkStore
      frames :: Framer IO a
frames = forall (m :: * -> *) a. Word32 -> Framer m a -> Framer m a
setChunkSize Word32
chunkSize' forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) frame.
MonadThrow m =>
Parser frame -> (frame -> m ()) -> ByteSource m -> Framer m frame
mkFramer Parser a
parser a -> IO ()
updateDst Word32 -> IO ByteString
src
  forall (m :: * -> *) frame. MonadThrow m => Framer m frame -> m ()
runFramer Framer IO a
frames forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(NoMoreInput
_e :: NoMoreInput) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

  [a]
got <- forall a. IORef a -> IO a
readIORef IORef [a]
dst
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [a]
got forall a. Eq a => a -> a -> Bool
== forall a. [a] -> [a]
reverse [a]
wanted


-- | Split a 'ByteString' into chunks of given size
chunksOfN :: Int -> ByteString -> [ByteString]
chunksOfN :: Int -> ByteString -> [ByteString]
chunksOfN Int
x ByteString
b =
  let go :: ByteString -> Maybe (ByteString, ByteString)
go ByteString
y =
        let taken :: ByteString
taken = Int -> ByteString -> ByteString
BS.take Int
x ByteString
y
         in if ByteString -> Bool
BS.null ByteString
taken then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just (ByteString
taken, Int -> ByteString -> ByteString
BS.drop Int
x ByteString
y)
   in forall b a. (b -> Maybe (a, b)) -> b -> [a]
unfoldr ByteString -> Maybe (ByteString, ByteString)
go ByteString
b


nextFrom' ::
  (Int -> [ByteString]) -> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' :: (Int -> [ByteString])
-> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' Int -> [ByteString]
initChunks IORef (Maybe [ByteString])
chunkStore Word32
chunkSize' = do
  forall a. IORef a -> IO a
readIORef IORef (Maybe [ByteString])
chunkStore forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe [ByteString]
Nothing -> do
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe [ByteString])
chunkStore forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Int -> [ByteString]
initChunks forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
chunkSize'
      (Int -> [ByteString])
-> IORef (Maybe [ByteString]) -> Word32 -> IO ByteString
nextFrom' Int -> [ByteString]
initChunks IORef (Maybe [ByteString])
chunkStore Word32
chunkSize'
    Just [] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
    Just (ByteString
x : [ByteString]
xs) -> do
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe [ByteString])
chunkStore forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just [ByteString]
xs
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
x


{- | A @'ByteSource'@ linked to a byte sink.

Provides a @ByteSource@ and @byte sink@ that emulate a responding endpoint.

The @responses@ are consumed each time the byte sink is invoked.

Whenever the sink is invoked, the head of the provided responses is removed
and starts to be returned in chunks by the @ByteSource@,
-}
linkedSrcAndSink :: [ByteString] -> IO (ByteSource IO, (ByteString -> IO ()))
linkedSrcAndSink :: [ByteString] -> IO (Word32 -> IO ByteString, ByteString -> IO ())
linkedSrcAndSink [ByteString]
responses = do
  IORef (Maybe ByteString)
refSrc <- forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing
  IORef [ByteString]
refSink <- forall a. a -> IO (IORef a)
newIORef [ByteString]
responses
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (IORef (Maybe ByteString) -> Word32 -> IO ByteString
ioRefByteSource IORef (Maybe ByteString)
refSrc, Bool
-> IORef [ByteString]
-> IORef (Maybe ByteString)
-> ByteString
-> IO ()
ioRefByteSink Bool
False IORef [ByteString]
refSink IORef (Maybe ByteString)
refSrc)


-- | Like 'linkedSrcAndSink', but prints the src and sink to output as debug
linkedSrcAndSink' :: [ByteString] -> IO (ByteSource IO, (ByteString -> IO ()))
linkedSrcAndSink' :: [ByteString] -> IO (Word32 -> IO ByteString, ByteString -> IO ())
linkedSrcAndSink' [ByteString]
responses = do
  IORef (Maybe ByteString)
refSrc <- forall a. a -> IO (IORef a)
newIORef forall a. Maybe a
Nothing
  IORef [ByteString]
refSink <- forall a. a -> IO (IORef a)
newIORef [ByteString]
responses
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (IORef (Maybe ByteString) -> Word32 -> IO ByteString
ioRefByteSource IORef (Maybe ByteString)
refSrc, Bool
-> IORef [ByteString]
-> IORef (Maybe ByteString)
-> ByteString
-> IO ()
ioRefByteSink Bool
True IORef [ByteString]
refSink IORef (Maybe ByteString)
refSrc)


ioRefByteSource :: IORef (Maybe ByteString) -> ByteSource IO
ioRefByteSource :: IORef (Maybe ByteString) -> Word32 -> IO ByteString
ioRefByteSource IORef (Maybe ByteString)
refSrc Word32
size = do
  forall a. IORef a -> IO a
readIORef IORef (Maybe ByteString)
refSrc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe ByteString
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
    Just ByteString
src -> do
      let taken :: ByteString
taken = Int -> ByteString -> ByteString
BS.take (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
size) ByteString
src
          rest :: ByteString
rest = Int -> ByteString -> ByteString
BS.drop (forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
size) ByteString
src
          stored :: Maybe ByteString
stored = if ByteString -> Bool
BS.null ByteString
taken then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just ByteString
rest
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc Maybe ByteString
stored
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
taken


ioRefByteSink :: Bool -> IORef [ByteString] -> IORef (Maybe ByteString) -> ByteString -> IO ()
ioRefByteSink :: Bool
-> IORef [ByteString]
-> IORef (Maybe ByteString)
-> ByteString
-> IO ()
ioRefByteSink Bool
debug IORef [ByteString]
refResponses IORef (Maybe ByteString)
refSrc ByteString
_ignored = do
  let asHex :: ByteString -> ByteString
asHex = Builder -> ByteString
toLazyByteString forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Builder
byteStringHex
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn forall a b. (a -> b) -> a -> b
$ ByteString
"bytesink got: " forall a. Semigroup a => a -> a -> a
<> (ByteString -> ByteString
asHex ByteString
_ignored)
  forall a. IORef a -> IO a
readIORef IORef [ByteString]
refResponses forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    [] -> do
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn ByteString
"bytesource has nothing"
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc forall a. Maybe a
Nothing
    (ByteString
x : [ByteString]
xs) -> do
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn forall a b. (a -> b) -> a -> b
$ ByteString
"bytesink will reply with: " forall a. Semigroup a => a -> a -> a
<> (ByteString -> ByteString
asHex ByteString
x)
      forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just ByteString
x
      forall a. IORef a -> a -> IO ()
writeIORef IORef [ByteString]
refResponses forall a b. (a -> b) -> a -> b
$ [ByteString]
xs