{-# LANGUAGE LambdaCase #-}
{-# 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,
) where

import Control.Exception (catch)
import qualified Data.Attoparsec.ByteString as A
import Data.Attoparsec.Framer
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
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 :: * -> *) a.
MonadThrow m =>
Parser a -> (a -> m ()) -> (Word32 -> m ByteString) -> Framer m a
mkFramer Parser a
parser a -> IO ()
updateDst Word32 -> IO ByteString
src
  forall (m :: * -> *) a. MonadThrow m => Framer m a -> 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