{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE LambdaCase #-}

-----------------------------------------------------------------------------
-- |
-- Module      :  $module
-- Copyright   :  (c) Powerweave Inc.
-- License     :  BSD-3-Clause
-- Maintainer  :  Laurent René de Cotret
-- Portability :  portable
--
-- This module defines a function, 'flakyTest', to declare a test
-- which intermittently fails. Flaky tests can be retries using retry policies
-- provided by the "Control.Retry" module (from the @retry@ package).
--
-- To dynamically retry based on the result of a test, see 'flakyTestWithRetryAction' instead.
--
-- For example, you can retry test cases from @tasty-hunit@ like so:
--
-- @
-- import Test.Tasty.HUnit ( testCase ) -- from tasty-hunit
-- 
-- myFlakyTest :: TestTree
-- myFlakyTest = 'flakyTest' ('limitRetries' 5 <> 'constantDelay' 1000) $ testCase "some test case" $ do ... 
-- @
--
-- In the example above, the test will be retried up to 5 times, with a delay of 1000 microseconds between tries,
-- if a failure occurs.
--
module Test.Tasty.Flaky (
    -- * Test wrappers
    flakyTest
    , flakyTestWithRetryAction

    -- * Re-exports
    -- 
    -- | The following functions allow to construct 'RetryPolicyM IO' 
    -- from the "Control.Retry" module.
    , constantDelay
    , exponentialBackoff
    , fullJitterBackoff
    , fibonacciBackoff
    , limitRetries

    -- * Policy Transformers
    , limitRetriesByDelay
    , limitRetriesByCumulativeDelay
    , capDelay

) where

import Control.Retry hiding (RetryPolicy)
import Data.Functor ( (<&>) )
import Data.Tagged (Tagged, retag )
import Test.Tasty.Providers ( IsTest(..), Progress, Result, TestTree )
import Test.Tasty.Runners ( TestTree(..), Result(..), Progress(..), emptyProgress, resultSuccessful )
import Test.Tasty.Options ( OptionDescription, OptionSet )


-- | A test tree of type @t@, with an associated retry policy
data FlakyTest t
    = MkFlakyTest (RetryStatus -> Result -> IO RetryAction) (RetryPolicyM IO) t


-- | Modify the delay of a RetryPolicy (in microseconds).
-- Does not change whether or not a retry is performed.
modifyRetryPolicyDelay :: Functor m => (Int -> Int) -> RetryPolicyM m -> RetryPolicyM m
modifyRetryPolicyDelay :: forall (m :: * -> *).
Functor m =>
(Int -> Int) -> RetryPolicyM m -> RetryPolicyM m
modifyRetryPolicyDelay Int -> Int
f (RetryPolicyM RetryStatus -> m (Maybe Int)
p) = (RetryStatus -> m (Maybe Int)) -> RetryPolicyM m
forall (m :: * -> *).
(RetryStatus -> m (Maybe Int)) -> RetryPolicyM m
RetryPolicyM ((RetryStatus -> m (Maybe Int)) -> RetryPolicyM m)
-> (RetryStatus -> m (Maybe Int)) -> RetryPolicyM m
forall a b. (a -> b) -> a -> b
$ \RetryStatus
stat -> (Int -> Int) -> Maybe Int -> Maybe Int
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Int -> Int
f (Maybe Int -> Maybe Int) -> m (Maybe Int) -> m (Maybe Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> RetryStatus -> m (Maybe Int)
p RetryStatus
stat


-- | Mark any test as flaky.
--
-- If this test is not successful, it will be retried according to the supplied @'RetryPolicyM' 'IO'@.
-- See "Control.Retry" for documentation on how to specify a @'RetryPolicyM' 'IO'@.
--
-- For example, you can retry test cases from @tasty-hunit@ like so:
--
-- @
-- import Test.Tasty.HUnit ( testCase ) -- from tasty-hunit
--
-- myFlakyTest :: TestTree
-- myFlakyTest = 'flakyTest' ('limitRetries' 5 <> 'constantDelay' 1000) $ testCase "some test case" $ do ...
-- @
--
-- To dynamically retry based on the result of a test, see 'flakyTestWithRetryAction' instead.
flakyTest :: (RetryPolicyM IO) -> TestTree -> TestTree
flakyTest :: RetryPolicyM IO -> TestTree -> TestTree
flakyTest = (RetryStatus -> Result -> IO RetryAction)
-> RetryPolicyM IO -> TestTree -> TestTree
flakyTestWithRetryAction (\RetryStatus
_ Result
_ -> RetryAction -> IO RetryAction
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure RetryAction
ConsultPolicy)


-- | Mark any test as flaky. Like 'flakyTest', but allows for overriding retry policies
-- based on test results. Also see 'RetryAction'.
--
-- For example, if you only want to retry a test if the error message contains @"some error message"@:
--
-- @
-- import Test.Tasty.HUnit ( testCase ) -- from tasty-hunit
-- import Data.List ( isInfixOf )
--
-- myFlakyTest :: TestTree
-- myFlakyTest 
--     = 'flakyTestWithRetryAction' 
--              retryAction 
--              ('constantDelay' 1000)
--                  $ testCase "some test case" $ do ...
--     where
--         retryAction :: 'RetryStatus' -> 'Result' -> IO 'RetryAction'
--         retryAction _ result
--             | "some error message" ``isInfixOf`` show result = pure `ConsultPolicy`
--             | otherwise = pure `DontRetry`
-- @
--
-- @since 0.1.2.0
flakyTestWithRetryAction :: (RetryStatus -> Result -> IO RetryAction)
                         -> (RetryPolicyM IO) 
                         -> TestTree -> TestTree
flakyTestWithRetryAction :: (RetryStatus -> Result -> IO RetryAction)
-> RetryPolicyM IO -> TestTree -> TestTree
flakyTestWithRetryAction RetryStatus -> Result -> IO RetryAction
retryAction RetryPolicyM IO
policy = \case
    (SingleTest TestName
name t
t)           -> TestName -> FlakyTest t -> TestTree
forall t. IsTest t => TestName -> t -> TestTree
SingleTest TestName
name ((RetryStatus -> Result -> IO RetryAction)
-> RetryPolicyM IO -> t -> FlakyTest t
forall t.
(RetryStatus -> Result -> IO RetryAction)
-> RetryPolicyM IO -> t -> FlakyTest t
MkFlakyTest RetryStatus -> Result -> IO RetryAction
retryAction RetryPolicyM IO
policy t
t)
    (TestGroup TestName
name [TestTree]
subtree)      -> TestName -> [TestTree] -> TestTree
TestGroup TestName
name ((TestTree -> TestTree) -> [TestTree] -> [TestTree]
forall a b. (a -> b) -> [a] -> [b]
map TestTree -> TestTree
go [TestTree]
subtree)
    (PlusTestOptions OptionSet -> OptionSet
modOption TestTree
t) -> (OptionSet -> OptionSet) -> TestTree -> TestTree
PlusTestOptions OptionSet -> OptionSet
modOption (TestTree -> TestTree
go TestTree
t)
    (WithResource ResourceSpec a
spec IO a -> TestTree
f)         -> ResourceSpec a -> (IO a -> TestTree) -> TestTree
forall a. ResourceSpec a -> (IO a -> TestTree) -> TestTree
WithResource ResourceSpec a
spec (IO a -> TestTree
f (IO a -> TestTree) -> (TestTree -> TestTree) -> IO a -> TestTree
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> TestTree -> TestTree
go)
    (AskOptions OptionSet -> TestTree
f)                -> (OptionSet -> TestTree) -> TestTree
AskOptions ((OptionSet -> TestTree) -> TestTree)
-> (OptionSet -> TestTree) -> TestTree
forall a b. (a -> b) -> a -> b
$ \OptionSet
optionSet -> TestTree -> TestTree
go (OptionSet -> TestTree
f OptionSet
optionSet)
    (After DependencyType
depType Expr
expr TestTree
t)        -> DependencyType -> Expr -> TestTree -> TestTree
After DependencyType
depType Expr
expr (TestTree -> TestTree
go TestTree
t)
  where
    go :: TestTree -> TestTree
go = (RetryStatus -> Result -> IO RetryAction)
-> RetryPolicyM IO -> TestTree -> TestTree
flakyTestWithRetryAction RetryStatus -> Result -> IO RetryAction
retryAction RetryPolicyM IO
policy


instance IsTest t => IsTest (FlakyTest t) where
    run :: IsTest t => OptionSet -> FlakyTest t -> (Progress -> IO ()) -> IO Result
    run :: IsTest t =>
OptionSet -> FlakyTest t -> (Progress -> IO ()) -> IO Result
run OptionSet
opts (MkFlakyTest RetryStatus -> Result -> IO RetryAction
retryAction RetryPolicyM IO
policy t
test) Progress -> IO ()
progressCallback = RetryStatus -> IO Result
go RetryStatus
defaultRetryStatus
        where
            -- The logic below mimics the `retry` package's Control.Retry.retrying
            -- with one major difference: we annotate the final result
            -- to report how many retries have been performed, regardless of
            -- the final result.
            go :: RetryStatus -> IO Result
            go :: RetryStatus -> IO Result
go RetryStatus
status = do
                Result
result <- OptionSet -> t -> (Progress -> IO ()) -> IO Result
forall t.
IsTest t =>
OptionSet -> t -> (Progress -> IO ()) -> IO Result
run OptionSet
opts t
test Progress -> IO ()
progressCallback
                let done :: IO Result
done = Result -> IO Result
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result -> IO Result) -> Result -> IO Result
forall a b. (a -> b) -> a -> b
$ RetryStatus -> Result -> Result
annotateResult RetryStatus
status Result
result
                    consultPolicy :: RetryPolicyM IO -> IO Result
consultPolicy RetryPolicyM IO
policy' = do
                        Maybe RetryStatus
rs <- RetryPolicyM IO -> RetryStatus -> IO (Maybe RetryStatus)
forall (m :: * -> *).
MonadIO m =>
RetryPolicyM m -> RetryStatus -> m (Maybe RetryStatus)
applyAndDelay RetryPolicyM IO
policy' RetryStatus
status
                        case Maybe RetryStatus
rs of
                            -- We are done: no more retries
                            Maybe RetryStatus
Nothing -> IO Result
done
                            -- At least one more retry
                            Just RetryStatus
rs' -> do
                                Progress -> IO ()
progressCallback (RetryStatus -> Progress
annotateProgress RetryStatus
status)
                                RetryStatus -> IO Result
go (RetryStatus -> IO Result) -> RetryStatus -> IO Result
forall a b. (a -> b) -> a -> b
$! RetryStatus
rs'

                if Result -> Bool
resultSuccessful Result
result
                then IO Result
done
                else do
                  RetryAction
retry <- RetryStatus -> Result -> IO RetryAction
retryAction RetryStatus
status Result
result
                  case RetryAction
retry of
                    RetryAction
DontRetry -> IO Result
done
                    RetryAction
ConsultPolicy -> RetryPolicyM IO -> IO Result
consultPolicy RetryPolicyM IO
policy
                    ConsultPolicyOverrideDelay Int
delay ->
                      RetryPolicyM IO -> IO Result
consultPolicy (RetryPolicyM IO -> IO Result) -> RetryPolicyM IO -> IO Result
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> RetryPolicyM IO -> RetryPolicyM IO
forall (m :: * -> *).
Functor m =>
(Int -> Int) -> RetryPolicyM m -> RetryPolicyM m
modifyRetryPolicyDelay (Int -> Int -> Int
forall a b. a -> b -> a
const Int
delay) RetryPolicyM IO
policy

            annotateProgress :: RetryStatus -> Progress
            annotateProgress :: RetryStatus -> Progress
annotateProgress RetryStatus
status
                -- Recall that `rsIterNumber` starts at 0, so the first attempt is rsIterNumber + 1
                = Progress
emptyProgress{progressText=mconcat ["Attempt #", show (rsIterNumber status + 1), " failed"]}

            annotateResult :: RetryStatus -> Result -> Result
            annotateResult :: RetryStatus -> Result -> Result
annotateResult RetryStatus
status Result
result
                = Result
result { resultDescription = resultDescription result <> annotate status }
                where
                    annotate :: RetryStatus -> String
                    annotate :: RetryStatus -> TestName
annotate (RetryStatus Int
iternum Int
cumdelay Maybe Int
_)
                        | Int
iternum Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = TestName
""
                        | Bool
otherwise    = [TestName] -> TestName
forall a. Monoid a => [a] -> a
mconcat [TestName
" [", Int -> TestName
forall a. Show a => a -> TestName
show Int
iternum, TestName
" retries, ", Int -> TestName
forall a. Show a => a -> TestName
show Int
cumdelay, TestName
" μs delay]"]


    testOptions :: Tagged (FlakyTest t) [OptionDescription]
    testOptions :: Tagged (FlakyTest t) [OptionDescription]
testOptions = Tagged t [OptionDescription]
-> Tagged (FlakyTest t) [OptionDescription]
forall {k1} {k2} (s :: k1) b (t :: k2). Tagged s b -> Tagged t b
retag (Tagged t [OptionDescription]
forall t. IsTest t => Tagged t [OptionDescription]
testOptions :: Tagged t [OptionDescription])