path: root/test-framework/Test/Framework.hs

{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImplicitParams #-}
{-# LANGUAGE TupleSections #-}
module Test.Framework (
  TestTree,
  testGroup,
  testGroupCollapse,
  testProperty,
  withResource,
  withResource',
  runTests,
  defaultMain,
  Options(..),

  -- * Compatibility
  TestName,
) where

import Control.Monad (forM)
import Control.Monad.Trans.State.Strict
import Control.Monad.IO.Class
import Data.List (isInfixOf, intercalate)
import Data.Maybe (isJust, mapMaybe, fromJust)
import Data.String (fromString)
import Data.Time.Clock
import System.Environment
import System.Exit
import System.IO (hFlush, hPutStrLn, stdout, stderr)
import Text.Read (readMaybe)

import qualified Hedgehog as H
import qualified Hedgehog.Internal.Config as H
import qualified Hedgehog.Internal.Property as H
import qualified Hedgehog.Internal.Report as H
import qualified Hedgehog.Internal.Runner as H
import qualified Hedgehog.Internal.Seed as H.Seed


data TestTree
  = Group Bool String [TestTree]
  | forall a. Resource (IO a) (a -> IO ()) (a -> TestTree)
  | HP String H.Property

type TestName = String

testGroup :: String -> [TestTree] -> TestTree
testGroup = Group False

testGroupCollapse :: String -> [TestTree] -> TestTree
testGroupCollapse = Group True

-- | The @a -> TestTree@ function must use the @a@ only inside properties: when
-- not actually running properties, it will be passed 'undefined'.
withResource :: IO a -> (a -> IO ()) -> (a -> TestTree) -> TestTree
withResource = Resource

-- | Same caveats as 'withResource'.
withResource' :: IO a -> (a -> TestTree) -> TestTree
withResource' make fun = withResource make (\_ -> return ()) fun

testProperty :: String -> H.Property -> TestTree
testProperty = HP

filterTree :: Options -> TestTree -> Maybe TestTree
filterTree (Options { optsPattern = pat }) = go []
  where
    go path (Group collapse name trees) =
      case mapMaybe (go (name:path)) trees of
        [] -> Nothing
        trees' -> Just (Group collapse name trees')
    go path (Resource make free fun) =
      case go path (fun undefined) of
        Nothing -> Nothing
        Just _ -> Just $ Resource make free (fromJust . go path . fun)
    go path hp@(HP name _)
      | pat `isInfixOf` renderPath (name:path) = Just hp
      | otherwise = Nothing

    renderPath comps = "^" ++ intercalate "/" (reverse comps) ++ "$"

computeMaxLen :: TestTree -> Int
computeMaxLen = go 0
  where
    go :: Int -> TestTree -> Int
    go indent (Group True name trees) = maximum (2*indent + length name : map (go (indent+1)) trees)
    go indent (Group False _ trees) = maximum (0 : map (go (indent+1)) trees)
    go indent (Resource _ _ fun) = go indent (fun undefined)
    go indent (HP name _) = 2 * indent + length name

data Stats = Stats
  { statsOK :: Int
  , statsTotal :: Int }
  deriving (Show)

initStats :: Stats
initStats = Stats 0 0

newtype M a = M (StateT Stats IO a)
  deriving newtype (Functor, Applicative, Monad, MonadIO)

modifyStats :: (Stats -> Stats) -> M ()
modifyStats f = M (modify f)

data Options = Options
  { optsPattern :: String
  , optsHelp :: Bool
  , optsHedgehogReplay :: Maybe (H.Skip, H.Seed) }
  deriving (Show)

defaultOptions :: Options
defaultOptions = Options "" False Nothing

parseOptions :: [String] -> Options -> Either String Options
parseOptions [] opts = pure opts
parseOptions ("-h":args) opts = parseOptions args opts { optsHelp = True }
parseOptions ("--help":args) opts = parseOptions args opts { optsHelp = True }
parseOptions ("-p":arg:args) opts
  | optsPattern opts == "" = parseOptions args opts { optsPattern = arg }
  | otherwise = Left "Multiple '-p' arguments given"
parseOptions ("--hedgehog-replay":arg:args) opts
  | Nothing <- optsHedgehogReplay opts =
      let parsed = do
            (skipstr, ' ' : seedstr) <- return $ span (/= ' ') arg
            (,) <$> H.skipDecompress skipstr <*> readMaybe seedstr
      in case parsed of
           Just res -> parseOptions args opts { optsHedgehogReplay = Just res }
           Nothing -> Left "Invalid argument to '--hedgehog-replay'"
  | otherwise = Left "Multiple '--hedgehog-replay' arguments given"
parseOptions (arg:_) _ = Left $ "Unrecognised argument: '" ++ arg ++ "'"

printUsage :: IO ()
printUsage = do
  progname <- getProgName
  putStr $ unlines
    ["Usage: " ++ progname ++ " [options]"
    ,"Options:"
    ,"  -h / --help    Show this help"
    ,"  -p PATTERN     Only tests whose path contains PATTERN are run. The path of a"
    ,"                 test looks like: '^group1/group2/testname$'."
    ,"  --hedgehog-replay '{skip} {seed}'"
    ,"                 Skip to a particular generated Hedgehog test. Should be used"
    ,"                 with -p. Overrides 'propertySkip' in 'PropertyConfig' if set."]

defaultMain :: TestTree -> IO ()
defaultMain tree = do
  args <- getArgs
  case parseOptions args defaultOptions of
    Left err -> die err
    Right opts
      | optsHelp opts -> printUsage >> exitSuccess
      | otherwise -> runTests opts tree >>= exitWith

runTests :: Options -> TestTree -> IO ExitCode
runTests options = \tree' ->
  case filterTree options tree' of
    Nothing -> do hPutStrLn stderr "No tests matched the given pattern."
                  return (ExitFailure 1)
    Just tree -> do
      let M m = let ?maxlen = computeMaxLen tree in go 0 id tree
      starttm <- getCurrentTime
      (success, stats) <- runStateT m initStats
      endtm <- getCurrentTime
      printStats stats (diffUTCTime endtm starttm)
      return (if isJust success then ExitSuccess else ExitFailure 1)
  where
    -- If all tests are successful, returns the number of output lines produced
    go :: (?maxlen :: Int) => Int -> (String -> String) -> TestTree -> M (Maybe Int)
    go indent path (Group collapse name trees) = do
      liftIO $ putStrLn (replicate (2 * indent) ' ' ++ name)
      starttm <- liftIO getCurrentTime
      mlns <- fmap (fmap sum . sequence) . forM trees $
                go (indent + 1) (path . (name++) . ('/':))
      endtm <- liftIO getCurrentTime
      case mlns of
        Just lns | collapse -> do
          let thislen = 2*indent + length name
          liftIO $ putStrLn $ concat (replicate (lns+1) "\x1B[A\x1B[2K") ++ "\x1B[G" ++
                              replicate (2 * indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ' ++
                              "\x1B[32mOK\x1B[0m" ++
                              prettyDuration False (realToFrac (diffUTCTime endtm starttm))
          return (Just 1)
        _ -> return mlns
    go indent path (Resource make cleanup fun) = do
      value <- liftIO make
      success <- go indent path (fun value)
      liftIO $ cleanup value
      return success
    go indent path (HP name (H.Property config test)) = do
      let thislen = 2*indent + length name
      liftIO $ putStr (replicate (2*indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' ')
      liftIO $ hFlush stdout

      (config', seed) <- case optsHedgehogReplay options of
                           Just (skip, seed) -> return (config { H.propertySkip = Just skip }, seed)
                           Nothing -> (config,) <$> H.Seed.random

      starttm <- liftIO getCurrentTime
      report <- liftIO $ H.checkReport config' 0 seed test (outputProgress (?maxlen + 2))
      endtm <- liftIO getCurrentTime

      liftIO $ printResult report (path name) (diffUTCTime endtm starttm)

      let ok = H.reportStatus report == H.OK
      modifyStats $ \stats -> stats { statsOK = fromEnum ok + statsOK stats
                                    , statsTotal = 1 + statsTotal stats }
      return (if ok then Just 1 else Nothing)

outputProgress :: Int -> H.Report H.Progress -> IO ()
outputProgress indent report = do
  str <- H.renderProgress H.EnableColor (Just (fromString "")) report
  putStr (replace '\n' "  " str ++ "\x1B[" ++ show (indent+1) ++ "G")
  hFlush stdout

printResult :: H.Report H.Result -> String -> NominalDiffTime -> IO ()
printResult report path timeTaken = do
  str <- H.renderResult H.EnableColor (Just (fromString "")) report
  case H.reportStatus report of
    H.OK -> putStrLn ("\x1B[K" ++ str ++ prettyDuration False (realToFrac timeTaken))
    H.Failed failure -> do
      let H.Report { H.reportTests = count, H.reportDiscards = discards } = report
          replayInfo = H.skipCompress (H.SkipToShrink count discards (H.failureShrinkPath failure)) ++
                       " " ++ show (H.reportSeed report)
          suffix = "\n    Flags to reproduce: `-p '" ++ path ++ "' --hedgehog-replay '" ++ replayInfo ++ "'`"
      putStrLn ("\x1B[K" ++ str ++ suffix)
    _ -> putStrLn ("\x1B[K" ++ str)

printStats :: Stats -> NominalDiffTime -> IO ()
printStats stats timeTaken
  | statsOK stats == statsTotal stats = do
      putStrLn $ "\x1B[32mAll " ++ show (statsTotal stats) ++
                 " tests passed." ++ prettyDuration True (realToFrac timeTaken) ++ "\x1B[0m"
  | otherwise =
      let nfailed = statsTotal stats - statsOK stats
      in putStrLn $ "\x1B[31mFailed " ++ show nfailed ++ " out of " ++ show (statsTotal stats) ++
                    " tests." ++ prettyDuration True (realToFrac timeTaken) ++ "\x1B[0m"

prettyDuration :: Bool -> Double -> String
prettyDuration False x | x < 0.5 = ""
prettyDuration _ x =
  let str = show (round (x * 100) :: Int)
      str' = replicate (3 - length str) '0' ++ str
      (pre, post) = splitAt (length str' - 2) str'
  in " (" ++ pre ++ "." ++ post ++ "s)"

replace :: Eq a => a -> [a] -> [a] -> [a]
replace x ys = concatMap (\y -> if y == x then ys else [y])