uniform-io/test/LimitedInput.hs

{-# LANGUAGE OverloadedStrings #-}

module LimitedInput where

import Distribution.TestSuite
import Base (simpleTest)
import System.IO.Error
import System.IO.Uniform.ByteString
import qualified System.IO.Uniform.Streamline as S
import System.Timeout (timeout)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C8
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
  ]

recvTest :: Bool -> ByteString -> Int -> S.Streamline IO ByteString -> IO Progress
recvTest readFirst input limit f = do
  r <- tryIOError $ test readFirst
  case r of
    Left e -> return . Finished . Fail $ "IO error: " ++ show e
    Right v -> return v
  where
    test recv = do
      to <- timeout 1000000 $ do
        (r, _) <- withByteStringIO input $ \t -> do
          S.withTarget t $ do
            when recv $ S.recieveN 1 >> return ()
            S.limitInput limit
            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
      case to of
        Nothing -> return . Finished . Fail $ "Execution blocked"
        Just v -> return v

parseAlphabet :: A.Parser ByteString
parseAlphabet = A.string "abcedefghijklmnopqrstuvwxyz"

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