{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}

-- |
-- Copyright: © 2018-2020 IOHK
-- License: Apache-2.0
--
-- Helper functions for testing.
--

module Test.Hspec.Extra
    ( aroundAll
    , it
    , itWithCustomTimeout
    , flakyBecauseOf
    ) where

import Prelude

import Control.Concurrent
    ( threadDelay )
import Control.Concurrent.Async
    ( AsyncCancelled, async, race, wait )
import Control.Concurrent.MVar
    ( MVar, newEmptyMVar, putMVar, takeMVar )
import Control.Exception
    ( SomeException, catch, throwIO )
import System.Environment
    ( lookupEnv )
import Test.Hspec
    ( ActionWith
    , Expectation
    , HasCallStack
    , Spec
    , SpecWith
    , afterAll
    , beforeAll
    , beforeWith
    , expectationFailure
    , pendingWith
    , specify
    )

-- | Run a 'bracket' resource acquisition function around all the specs. The
-- bracket opens before the first test case and closes after the last test case.
--
-- It works by actually spawning a new thread responsible for the resource
-- acquisition, passing the resource along to the parent threads via a shared
-- MVar. Then, there's a bit of logic to synchronize both threads and make sure
-- that:
--
-- a) The 'Resource Owner' thread is terminated when the main thread is done
--    with the resource.
--
-- b) The 'Main Thread' only exists when the resource owner has released the
--    resource. Exiting the main thread before the 'Resource Owner' has
--    released the resource could left a hanging resource open. This is
--    particularly annoying when the resource is a running process!
--
--     Main Thread            Resource Owner
--          x
--          |         Spawn
--          |----------------------->x
--          |                        |
--          |                        |-- Acquire resource
--          |     Send Resource      |
--          |<-----------------------|
--          |                        |
--          |                        |
--         ...                      ... Await main thread signal
--          |                        |
--          |                        |
--          |      Send Signal       |
--          |----------------------->|
--          |                        |
--          |                       ... Release resource
--          |      Send Done         |
--          |<-----------------------|
--          |                       Exit
--          |
--         Exit
--
aroundAll
    :: forall a.
       (HasCallStack)
    => (ActionWith a -> IO ())
    -> SpecWith a
    -> Spec
aroundAll :: (ActionWith a -> IO ()) -> SpecWith a -> Spec
aroundAll ActionWith a -> IO ()
acquire =
    IO (a, IO ()) -> SpecWith (a, IO ()) -> Spec
forall a. IO a -> SpecWith a -> Spec
beforeAll IO (a, IO ())
setup (SpecWith (a, IO ()) -> Spec)
-> (SpecWith a -> SpecWith (a, IO ())) -> SpecWith a -> Spec
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ActionWith (a, IO ()) -> SpecWith (a, IO ()) -> SpecWith (a, IO ())
forall a. HasCallStack => ActionWith a -> SpecWith a -> SpecWith a
afterAll ActionWith (a, IO ())
forall a b. (a, b) -> b
snd (SpecWith (a, IO ()) -> SpecWith (a, IO ()))
-> (SpecWith a -> SpecWith (a, IO ()))
-> SpecWith a
-> SpecWith (a, IO ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, IO ()) -> IO a) -> SpecWith a -> SpecWith (a, IO ())
forall b a. (b -> IO a) -> SpecWith a -> SpecWith b
beforeWith (a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> IO a) -> ((a, IO ()) -> a) -> (a, IO ()) -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, IO ()) -> a
forall a b. (a, b) -> a
fst)
  where
    setup :: IO (a, IO ())
    setup :: IO (a, IO ())
setup = do
        MVar a
resource <- IO (MVar a)
forall a. IO (MVar a)
newEmptyMVar
        MVar ()
release  <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
        MVar ()
done     <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar

        Async ()
pid <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ do
            ActionWith a -> IO ()
acquire (ActionWith a -> IO ()) -> ActionWith a -> IO ()
forall a b. (a -> b) -> a -> b
$ \a
a -> do
                MVar a -> ActionWith a
forall a. MVar a -> a -> IO ()
putMVar MVar a
resource a
a
                MVar () -> IO ()
await MVar ()
release
            MVar () -> IO ()
unlock MVar ()
done

        IO () -> IO a -> IO (Either () a)
forall a b. IO a -> IO b -> IO (Either a b)
race (Async () -> IO ()
forall a. Async a -> IO a
wait Async ()
pid) (MVar a -> IO a
forall a. MVar a -> IO a
takeMVar MVar a
resource) IO (Either () a) -> (Either () a -> IO (a, IO ())) -> IO (a, IO ())
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Left ()
_ ->
                IOError -> IO (a, IO ())
forall e a. Exception e => e -> IO a
throwIO (IOError -> IO (a, IO ())) -> IOError -> IO (a, IO ())
forall a b. (a -> b) -> a -> b
$ String -> IOError
userError String
"aroundAll: failed to setup"
            Right a
a -> (a, IO ()) -> IO (a, IO ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((a, IO ()) -> IO (a, IO ())) -> (a, IO ()) -> IO (a, IO ())
forall a b. (a -> b) -> a -> b
$ (a
a,) (IO () -> (a, IO ())) -> IO () -> (a, IO ())
forall a b. (a -> b) -> a -> b
$ do
                MVar () -> IO ()
unlock MVar ()
release
                MVar () -> IO ()
await MVar ()
done

-- | A drop-in replacement for 'it' that'll automatically retry a scenario once
-- if it fails, to cope with potentially flaky tests.
--
-- It also has a timeout of 10 minutes.
it :: HasCallStack => String -> ActionWith ctx -> SpecWith ctx
it :: String -> ActionWith ctx -> SpecWith ctx
it = Int -> String -> ActionWith ctx -> SpecWith ctx
forall ctx.
HasCallStack =>
Int -> String -> ActionWith ctx -> SpecWith ctx
itWithCustomTimeout (Int
10Int -> Int -> Int
forall a. Num a => a -> a -> a
*Int
minute)
  where
    minute :: Int
minute = Int
60

-- | Like @it@ but with a custom timeout, which makes it realistic to test.
itWithCustomTimeout
    :: HasCallStack
    => Int -- ^ Timeout in seconds.
    -> String
    -> ActionWith ctx
    -> SpecWith ctx
itWithCustomTimeout :: Int -> String -> ActionWith ctx -> SpecWith ctx
itWithCustomTimeout Int
sec String
title ActionWith ctx
action = String -> ActionWith ctx -> SpecWith (Arg (ActionWith ctx))
forall a.
(HasCallStack, Example a) =>
String -> a -> SpecWith (Arg a)
specify String
title (ActionWith ctx -> SpecWith (Arg (ActionWith ctx)))
-> ActionWith ctx -> SpecWith (Arg (ActionWith ctx))
forall a b. (a -> b) -> a -> b
$ \ctx
ctx -> Int -> IO () -> IO ()
timeout Int
sec (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    ActionWith ctx
action ctx
ctx
        IO () -> (AsyncCancelled -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(AsyncCancelled
_ :: AsyncCancelled) -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ())
        IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(SomeException
e :: SomeException)  -> ActionWith ctx
action ctx
ctx
        IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` (\(SomeException
_ :: SomeException)  -> SomeException -> IO ()
forall e a. Exception e => e -> IO a
throwIO SomeException
e))
  where
    timeout :: Int -> IO () -> IO ()
timeout Int
t IO ()
act =
        IO () -> IO () -> IO (Either () ())
forall a b. IO a -> IO b -> IO (Either a b)
race (Int -> IO ()
threadDelay (Int -> Int
micro Int
t)) IO ()
act IO (Either () ()) -> (Either () () -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            Right () ->
                () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            Left () ->
                HasCallStack => String -> IO ()
String -> IO ()
expectationFailure (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$  String
"timed out in " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
t String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
" seconds"
      where
        micro :: Int -> Int
micro = (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000) (Int -> Int) -> (Int -> Int) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1000)

-- | Some helper to help readability on the thread synchronization above.
await :: MVar () -> IO ()
await :: MVar () -> IO ()
await = MVar () -> IO ()
forall a. MVar a -> IO a
takeMVar

-- | Some helper to help readability on the thread synchronization above.
unlock  :: MVar () -> IO ()
unlock :: MVar () -> IO ()
unlock = (MVar () -> () -> IO ()) -> () -> MVar () -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar ()

-- | Mark a test pending because of flakiness, with given reason. Unless the
-- RUN_FLAKY_TESTS environment variable is set.
flakyBecauseOf :: String -> Expectation
flakyBecauseOf :: String -> IO ()
flakyBecauseOf String
ticketOrReason =
    String -> IO (Maybe String)
lookupEnv String
"RUN_FLAKY_TESTS" IO (Maybe String) -> (Maybe String -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Just String
_ -> () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Maybe String
Nothing -> HasCallStack => String -> IO ()
String -> IO ()
pendingWith (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Flaky: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
ticketOrReason