test-framework: Support --hedgehog-shrinks

1 files changed, 20 insertions, 5 deletions

diff --git a/test-framework/Test/Framework.hs b/test-framework/Test/Framework.hs
index 7380df1..e0dc4b3 100644
--- a/test-framework/Test/Framework.hs
+++ b/test-framework/Test/Framework.hs
@@ -106,11 +106,13 @@ modifyStats f = M (modify f)
 data Options = Options
   { optsPattern :: String
   , optsHelp :: Bool
-  , optsHedgehogReplay :: Maybe (H.Skip, H.Seed) }
+  , optsHedgehogReplay :: Maybe (H.Skip, H.Seed)
+  , optsHedgehogShrinks :: Maybe Int
+  }
   deriving (Show)
 
 defaultOptions :: Options
-defaultOptions = Options "" False Nothing
+defaultOptions = Options "" False Nothing Nothing
 
 parseOptions :: [String] -> Options -> Either String Options
 parseOptions [] opts = pure opts
@@ -128,6 +130,10 @@ parseOptions ("--hedgehog-replay":arg:args) opts
            Just res -> parseOptions args opts { optsHedgehogReplay = Just res }
            Nothing -> Left "Invalid argument to '--hedgehog-replay'"
   | otherwise = Left "Multiple '--hedgehog-replay' arguments given"
+parseOptions ("--hedgehog-shrinks":arg:args) opts =
+  case readMaybe arg of
+    Just n -> parseOptions args opts { optsHedgehogShrinks = Just n }
+    Nothing -> Left "Invalid argument to '--hedgehog-shrinks'"
 parseOptions (arg:_) _ = Left $ "Unrecognised argument: '" ++ arg ++ "'"
 
 printUsage :: IO ()
@@ -195,9 +201,8 @@ runTests options = \tree' ->
       let outputPrefix = replicate (2*indent) ' ' ++ name ++ ": " ++ replicate (?maxlen - thislen) ' '
       when ?istty $ liftIO $ putStr outputPrefix >> hFlush stdout
 
-      (config', seed) <- case optsHedgehogReplay options of
-                           Just (skip, seed) -> return (config { H.propertySkip = Just skip }, seed)
-                           Nothing -> (config,) <$> H.Seed.random
+      let (config', seedfun) = applyHedgehogOptions options config
+      seed <- seedfun
 
       starttm <- liftIO getCurrentTime
       report <- liftIO $ H.checkReport config' 0 seed test (outputProgress (?maxlen + 2))
@@ -213,6 +218,16 @@ runTests options = \tree' ->
                                     , statsTotal = 1 + statsTotal stats }
       return (if ok then Just 1 else Nothing)
 
+applyHedgehogOptions :: MonadIO m => Options -> H.PropertyConfig -> (H.PropertyConfig, m H.Seed)
+applyHedgehogOptions opts config0 =
+  let (f1, seedfun) = case optsHedgehogReplay opts of
+                        Just (skip, seed) -> (\c -> c { H.propertySkip = Just skip }, return seed)
+                        Nothing -> (id, H.Seed.random)
+      f2 = case optsHedgehogShrinks opts of
+             Just n -> \c -> c { H.propertyShrinkLimit = H.ShrinkLimit n }
+             Nothing -> id
+  in (f2 . f1 $ config0, seedfun)
+
 outputProgress :: (?istty :: Bool) => Int -> H.Report H.Progress -> IO ()
 outputProgress indent report
   | ?istty = do