WIP exception catching in test-frameworktest-catch-exc

2 files changed, 38 insertions, 20 deletions

diff --git a/chad-fast.cabal b/chad-fast.cabal
index e201683..2b8f6c7 100644
--- a/chad-fast.cabal
+++ b/chad-fast.cabal
@@ -80,7 +80,8 @@ library test-framework
     base,
     hedgehog,
     time,
-    transformers
+    transformers,
+    unliftio-core
   hs-source-dirs: test-framework
   default-language: Haskell2010
   ghc-options: -Wall -threaded
diff --git a/test-framework/Test/Framework.hs b/test-framework/Test/Framework.hs
index 4c7799b..53f5e7b 100644
--- a/test-framework/Test/Framework.hs
+++ b/test-framework/Test/Framework.hs
@@ -1,4 +1,7 @@
+{-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingVia #-}
 {-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ImplicitParams #-}
@@ -18,12 +21,16 @@ module Test.Framework (
 ) where
 
 import Control.Monad (forM)
-import Control.Monad.Trans.State.Strict
+import Control.Monad.Trans.Writer.CPS
 import Control.Monad.IO.Class
+import Control.Monad.IO.Unlift
+import Data.IORef (newIORef, readIORef, atomicModifyIORef')
 import Data.List (isInfixOf)
 import Data.Maybe (isJust, mapMaybe, fromJust)
+import Data.Monoid (Sum(..))
 import Data.String (fromString)
 import Data.Time.Clock
+import GHC.Generics (Generic, Generically(..))
 import System.Environment
 import System.Exit
 import System.IO (hFlush, hPutStrLn, stdout, stderr)
@@ -85,20 +92,6 @@ computeMaxLen = go 0
     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 }
@@ -135,15 +128,40 @@ defaultMain tree = do
       | optsHelp opts -> printUsage >> exitSuccess
       | otherwise -> runTests opts tree >>= exitWith
 
+data Stats = Stats
+  { statsOK :: Sum Int
+  , statsTotal :: Sum Int }
+  deriving (Show, Generic)
+  deriving (Semigroup, Monoid) via (Generically Stats)
+
+newtype M a = M (WriterT Stats IO a)
+  deriving newtype (Functor, Applicative, Monad, MonadIO)
+
+-- | Not totally exception-safe (may lose writes if an exception gets thrown),
+-- but we don't care about that here.
+instance MonadUnliftIO M where
+  withRunInIO f = M $ writerT $ do
+    accum <- newIORef mempty
+    res <- f $ \(M w) -> do (x, s') <- runWriterT w
+                            atomicModifyIORef' accum (\s -> (s <> s', x))
+    output <- readIORef accum
+    return (res, output)
+
+tellStats :: Stats -> M ()
+tellStats s = M (tell s)
+
+runM :: M a -> IO (a, Stats)
+runM (M w) = runWriterT w
+
 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
+      let !maxlen = computeMaxLen tree
       starttm <- getCurrentTime
-      (success, stats) <- runStateT m initStats
+      (success, stats) <- runM $ let ?maxlen = maxlen in go 0 tree
       endtm <- getCurrentTime
       printStats stats (diffUTCTime endtm starttm)
       return (if isJust success then ExitSuccess else ExitFailure 1)
@@ -183,8 +201,7 @@ runTests options = \tree' ->
       liftIO $ printResult report (diffUTCTime endtm starttm)
 
       let ok = H.reportStatus report == H.OK
-      modifyStats $ \stats -> stats { statsOK = fromEnum ok + statsOK stats
-                                    , statsTotal = 1 + statsTotal stats }
+      tellStats $ Stats { statsOK = Sum (fromEnum ok), statsTotal = 1 }
       return (if ok then Just 1 else Nothing)
 
 outputProgress :: Int -> H.Report H.Progress -> IO ()