{-# 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 <- Maybe [ByteString] -> IO (IORef (Maybe [ByteString]))
forall a. a -> IO (IORef a)
newIORef Maybe [ByteString]
forall a. Maybe a
Nothing
  IORef [a]
dst <- [a] -> IO (IORef [a])
forall a. a -> IO (IORef a)
newIORef []
  let updateDst :: a -> IO ()
updateDst a
x = IORef [a] -> ([a] -> [a]) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' IORef [a]
dst ((:) a
x)
      mkChunks :: Int -> [ByteString]
mkChunks Int
n = [[ByteString]] -> [ByteString]
forall a. Monoid a => [a] -> a
mconcat ([[ByteString]] -> [ByteString]) -> [[ByteString]] -> [ByteString]
forall a b. (a -> b) -> a -> b
$ (a -> [ByteString]) -> [a] -> [[ByteString]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> ByteString -> [ByteString]
chunksOfN Int
n (ByteString -> [ByteString])
-> (a -> ByteString) -> a -> [ByteString]
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 = Word32 -> Framer IO a -> Framer IO a
forall (m :: * -> *) a. Word32 -> Framer m a -> Framer m a
setChunkSize Word32
chunkSize' (Framer IO a -> Framer IO a) -> Framer IO a -> Framer IO a
forall a b. (a -> b) -> a -> b
$ Parser a
-> (a -> IO ()) -> (Word32 -> IO ByteString) -> Framer IO a
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
  Framer IO a -> IO ()
forall (m :: * -> *) frame. MonadThrow m => Framer m frame -> m ()
runFramer Framer IO a
frames IO () -> (NoMoreInput -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(NoMoreInput
_e :: NoMoreInput) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())

  [a]
got <- IORef [a] -> IO [a]
forall a. IORef a -> IO a
readIORef IORef [a]
dst
  Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ [a]
got [a] -> [a] -> Bool
forall a. Eq a => a -> a -> Bool
== [a] -> [a]
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 Maybe (ByteString, ByteString)
forall a. Maybe a
Nothing else (ByteString, ByteString) -> Maybe (ByteString, ByteString)
forall a. a -> Maybe a
Just (ByteString
taken, Int -> ByteString -> ByteString
BS.drop Int
x ByteString
y)
   in (ByteString -> Maybe (ByteString, ByteString))
-> ByteString -> [ByteString]
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
  IORef (Maybe [ByteString]) -> IO (Maybe [ByteString])
forall a. IORef a -> IO a
readIORef IORef (Maybe [ByteString])
chunkStore IO (Maybe [ByteString])
-> (Maybe [ByteString] -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe [ByteString]
Nothing -> do
      IORef (Maybe [ByteString]) -> Maybe [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe [ByteString])
chunkStore (Maybe [ByteString] -> IO ()) -> Maybe [ByteString] -> IO ()
forall a b. (a -> b) -> a -> b
$ [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just ([ByteString] -> Maybe [ByteString])
-> [ByteString] -> Maybe [ByteString]
forall a b. (a -> b) -> a -> b
$ Int -> [ByteString]
initChunks (Int -> [ByteString]) -> Int -> [ByteString]
forall a b. (a -> b) -> a -> b
$ Word32 -> Int
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 [] -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
BS.empty
    Just (ByteString
x : [ByteString]
xs) -> do
      IORef (Maybe [ByteString]) -> Maybe [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe [ByteString])
chunkStore (Maybe [ByteString] -> IO ()) -> Maybe [ByteString] -> IO ()
forall a b. (a -> b) -> a -> b
$ [ByteString] -> Maybe [ByteString]
forall a. a -> Maybe a
Just [ByteString]
xs
      ByteString -> IO ByteString
forall a. a -> IO a
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 <- Maybe ByteString -> IO (IORef (Maybe ByteString))
forall a. a -> IO (IORef a)
newIORef Maybe ByteString
forall a. Maybe a
Nothing
  IORef [ByteString]
refSink <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef [ByteString]
responses
  (Word32 -> IO ByteString, ByteString -> IO ())
-> IO (Word32 -> IO ByteString, ByteString -> IO ())
forall a. a -> IO a
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 <- Maybe ByteString -> IO (IORef (Maybe ByteString))
forall a. a -> IO (IORef a)
newIORef Maybe ByteString
forall a. Maybe a
Nothing
  IORef [ByteString]
refSink <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef [ByteString]
responses
  (Word32 -> IO ByteString, ByteString -> IO ())
-> IO (Word32 -> IO ByteString, ByteString -> IO ())
forall a. a -> IO a
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
  IORef (Maybe ByteString) -> IO (Maybe ByteString)
forall a. IORef a -> IO a
readIORef IORef (Maybe ByteString)
refSrc IO (Maybe ByteString)
-> (Maybe ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    Maybe ByteString
Nothing -> ByteString -> IO ByteString
forall a. a -> IO a
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 (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word32
size) ByteString
src
          rest :: ByteString
rest = Int -> ByteString -> ByteString
BS.drop (Word32 -> Int
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 Maybe ByteString
forall a. Maybe a
Nothing else ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
rest
      IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc Maybe ByteString
stored
      ByteString -> IO ByteString
forall a. a -> IO a
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 (Builder -> ByteString)
-> (ByteString -> Builder) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Builder
byteStringHex
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString
"bytesink got: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (ByteString -> ByteString
asHex ByteString
_ignored)
  IORef [ByteString] -> IO [ByteString]
forall a. IORef a -> IO a
readIORef IORef [ByteString]
refResponses IO [ByteString] -> ([ByteString] -> IO ()) -> IO ()
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    [] -> do
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn ByteString
"bytesource has nothing"
      IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc Maybe ByteString
forall a. Maybe a
Nothing
    (ByteString
x : [ByteString]
xs) -> do
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> IO ()
C8.putStrLn (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString
"bytesink will reply with: " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (ByteString -> ByteString
asHex ByteString
x)
      IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe ByteString)
refSrc (Maybe ByteString -> IO ()) -> Maybe ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
x
      IORef [ByteString] -> [ByteString] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef [ByteString]
refResponses ([ByteString] -> IO ()) -> [ByteString] -> IO ()
forall a b. (a -> b) -> a -> b
$ [ByteString]
xs