uniform-io/test/LimitedInput.hs

{-# LANGUAGE OverloadedStrings #-}

module LimitedInput where

import Distribution.TestSuite
import Base (simpleTest)
import System.IO.Uniform.ByteString
import qualified System.IO.Uniform.Streamline as S
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.Attoparsec.ByteString as A
import Control.Monad

tests :: IO [Test]
tests = return [
  simpleTest "recieveLine" $ recvTest False "A test\n" 3 S.recieveLine,
  simpleTest "recieveN" $ recvTest False "A test\n" 3 $ S.recieveN 30,
  simpleTest "runAttoparsec" $ recvTest False "abcedefghijklmnopqrstuvwxyz" 3 $ runParser parseAlphabet,
  simpleTest "recieveLine" $ recvTest True "A test\n" 3 S.recieveLine,
  simpleTest "recieveN" $ recvTest True "A test\n" 3 $ S.recieveN 30,
  simpleTest "runAttoparsec" $ recvTest True "abcedefghijklmnopqrstuvwxyz" 3 $ runParser parseAlphabet,
  simpleTest "Attoparsec failure gets pushed back" attoparsecPushBack
  ]

recvTest :: Bool -> ByteString -> Int -> S.Streamline IO ByteString -> IO Progress
recvTest readFirst input limit f = fstIO $ withByteStringIO input $ \t -> S.withTarget t $ do
    when readFirst $ S.recieveN 1 >> return ()
    S.limitInput limit
    r <- f
    let extra = if readFirst then 1 else 0
    if BS.length r <= limit + extra
      then return . Finished $ Pass
      else return . Finished . Fail $ "Got too much text: " ++ show r

fstIO :: IO (a, b) -> IO a
fstIO f = fst <$> f

alphabet :: ByteString
alphabet = "abcedefghijklmnopqrstuvwxyz"

parseAlphabet :: A.Parser ByteString
parseAlphabet = A.string alphabet

parseNotAlphabet :: A.Parser ByteString
parseNotAlphabet = A.string . BS.reverse $ alphabet

runParser :: A.Parser ByteString -> S.Streamline IO ByteString
runParser p = fst <$> S.runAttoparsecAndReturn p

attoparsecPushBack :: IO Progress
attoparsecPushBack = fstIO $ withByteStringIO alphabet $ \t -> S.withTarget t $ do
  S.runAttoparsec parseNotAlphabet
  abc <- S.recieveN 3
  if abc == "abc"
    then return . Finished $ Pass
    else return . Finished . Fail $ "Read wrong data: " ++ show abc