test-framework: Allow filtering tests by substring

1 files changed, 67 insertions, 12 deletions

diff --git a/test-framework/Test/Framework.hs b/test-framework/Test/Framework.hs
index 9622686..4c7799b 100644
--- a/test-framework/Test/Framework.hs
+++ b/test-framework/Test/Framework.hs
@@ -11,6 +11,7 @@ module Test.Framework (
   withResource',
   runTests,
   defaultMain,
+  Options(..),
 
   -- * Compatibility
   TestName,
@@ -19,11 +20,13 @@ module Test.Framework (
 import Control.Monad (forM)
 import Control.Monad.Trans.State.Strict
 import Control.Monad.IO.Class
-import Data.Maybe (isJust)
+import Data.List (isInfixOf)
+import Data.Maybe (isJust, mapMaybe, fromJust)
 import Data.String (fromString)
 import Data.Time.Clock
+import System.Environment
 import System.Exit
-import System.IO (hFlush, stdout)
+import System.IO (hFlush, hPutStrLn, stdout, stderr)
 
 import qualified Hedgehog as H
 import qualified Hedgehog.Internal.Config as H
@@ -58,6 +61,21 @@ 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 (path++"/"++name)) 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` (path ++ "/" ++ name) = Just hp
+      | otherwise = Nothing
+
 computeMaxLen :: TestTree -> Int
 computeMaxLen = go 0
   where
@@ -81,17 +99,54 @@ newtype M a = M (StateT Stats IO a)
 modifyStats :: (Stats -> Stats) -> M ()
 modifyStats f = M (modify f)
 
+data Options = Options
+  { optsPattern :: String
+  , optsHelp :: Bool }
+  deriving (Show)
+
+defaultOptions :: Options
+defaultOptions = Options "" False
+
+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 (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'."]
+
 defaultMain :: TestTree -> IO ()
-defaultMain tree = runTests tree >>= exitWith
-
-runTests :: TestTree -> IO ExitCode
-runTests = \tree -> do
-  let M m = let ?maxlen = computeMaxLen tree in go 0 tree
-  starttm <- getCurrentTime
-  (success, stats) <- runStateT m initStats
-  endtm <- getCurrentTime
-  printStats stats (diffUTCTime endtm starttm)
-  return (if isJust success then ExitSuccess else ExitFailure 1)
+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 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 -> TestTree -> M (Maybe Int)