{-# OPTIONS_GHC -Wno-redundant-constraints #-}
{-# LANGUAGE ScopedTypeVariables #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  $module
-- Copyright   :  (c) Powerweave Inc.
-- License     :  BSD-3-Clause
-- Maintainer  :  Laurent René de Cotret
-- Portability :  portable
--
-- This module defines a single 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).
--
--
-- 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 wrapper
    flakyTest

    -- * 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(..), resultSuccessful )
import Test.Tasty.Options ( OptionDescription, OptionSet )


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


-- | Mark any test as flaky.
--
-- If this test is not successful, it will be retries 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 ... 
-- @
--
-- You can retry individual tests (like the example above), or retry entire groups by wrapping
-- 'Test.Tasty.testGroup'.
--
flakyTest :: (RetryPolicyM IO) -> TestTree -> TestTree
flakyTest :: RetryPolicyM IO -> TestTree -> TestTree
flakyTest RetryPolicyM IO
policy (SingleTest TestName
name t
t) = TestName -> FlakyTest t -> TestTree
forall t. IsTest t => TestName -> t -> TestTree
SingleTest TestName
name (RetryPolicyM IO -> t -> FlakyTest t
forall t. RetryPolicyM IO -> t -> FlakyTest t
MkFlakyTest RetryPolicyM IO
policy t
t)
flakyTest RetryPolicyM IO
policy (TestGroup TestName
name [TestTree]
subtree) = TestName -> [TestTree] -> TestTree
TestGroup TestName
name ((TestTree -> TestTree) -> [TestTree] -> [TestTree]
forall a b. (a -> b) -> [a] -> [b]
map (RetryPolicyM IO -> TestTree -> TestTree
flakyTest RetryPolicyM IO
policy) [TestTree]
subtree)
flakyTest RetryPolicyM IO
policy (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
<&> RetryPolicyM IO -> TestTree -> TestTree
flakyTest RetryPolicyM IO
policy)
flakyTest RetryPolicyM IO
_ TestTree
other = TestTree
other


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 RetryPolicyM IO
policy t
test) Progress -> IO ()
callback = 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 ()
callback
                let 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
                            Maybe RetryStatus
Nothing -> 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
                            Just RetryStatus
rs' -> 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 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
                    else RetryPolicyM IO -> IO Result
consultPolicy RetryPolicyM IO
policy
            
            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
" [First try]"
                        | Bool
otherwise    = [TestName] -> TestName
forall a. Monoid a => [a] -> a
mconcat [TestName
" [", TestName
"Retried ", Int -> TestName
forall a. Show a => a -> TestName
show Int
iternum, TestName
" times, total delay of ", Int -> TestName
forall a. Show a => a -> TestName
show Int
cumdelay, TestName
" microseconds]"]


    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])