Simplify itWithCustomTimeout implementation

with Matthias suggestion.

Author: Anviking

lib/test-utils/src/Test/Hspec/Extra.hs

@@ -2,7 +2,6 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TupleSections #-}
-{-# LANGUAGE TypeApplications #-}
 
 -- |
 -- Copyright: © 2018-2020 IOHK
@@ -22,13 +21,11 @@ import Prelude
 import Control.Concurrent
     ( threadDelay )
 import Control.Concurrent.Async
-    ( AsyncCancelled, async, race, race_, wait )
+    ( AsyncCancelled, async, race, wait )
 import Control.Concurrent.MVar
     ( MVar, newEmptyMVar, putMVar, takeMVar )
 import Control.Exception
-    ( SomeException, fromException, throwIO, try )
-import Data.Maybe
-    ( isJust )
+    ( SomeException, catch, throwIO )
 import Test.Hspec
     ( ActionWith
     , HasCallStack
@@ -115,36 +112,29 @@ aroundAll acquire =
 it :: HasCallStack => String -> ActionWith ctx -> SpecWith ctx
 it = itWithCustomTimeout (10*minute)
   where
-    minute = 60*1000*1000
+    minute = 60
 
 -- | Like @it@ but with a custom timeout, which makes it realistic to test.
-itWithCustomTimeout :: HasCallStack => Int -> String -> ActionWith ctx -> SpecWith ctx
-itWithCustomTimeout micro title action = specify title $ \ctx -> do
-    race_ timeout $ do
-        res1 <- try $ action ctx
-        case res1 of
-            Right r1 -> return r1
-            Left (e1 :: SomeException)
-                -- If we time out during the first attempt, then we /should/ time
-                -- out.
-                | isAsyncCancelled e1 -> throwIO e1
-                | otherwise           -> do
-                res2 <- try $ action ctx
-                case res2 of
-                    -- If the second try fails, return the first error. The
-                    -- second error might not be helpful.
-                    Left (e2 :: SomeException)
-                        | isAsyncCancelled e2 -> throwIO e1
-                        | otherwise           -> throwIO e1
-                    Right r2 -> return r2
+itWithCustomTimeout
+    :: HasCallStack
+    => Int -- ^ Timeout in seconds.
+    -> String
+    -> ActionWith ctx
+    -> SpecWith ctx
+itWithCustomTimeout sec title action = specify title $ \ctx -> timeout sec $ do
+    action ctx
+        `catch` (\(_ :: AsyncCancelled) -> return ())
+        `catch` (\(e :: SomeException)  -> action ctx
+        `catch` (\(_ :: SomeException)  -> throwIO e))
   where
-    isAsyncCancelled = isJust . fromException @AsyncCancelled
-    timeout = do
-        threadDelay micro
-        expectationFailure
-            $ "timed out in "
-            <> show (micro `div` 1000000)
-            <> " seconds"
+    timeout t act =
+        race (threadDelay (micro t)) act >>= \case
+            Right () ->
+                return ()
+            Left () ->
+                expectationFailure $  "timed out in " <> show t <> " seconds"
+      where
+        micro = (* 1000) . (* 1000)
 
 -- | Some helper to help readability on the thread synchronization above.
 await :: MVar () -> IO ()

lib/test-utils/test/Test/Hspec/ExtraSpec.hs

@@ -56,11 +56,11 @@ spec = do
                     outcomes <- newIORef [failure, noRetry]
                     (dynamically outcomes) `shouldMatchHSpecIt` failure
         it "can time out" $ do
-            let s = 1000*1000
+            let micro = (1000*1000 *)
             let timeout = do
-                    threadDelay (10 * s)
+                    threadDelay (micro 10)
                     expectationFailure "should have timed out"
-            res <- run (Extra.itWithCustomTimeout $ 2 * s) timeout
+            res <- run (Extra.itWithCustomTimeout 2) timeout
             res `shouldContain` "timed out in 2 seconds"
 
   where