content: {-# LANGUAGE DuplicateRecordFields #-}
  {-# LANGUAGE NamedFieldPuns        #-}
  {-# LANGUAGE NoImplicitPrelude     #-}
  {-# LANGUAGE OverloadedStrings     #-}
  module Test.Util where
  
  import           RIO
  import qualified RIO.Text as T
  
  import Test.Tasty.HUnit (assertBool, assertFailure)
  import Text.Show.Pretty (ppShow)
  
  
  import qualified Control.Concurrent.Capataz       as SUT
  import qualified Control.Concurrent.Capataz.Event as SUT
  
  --------------------------------------------------------------------------------
  -- Util
  
  -- | Utility function that gets the type name of a Record through it's Show
  -- output.
  fetchRecordName :: Show a => a -> Text
  fetchRecordName = T.takeWhile (/= ' ') . tshow
  
  -- | Composes two predicate functions together with a boolean AND
  andP :: [a -> Bool] -> a -> Bool
  andP predList a = all ($ a) predList
  
  --------------------------------------------------------------------------------
  -- Assertions and Testers
  
  -- | This record duplicate the same event names as the ones found in the
  -- "CapatazEvent" type, we use this to avoid using Text comparisons on assertion
  -- helper functions. The "CapatazEvent" record is imported qualified, so there
  -- is no conflict happening.
  data EventType
    = InvalidCapatazStatusReached
    | SupervisorStatusChanged
    | ProcessTerminated
    | ProcessStarted
    | ProcessRestarted
    | ProcessCompleted
    | ProcessFailed
    | ProcessCallbackExecuted
    | ProcessTerminationStarted
    | ProcessTerminationFinished
    | CapatazFailed
    | CapatazTerminated
    deriving (Show, Eq)
  
  toEventType :: SUT.CapatazEvent -> EventType
  toEventType ev = case ev of
    SUT.InvalidSupervisorStatusReached{} -> InvalidCapatazStatusReached
    SUT.SupervisorStatusChanged{}        -> SupervisorStatusChanged
    SUT.ProcessTerminated{}              -> ProcessTerminated
    SUT.ProcessStarted{}                 -> ProcessStarted
    SUT.ProcessRestarted{}               -> ProcessRestarted
    SUT.ProcessCompleted{}               -> ProcessCompleted
    SUT.ProcessFailed{}                  -> ProcessFailed
    SUT.ProcessCallbackExecuted{}        -> ProcessCallbackExecuted
    SUT.ProcessTerminationStarted{}      -> ProcessTerminationStarted
    SUT.ProcessTerminationFinished{}     -> ProcessTerminationFinished
    SUT.CapatazFailed{}                  -> CapatazFailed
    SUT.CapatazTerminated{}              -> CapatazTerminated
  
  -- | Predicate function to assert "CapatazEvent" types
  assertEventType :: EventType -> SUT.CapatazEvent -> Bool
  assertEventType evType ev = toEventType ev == evType
  
  -- | Predicate function to assert "CapatazEvent" process type
  assertProcessType :: SUT.ProcessType -> SUT.CapatazEvent -> Bool
  assertProcessType processTy ev = case ev of
    SUT.ProcessFailed { processType }           -> processTy == processType
    SUT.ProcessTerminated { processType }       -> processTy == processType
    SUT.ProcessStarted { processType }          -> processTy == processType
    SUT.ProcessCallbackExecuted { processType } -> processTy == processType
    SUT.ProcessRestarted { processType }        -> processTy == processType
    _                                           -> False
  
  -- | Predicate function to assert "CapatazEvent" worker name
  assertProcessName :: Text -> SUT.CapatazEvent -> Bool
  assertProcessName processName' ev = case ev of
    SUT.ProcessRestarted { processName }  -> processName' == processName
    SUT.ProcessFailed { processName }     -> processName' == processName
    SUT.ProcessTerminated { processName } -> processName' == processName
    SUT.ProcessStarted { processName }    -> processName' == processName
    _                                     -> False
  
  -- | Predicate function to assert type of an error inside a "CapatazEvent"
  assertErrorType :: Text -> SUT.CapatazEvent -> Bool
  assertErrorType errType ev = case ev of
    SUT.CapatazFailed { supervisorError } ->
      fetchRecordName supervisorError == errType
    SUT.ProcessFailed { processError } -> fetchRecordName processError == errType
    SUT.ProcessCallbackExecuted { processCallbackError } ->
      case processCallbackError of
        Nothing            -> False
        Just originalError -> fetchRecordName originalError == errType
    _ -> False
  
  -- | Predicate function to assert type of callback executed inside a
  -- "CapatazEvent"
  assertCallbackType :: SUT.CallbackType -> SUT.CapatazEvent -> Bool
  assertCallbackType cbType ev = case ev of
    SUT.ProcessFailed { processError } -> case fromException processError of
      Just SUT.ProcessCallbackFailed { processCallbackType } ->
        cbType == processCallbackType
      _ -> False
    SUT.ProcessCallbackExecuted { processCallbackType } ->
      cbType == processCallbackType
    _ -> False
  
  -- | Predicate function to assert restart count inside a "CapatazEvent"
  assertRestartCount :: (Int -> Bool) -> SUT.CapatazEvent -> Bool
  assertRestartCount predFn ev = case ev of
    SUT.ProcessRestarted { processRestartCount } -> predFn processRestartCount
    _                                            -> False
  
  -- | Predicate function to assert a Capataz status change
  assertSupervisorStatusChanged
    :: SUT.SupervisorStatus -> SUT.SupervisorStatus -> SUT.CapatazEvent -> Bool
  assertSupervisorStatusChanged fromEv toEv ev = case ev of
    SUT.SupervisorStatusChanged { prevSupervisorStatus, newSupervisorStatus } ->
      fromEv == prevSupervisorStatus && toEv == newSupervisorStatus
    _ -> False
  
  -- | Predicate function to assert process was supervised by a a supervisor
  -- indicated by name
  assertSupervisorName :: Text -> SUT.CapatazEvent -> Bool
  assertSupervisorName supervisorName' ev = case ev of
    SUT.SupervisorStatusChanged { supervisorName } ->
      supervisorName' == supervisorName
    SUT.ProcessStarted { supervisorName }    -> supervisorName' == supervisorName
    SUT.ProcessTerminated { supervisorName } -> supervisorName' == supervisorName
    SUT.ProcessRestarted { supervisorName }  -> supervisorName' == supervisorName
    SUT.ProcessCompleted { supervisorName }  -> supervisorName' == supervisorName
    SUT.ProcessCallbackExecuted { supervisorName } ->
      supervisorName' == supervisorName
    SUT.ProcessTerminationStarted { supervisorName } ->
      supervisorName' == supervisorName
    SUT.ProcessTerminationFinished { supervisorName } ->
      supervisorName' == supervisorName
    _ -> False
  
  -- | Predicate function to assert a supervisor was started
  assertSupervisorStarted :: Text -> SUT.CapatazEvent -> Bool
  assertSupervisorStarted supervisorName = andP
    [ assertEventType ProcessStarted
    , assertProcessType SUT.SupervisorType
    , assertProcessName supervisorName
    ]
  
  -- | Predicate function to assert a supervisor was terminated
  assertSupervisorTerminated :: Text -> SUT.CapatazEvent -> Bool
  assertSupervisorTerminated supervisorName = andP
    [ assertEventType ProcessTerminated
    , assertProcessType SUT.SupervisorType
    , assertProcessName supervisorName
    ]
  
  assertSupervisorFailed :: Text -> SUT.CapatazEvent -> Bool
  assertSupervisorFailed supervisorName = andP
    [ assertEventType ProcessFailed
    , assertProcessType SUT.SupervisorType
    , assertProcessName supervisorName
    ]
  
  assertSupervisorRestarted :: Text -> SUT.CapatazEvent -> Bool
  assertSupervisorRestarted supervisorName = andP
    [ assertEventType ProcessRestarted
    , assertProcessType SUT.SupervisorType
    , assertProcessName supervisorName
    ]
  
  -- | Predicate function to assert a worker was started
  assertWorkerStarted :: Text -> SUT.CapatazEvent -> Bool
  assertWorkerStarted workerName = andP
    [ assertEventType ProcessStarted
    , assertProcessType SUT.WorkerType
    , assertProcessName workerName
    ]
  
  -- | Predicate function to assert a worker was terminated
  assertWorkerTerminated :: Text -> SUT.CapatazEvent -> Bool
  assertWorkerTerminated workerName = andP
    [ assertEventType ProcessTerminated
    , assertProcessType SUT.WorkerType
    , assertProcessName workerName
    ]
  
  assertWorkerFailed :: Text -> SUT.CapatazEvent -> Bool
  assertWorkerFailed workerName = andP
    [ assertEventType ProcessFailed
    , assertProcessType SUT.WorkerType
    , assertProcessName workerName
    ]
  
  assertWorkerRestarted :: Text -> SUT.CapatazEvent -> Bool
  assertWorkerRestarted workerName = andP
    [ assertEventType ProcessRestarted
    , assertProcessType SUT.WorkerType
    , assertProcessName workerName
    ]
  
  -- | Predicate function to assert a capataz thread failed with error type
  assertCapatazFailedWith :: Text -> SUT.CapatazEvent -> Bool
  assertCapatazFailedWith errorName =
    andP [assertEventType CapatazFailed, assertErrorType errorName]
  
  rootSupervisorName :: Text
  rootSupervisorName = "capataz-root-supervisor"
  
  assertRootSupervisor :: SUT.CapatazEvent -> Bool
  assertRootSupervisor = assertSupervisorName rootSupervisorName
  
  
  --------------------------------------------------------------------------------
  
  -- | Exception used to test failures inside Worker sub-routines
  data RestartingWorkerError
    = RestartingWorkerError
    deriving (Show)
  
  instance Exception RestartingWorkerError
  
  -- | Exception used to test failures inside Worker callback sub-routines
  data TimeoutError
    = TimeoutError
    deriving (Show)
  
  instance Exception TimeoutError
  
  -- | Utility function to create a Worker sub-routine that fails at least a
  -- number of times
  mkFailingSubRoutine
    :: (MonadIO m, MonadThrow m1, MonadIO m1)
    => Int  -- ^ Number of times the Worker sub-routine will fail
    -> m (m1 ()) -- ^ Sub-routine used on worker creation
  mkFailingSubRoutine failCount = do
    countRef <- newIORef failCount
    let subRoutine = do
          shouldFail <- atomicModifyIORef' countRef
                                           (\count -> (count - 1, count > 0))
          when shouldFail (throwM RestartingWorkerError)
  
    return subRoutine
  
  -- | A sub-routine that will complete for `initCount` amount of times. This
  -- function works great when testing `Permanent` strategies, as you would like
  -- to assert restart events once (if it keeps completing it will fill up the log
  -- with restart events)
  mkCompletingBeforeNRestartsSubRoutine :: MonadIO m => Int -> m (m ())
  mkCompletingBeforeNRestartsSubRoutine initCount = do
    countRef <- newIORef initCount
    let subRoutine = do
          shouldStop <- atomicModifyIORef' countRef
                                           (\count -> (count - 1, count > 0))
          if shouldStop then return () else forever $ threadDelay 1000100
    return subRoutine
  
  -- | A sub-routine that will complete once. This function works great when
  -- testing `Permanent` strategies, as you would like to assert restart events
  -- once (if it keeps completing it will fill up the log with restart events)
  mkCompletingOnceSubRoutine :: MonadIO m => m (m ())
  mkCompletingOnceSubRoutine = mkCompletingBeforeNRestartsSubRoutine 1
  
  -- | Utility function to build a test environment for a Capataz execution.
  -- It is composed by:
  --
  -- * List of assertions that represent events that should be triggered by the
  --   capataz instance in order
  --
  -- * A function to modify the default "CapatazOptions", this utility function injects
  --   a special @notifyEvent@ callback to execute given assertions.
  testCapatazStreamWithOptions
    :: (MonadUnliftIO m)
    => (SUT.CapatazOptions (RIO LogFunc) -> SUT.CapatazOptions (RIO LogFunc)) -- ^ Function to modify default
                                                  -- @CapatazOptions@
    -> [SUT.CapatazEvent -> Bool] -- ^ Assertions happening before setup function
                                  -- is called
    -> (SUT.Capataz (RIO LogFunc) -> RIO LogFunc ()) -- ^ Function used to test public the supervisor
                              -- public API (a.k.a setup function)
    -> [SUT.CapatazEvent -> Bool] -- ^ Assertions happening after the setup
                                  -- function
    -> [SUT.CapatazEvent -> Bool] -- ^ Assertions happening after the capataz
                                  -- record is terminated
    -> Maybe (SUT.CapatazEvent -> Bool) -- ^ An assertion checked across all
                                        -- @CapatazEvents@ that happened in a
                                        -- test, great when testing that an event
                                        -- __did not__ happen
    -> m ()
  testCapatazStreamWithOptions optionModFn preSetupAssertion setupFn postSetupAssertions postTeardownAssertions mAllEventsAssertion
    = do
  
      eventStream     <- newTQueueIO
      accRef          <- newIORef []
      pendingCountVar <- newIORef
        (sum $ fmap
          length
          [preSetupAssertion, postSetupAssertions, postTeardownAssertions]
        )
  
      logOptions <- logOptionsHandle stdout False
      withLogFunc logOptions $ \logFunc -> runRIO logFunc $ do
        capataz <- SUT.forkCapataz
          rootSupervisorName
          ( set
              SUT.onSystemEventL
              (\ev -> logDebug (display ev) >> trackEvent accRef eventStream ev)
          . optionModFn
          )
  
        -- We check preSetup assertions are met before we execute the setup
        -- function. This serves to test initialization of capataz instance
        runAssertions "PRE-SETUP"
                      (eventStream, accRef)
                      pendingCountVar
                      preSetupAssertion
                      capataz
  
        -- We execute the setup sub-routine, which is going to use the Capataz public
        -- API to assert events
        setupResult <- try (setupFn capataz)
  
        case setupResult of
          -- If the sub-routine fails, show exception
          Left  err -> liftIO $ assertFailure (show (err :: SomeException))
          Right _   -> do
            -- We now run post-setup assertions
            runAssertions "POST-SETUP"
                          (eventStream, accRef)
                          pendingCountVar
                          postSetupAssertions
                          capataz
  
            -- We now shutdown the capataz instance
            SUT.terminateCapataz_ capataz
  
            -- We run assertions for after the capataz has been shut down
            runAssertions "POST-TEARDOWN"
                          (eventStream, accRef)
                          pendingCountVar
                          postTeardownAssertions
                          capataz
  
            -- Lastly, we check if there is a function that we want to execute
            -- across all events that happened in the test, this is to assert the
            -- absence of an event
            case mAllEventsAssertion of
              Nothing                 -> return ()
              Just allEventsAssertion -> do
                events <- reverse <$> readIORef accRef
                liftIO $ assertBool
                  ("On AFTER-TEST, expected all events to match predicate, but didn't ("
                  <> show (length events)
                  <> " events tried)\n"
                  <> ppShow (zip ([0 ..] :: [Int]) events)
                  )
                  (all allEventsAssertion events)
   where
    -- Utility functions that runs the readEventLoop function with a timeout
    -- of a second, this way we can guarantee assertions are met without having
    -- to add @threadDelays@ to the test execution
    runAssertions stageName (eventStream, accRef) pendingCountVar assertions capataz
      = do
        raceResult <- race
          (liftIO $ threadDelay 1000100)
          (readEventLoop eventStream pendingCountVar assertions)
        case raceResult of
          Left _ -> do
            events       <- reverse <$> readIORef accRef
            pendingCount <- readIORef pendingCountVar
            SUT.terminateCapataz_ capataz
            liftIO $ assertFailure
              (  "On "
              <> stageName
              <> " stage, expected all assertions to match, but didn't ("
              <> show pendingCount
              <> " assertions remaining, "
              <> show (length events)
              <> " events tried)\n"
              <> ppShow (zip ([0 ..] :: [Int]) events)
              )
          Right _ -> return ()
  
  
    -- Sub-routine that accumulates all events that have happened in the Capataz
    -- instance so far
    trackEvent accRef eventStream event = do
      atomicModifyIORef' accRef (\old -> (event : old, ()))
      atomically $ writeTQueue eventStream event
  
    -- Sub-routine that reads the event stream, and ensures that all assertions
    -- are executed, this loop won't stop until all assertions are met
    readEventLoop eventStream pendingCount assertions = do
      writeIORef pendingCount (length assertions)
      case assertions of
        []                          -> return ()
        (assertionFn : assertions1) -> do
          event <- atomically $ readTQueue eventStream
          if assertionFn event
            then readEventLoop eventStream pendingCount assertions1
            else readEventLoop eventStream pendingCount assertions
  
  
  -- | A version of "testCapatazStreamWithOptions" that does not receive the
  -- function that modifies a "CapatazOptions" record.
  testCapatazStream
    :: (MonadUnliftIO m)
    => [SUT.CapatazEvent -> Bool] -- ^ Assertions happening before setup function
                                  -- is called
    -> (SUT.Capataz (RIO LogFunc) -> RIO LogFunc ()) -- ^ Function used to test public the supervisor
                              -- public API (a.k.a setup function)
    -> [SUT.CapatazEvent -> Bool] -- ^ Assertions happening after the setup
                                  -- function
    -> [SUT.CapatazEvent -> Bool] -- ^ Assertions happening after the capataz
                                  -- record is terminated
    -> Maybe (SUT.CapatazEvent -> Bool) -- ^ An assertion checked across all
                                        -- @CapatazEvents@ that happened in a
                                        -- test, great when testing that an event
                                        -- __did not__ happen
    -> m ()
  testCapatazStream = testCapatazStreamWithOptions id
  