{-# LANGUAGE RecordWildCards #-} module Test.Main ( -- * Utilities for testing your main function captureProcessResult , ProcessResult(..) , withStdin -- ** Re-export from System.Environment , withArgs -- ** Re-export from System.Exit , ExitCode(..) ) where import qualified Control.Exception as E import qualified Data.ByteString as B import GHC.IO.Handle (hDuplicate, hDuplicateTo) import System.Directory (removeFile, getTemporaryDirectory) import System.Environment (withArgs) import System.Exit (ExitCode(ExitSuccess)) 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(..)) -- | -- Capture stdout, stderr, and exit code of the given IO action. -- -- >>> let main = putStr "hello" -- >>> captureProcessResult main -- ProcessResult {prStdout = "hello", prStderr = "", prExitCode = ExitSuccess} -- -- If the IO action exit with error message, the exit code of result is 'ExitFailure'. -- -- >>> import System.IO -- >>> import System.Exit -- >>> let main = hPutStr stderr "OMG!" >> exitWith (ExitFailure 1) -- >>> captureProcessResult main -- ProcessResult {prStdout = "", prStderr = "OMG!", prExitCode = ExitFailure 1} 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 <- either id (const ExitSuccess) <$> E.try 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 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 setEnc from hSetBuffering from saveBuf setEnc h = maybe (hSetBinaryMode h True) (hSetEncoding h) saveEnc E.bracket redirect restore (const action) -- | -- Pass the ByteString to stdin of the given IO action. -- -- >>> import Data.ByteString.Char8 () -- >>> :set -XOverloadedStrings -- >>> let main = putStrLn . reverse =<< getLine -- >>> withStdin "abcde" main -- edcba 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)