{-# LANGUAGE RecordWildCards #-}
module Test.Main
(
captureProcessResult
, ProcessResult(..)
, withStdin
, withEnv
, withArgs
, ExitCode(..)
) where
import qualified Control.Exception as E
import Control.Monad (mapM, mapM_)
import qualified Data.ByteString as B
import GHC.IO.Handle (hDuplicate, hDuplicateTo)
import System.Directory (removeFile, getTemporaryDirectory)
import System.Environment (withArgs, setEnv, unsetEnv, lookupEnv)
import System.Exit (ExitCode(ExitSuccess, ExitFailure))
import System.IO
( Handle
, SeekMode(AbsoluteSeek)
, stdin
, stderr
, stdout
, hClose
, hFlush
, hSetBinaryMode
, hGetBuffering
, hSetBuffering
, hGetEncoding
, hSetEncoding
, hSeek
, openBinaryTempFile
, withBinaryFile
, IOMode(ReadMode)
)
import Test.Main.Internal (ProcessResult(..))
captureProcessResult :: IO () -> IO ProcessResult
captureProcessResult action = do
tDir <- getTemporaryDirectory
withBinaryTmpFile tDir "test-stdout" $ \(_oPath, oHd) ->
withBinaryTmpFile tDir "test-stderr" $ \(_ePath, eHd) ->
redirectingHandle stdout oHd $
redirectingHandle stderr eHd $ do
(prExitCode, prException) <- captureExitCodeAndException action
prStdout <- readFromHead oHd stdout
prStderr <- readFromHead eHd stderr
return ProcessResult {..}
where
readFromHead tmpH stdH = do
hFlush stdH
hSeek tmpH AbsoluteSeek 0
B.hGetContents tmpH
captureExitCodeAndException act =
(act >> return (ExitSuccess, Nothing))
`E.catches` [E.Handler forExitCode, E.Handler forSomeException]
forExitCode :: ExitCode -> IO (ExitCode, Maybe E.SomeException)
forExitCode eCode = return (eCode, Nothing)
forSomeException :: E.SomeException -> IO (ExitCode, Maybe E.SomeException)
forSomeException ex = return (ExitFailure 1, Just ex)
withBinaryTmpFile :: FilePath -> String -> ((FilePath, Handle) -> IO a) -> IO a
withBinaryTmpFile parent name =
E.bracket
(openBinaryTempFile parent name)
(\(path, hd) -> do
hClose hd
removeFile path `E.catch` doNothing
)
where
doNothing :: IOError-> IO ()
doNothing _ = return ()
redirectingHandle :: Handle -> Handle -> IO r -> IO r
redirectingHandle from to action = do
saveEnc <- hGetEncoding from
saveBuf <- hGetBuffering from
let redirect = do
save <- hDuplicate from
hDuplicateTo to from
setEnc to
return save
restore save = do
hDuplicateTo save from
hClose save
setEnc from
hSetBuffering from saveBuf
setEnc h =
maybe (hSetBinaryMode h True) (hSetEncoding h) saveEnc
E.bracket redirect restore (const action)
withStdin :: B.ByteString -> IO a -> IO a
withStdin bs action =
E.bracket
prepareInputFile
removeFile
(\inPath ->
withBinaryFile inPath ReadMode $ \tmpHd ->
redirectingHandle stdin tmpHd action
)
where
prepareInputFile = do
tDir <- getTemporaryDirectory
E.bracket
(openBinaryTempFile tDir "test-stdin")
(\(_path, hd) -> hClose hd)
(\(path, hd) -> B.hPut hd bs >> return path)
withEnv :: [(String, Maybe String)] -> IO a -> IO a
withEnv newVarVals action =
E.bracket
(saveReplaces newVarVals)
replaces
(const action)
where
saveReplaces = mapM $ \varVal@(var, _val) ->
((,) <$> pure var <*> lookupEnv var) <* replace varVal
replaces = mapM_ replace
replace (var, Just val) = setEnv var val
replace (var, Nothing) = unsetEnv var