Main.hs

{-# 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)