StartTree.hs

{-# LANGUAGE CPP #-}
{-# LANGUAGE TupleSections #-}

module Test.Sandwich.Interpreters.StartTree (
  startTree
  , runNodesSequentially
  , markAllChildrenWithResult
  ) where


import Control.Concurrent.Async
import Control.Concurrent.MVar
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.IO.Unlift
import Control.Monad.Logger
import Control.Monad.Trans.Reader
import Data.IORef
import qualified Data.List as L
import Data.Sequence hiding ((:>))
import qualified Data.Set as S
import Data.String.Interpolate
import qualified Data.Text as T
import Data.Time.Clock
import Data.Typeable
import GHC.Stack
import System.Directory
import System.FilePath
import System.IO
import Test.Sandwich.Formatters.Print
import Test.Sandwich.Formatters.Print.CallStacks
import Test.Sandwich.Formatters.Print.FailureReason
import Test.Sandwich.Formatters.Print.Logs
import Test.Sandwich.Formatters.Print.Printing
import Test.Sandwich.Interpreters.RunTree.Logging
import Test.Sandwich.Interpreters.RunTree.Util
import Test.Sandwich.RunTree
import Test.Sandwich.TestTimer
import Test.Sandwich.Types.RunTree
import Test.Sandwich.Types.Spec
import Test.Sandwich.Types.TestTimer
import Test.Sandwich.Util
import UnliftIO.Exception
import UnliftIO.STM


baseContextFromCommon :: RunNodeCommonWithStatus s l t -> BaseContext -> BaseContext
baseContextFromCommon (RunNodeCommonWithStatus {..}) bc@(BaseContext {}) =
  bc { baseContextPath = runTreeFolder }

startTree :: (MonadIO m, HasBaseContext context) => RunNode context -> context -> m (Async Result)
startTree node@(RunNodeBefore {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    (timed runTreeRecordTime (getBaseContext ctx) (runTreeLabel <> " (setup)") (runExampleM runNodeBefore ctx runTreeLogs (Just [i|Exception in before '#{runTreeLabel}' handler|]))) >>= \case
      (result@(Failure fr@(Pending {})), _setupStartTime, setupFinishTime) -> do
        markAllChildrenWithResult runNodeChildren ctx (Failure fr)
        return (result, mkSetupTimingInfo setupFinishTime)
      (result@(Failure fr), _setupStartTime, setupFinishTime) -> do
        markAllChildrenWithResult runNodeChildren ctx (Failure $ GetContextException Nothing (SomeExceptionWithEq $ toException fr))
        return (result, mkSetupTimingInfo setupFinishTime)
      (Success, _setupStartTime, setupFinishTime) -> do
        void $ runNodesSequentially runNodeChildren ctx
        return (Success, mkSetupTimingInfo setupFinishTime)
      (Cancelled, _setupStartTime, setupFinishTime) -> do
        return (Cancelled, mkSetupTimingInfo setupFinishTime)
      (DryRun, _setupStartTime, setupFinishTime) -> do
        void $ runNodesSequentially runNodeChildren ctx
        return (DryRun, mkSetupTimingInfo setupFinishTime)
startTree node@(RunNodeAfter {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    result <- liftIO $ newIORef (Success, emptyExtraTimingInfo)
    finally (void $ runNodesSequentially runNodeChildren ctx)
            (do
                (ret, teardownStartTime, _teardownFinishTime) <- timed runTreeRecordTime (getBaseContext ctx) (runTreeLabel <> " (teardown)") $
                  runExampleM runNodeAfter ctx runTreeLogs (Just [i|Exception in after '#{runTreeLabel}' handler|])
                writeIORef result (ret, mkTeardownTimingInfo teardownStartTime)
            )
    liftIO $ readIORef result
startTree node@(RunNodeIntroduce {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    result <- liftIO $ newIORef (Success, emptyExtraTimingInfo)
    bracket (do
                let asyncExceptionResult e = Failure $ GotAsyncException Nothing (Just [i|introduceWith #{runTreeLabel} alloc handler got async exception|]) (SomeAsyncExceptionWithEq e)
                flip withException (\(e :: SomeAsyncException) -> markAllChildrenWithResult runNodeChildrenAugmented ctx (asyncExceptionResult e)) $
                  timed runTreeRecordTime (getBaseContext ctx) (runTreeLabel <> " (setup)") $
                    runExampleM' runNodeAlloc ctx runTreeLogs (Just [i|Failure in introduce '#{runTreeLabel}' allocation handler|])
            )
            (\(ret, setupStartTime, setupFinishTime) -> case ret of
                Left failureReason -> writeIORef result (Failure failureReason, mkSetupTimingInfo setupStartTime)
                Right intro -> do
                  teardownStartTime <- getCurrentTime
                  addTeardownStartTimeToStatus runTreeStatus teardownStartTime
                  (ret', _, _) <- timed runTreeRecordTime (getBaseContext ctx) (runTreeLabel <> " (teardown)") $
                    runExampleM (runNodeCleanup intro) ctx runTreeLogs (Just [i|Failure in introduce '#{runTreeLabel}' cleanup handler|])
                  writeIORef result (ret', ExtraTimingInfo (Just setupFinishTime) (Just teardownStartTime))
            )
            (\(ret, _setupStartTime, setupFinishTime) -> do
                addSetupFinishTimeToStatus runTreeStatus setupFinishTime
                case ret of
                  Left failureReason@(Pending {}) -> do
                    -- TODO: add note about failure in allocation
                    markAllChildrenWithResult runNodeChildrenAugmented ctx (Failure failureReason)
                  Left failureReason -> do
                    -- TODO: add note about failure in allocation
                    markAllChildrenWithResult runNodeChildrenAugmented ctx (Failure $ GetContextException Nothing (SomeExceptionWithEq $ toException failureReason))
                  Right intro -> do
                    -- Special hack to modify the test timer profile via an introduce, without needing to track it everywhere.
                    -- It would be better to track the profile at the type level
                    let ctxFinal = case cast intro of
                          Just (TestTimerProfile t) -> modifyBaseContext ctx (\bc -> bc { baseContextTestTimerProfile = t })
                          Nothing -> ctx

                    void $ runNodesSequentially runNodeChildrenAugmented ((LabelValue intro) :> ctxFinal)
            )
    readIORef result
startTree node@(RunNodeIntroduceWith {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  didRunWrappedAction <- liftIO $ newIORef (Left (), emptyExtraTimingInfo)
  runInAsync node ctx $ do
    let wrappedAction = do
          let failureResult e = case fromException e of
                Just fr@(Pending {}) -> Failure fr
                _ -> Failure $ Reason Nothing [i|introduceWith '#{runTreeLabel}' handler threw exception|]
          flip withException (\e -> recordExceptionInStatus runTreeStatus e >> markAllChildrenWithResult runNodeChildrenAugmented ctx (failureResult e)) $ do
            beginningCleanupVar <- liftIO $ newIORef Nothing

            -- Record a start event in the test timing, if configured
            -- TODO: how to do this? The speedscope viewer crashes if you don't properly nest events
            -- let tt = if runTreeRecordTime then getTestTimer (getBaseContext ctx) else NullTestTimer
            -- let setupLabel = runTreeLabel <> " (setup)"
            -- let teardownLabel = runTreeLabel <> " (teardown)"
            -- handleStartEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack setupLabel)

            results <- runNodeIntroduceAction $ \intro -> do
              -- Record the end event in the test timing
              -- TODO: do we need to deal with exceptions here and in teardown? I think we we fail to emit the
              -- end event, the time profile won't be viewable.
              -- handleEndEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack setupLabel)
              setupFinishTime <- liftIO getCurrentTime
              addSetupFinishTimeToStatus runTreeStatus setupFinishTime

              (results, _, _) <- timed runTreeRecordTime (getBaseContext ctx) (runTreeLabel <> " (body)") $
                liftIO $ runNodesSequentially runNodeChildrenAugmented (LabelValue intro :> ctx)

              teardownStartTime <- liftIO getCurrentTime
              addTeardownStartTimeToStatus runTreeStatus teardownStartTime

              liftIO $ writeIORef beginningCleanupVar (Just teardownStartTime)
              -- handleStartEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack teardownLabel)
              liftIO $ writeIORef didRunWrappedAction (Right results, mkSetupTimingInfo setupFinishTime)
              return results

            liftIO (readIORef beginningCleanupVar) >>= \case
              Nothing -> return ()
              Just teardownStartTime -> do
                -- handleEndEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack teardownLabel)
                liftIO $ modifyIORef' didRunWrappedAction $ \(ret, timingInfo) ->
                  (ret, timingInfo { teardownStartTime = Just teardownStartTime })

            return results

          liftIO (readIORef didRunWrappedAction) >>= \case
            (Left (), timingInfo) -> return (Failure $ Reason Nothing [i|introduceWith '#{runTreeLabel}' handler didn't call action|], timingInfo)
            (Right _, timingInfo) -> return (Success, timingInfo)
    runExampleM' wrappedAction ctx runTreeLogs (Just [i|Exception in introduceWith '#{runTreeLabel}' handler|]) >>= \case
      Left err -> return (Failure err, emptyExtraTimingInfo)
      Right x -> pure x
startTree node@(RunNodeAround {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  didRunWrappedAction <- liftIO $ newIORef (Left (), emptyExtraTimingInfo)
  runInAsync node ctx $ do
    let wrappedAction = do
          let failureResult e = case fromException e of
                Just fr@(Pending {}) -> Failure fr
                _ -> Failure $ Reason Nothing [i|around '#{runTreeLabel}' handler threw exception|]
          flip withException (\e -> recordExceptionInStatus runTreeStatus e >> markAllChildrenWithResult runNodeChildren ctx (failureResult e)) $ do
            runNodeActionWith $ do
              setupFinishTime <- liftIO getCurrentTime
              addSetupFinishTimeToStatus runTreeStatus setupFinishTime
              results <- liftIO $ runNodesSequentially runNodeChildren ctx
              liftIO $ writeIORef didRunWrappedAction (Right results, mkSetupTimingInfo setupFinishTime)
              return results

          (liftIO $ readIORef didRunWrappedAction) >>= \case
            (Left (), timingInfo) -> return (Failure $ Reason Nothing [i|around '#{runTreeLabel}' handler didn't call action|], timingInfo)
            (Right _, timingInfo) -> return (Success, timingInfo)
    runExampleM' wrappedAction ctx runTreeLogs (Just [i|Exception in introduceWith '#{runTreeLabel}' handler|]) >>= \case
      Left err -> return (Failure err, emptyExtraTimingInfo)
      Right x -> pure x
startTree node@(RunNodeDescribe {..}) ctx' = do
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    ((L.length . L.filter isFailure) <$> runNodesSequentially runNodeChildren ctx) >>= \case
      0 -> return (Success, emptyExtraTimingInfo)
      n -> return (Failure (ChildrenFailed Nothing n), emptyExtraTimingInfo)
startTree node@(RunNodeParallel {..}) ctx' = do
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    ((L.length . L.filter isFailure) <$> runNodesConcurrently runNodeCommon runNodeChildren ctx) >>= \case
      0 -> return (Success, emptyExtraTimingInfo)
      n -> return (Failure (ChildrenFailed Nothing n), emptyExtraTimingInfo)
startTree node@(RunNodeIt {..}) ctx' = do
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    (, emptyExtraTimingInfo) <$> runExampleM runNodeExample ctx (runTreeLogs runNodeCommon) Nothing

-- * Util

runInAsync :: (HasBaseContext context, MonadIO m) => RunNode context -> context -> IO (Result, ExtraTimingInfo) -> m (Async Result)
runInAsync node ctx action = do
  let RunNodeCommonWithStatus {..} = runNodeCommon node
  let bc@(BaseContext {..}) = getBaseContext ctx
  let timerFn = if runTreeRecordTime then timeAction' (getTestTimer bc) baseContextTestTimerProfile (T.pack runTreeLabel) else id
  startTime <- liftIO getCurrentTime
  mvar <- liftIO newEmptyMVar
  myAsync <- liftIO $ asyncWithUnmask $ \unmask -> do
    flip withException (recordExceptionInStatus runTreeStatus) $ unmask $ do
      readMVar mvar
      (result, extraTimingInfo) <- timerFn action
      endTime <- liftIO getCurrentTime
      liftIO $ atomically $ writeTVar runTreeStatus $ Done startTime (setupFinishTime extraTimingInfo) (teardownStartTime extraTimingInfo) endTime result

      whenFailure result $ \reason -> do
        -- Make sure the folder exists, if configured
        whenJust baseContextPath $ createDirectoryIfMissing True

        -- Create error symlink when configured to
        case node of
          RunNodeDescribe {} -> return () -- These are just noisy so don't create them
          RunNodeParallel {} -> return () -- These are just noisy so don't create them
          _ -> do
            whenJust baseContextErrorSymlinksDir $ \errorsDir ->
              whenJust baseContextPath $ \dir -> do
                whenJust baseContextRunRoot $ \runRoot -> do
                  let symlinkBaseName = case runTreeLoc of
                        Nothing -> takeFileName dir
                        Just loc -> [i|#{srcLocFile loc}_line#{srcLocStartLine loc}_#{takeFileName dir}|]
                  let symlinkPath = errorsDir </> (nodeToFolderName symlinkBaseName 9999999 runTreeId)

                  -- Delete the symlink if it's already present. This can happen when re-running
                  -- a previously failed test
                  exists <- doesPathExist symlinkPath
                  when exists $ removePathForcibly symlinkPath

#ifndef mingw32_HOST_OS
                  -- Get a relative path from the error dir to the results dir. System.FilePath doesn't want to
                  -- introduce ".." components, so we have to do it ourselves
                  let errorDirDepth = L.length $ splitPath $ makeRelative runRoot errorsDir

                  -- Don't do createDirectoryLink on Windows, as creating symlinks is generally not allowed for users.
                  -- See https://security.stackexchange.com/questions/10194/why-do-you-have-to-be-an-admin-to-create-a-symlink-in-windows
                  -- TODO: could we detect if this permission is available?
                  let relativePath = joinPath (L.replicate errorDirDepth "..") </> (makeRelative runRoot dir)
                  liftIO $ createDirectoryLink relativePath symlinkPath
#endif

        -- Write failure info
        whenJust baseContextPath $ \dir -> do
          withFile (dir </> "failure.txt") AppendMode $ \h -> do
            -- Use the PrintFormatter to format failure.txt nicely
            let pf = defaultPrintFormatter {
                  printFormatterUseColor = False
                  , printFormatterLogLevel = Just LevelDebug
                  , printFormatterIncludeCallStacks = True
                  }
            flip runReaderT (pf, 0, h) $ do
              printFailureReason reason
              whenJust (failureCallStack reason) $ \cs -> do
                p "\n"
                printCallStack cs
              p "\n"
              printLogs runTreeLogs

      return result
  liftIO $ atomically $ writeTVar runTreeStatus $ Running startTime Nothing Nothing myAsync
  liftIO $ putMVar mvar ()
  return myAsync  -- TODO: fix race condition with writing to runTreeStatus (here and above)

-- | Run a list of children sequentially, cancelling everything on async exception
runNodesSequentially :: HasBaseContext context => [RunNode context] -> context -> IO [Result]
runNodesSequentially children ctx =
  flip withException (\(e :: SomeAsyncException) -> cancelAllChildrenWith children e) $
    forM (L.filter (shouldRunChild ctx) children) $ \child ->
      startTree child ctx >>= wait

-- | Run a list of children concurrently, cancelling everything on async exception
runNodesConcurrently :: forall context. HasBaseContext context => RunNodeCommon -> [RunNode context] -> context -> IO [Result]
runNodesConcurrently (RunNodeCommonWithStatus {runTreeLabel, runTreeId}) children ctx =
  flip withException (\(e :: SomeAsyncException) -> cancelAllChildrenWith children e) $
    mapM wait =<< sequence [startTree child (modifyTimingProfile ix ctx)
                           | (child, ix) <- L.zip runnableChildren [0..]]
  where
    runnableChildren = L.filter (shouldRunChild ctx) children

    leftPadWithZeros :: Int -> String
    leftPadWithZeros num = L.replicate (L.length (show (L.length runnableChildren)) - L.length (show num)) '0' <> show num

    modifyTimingProfile :: Int -> context -> context
    modifyTimingProfile n = flip modifyBaseContext (modifyTimingProfile' n)

    modifyTimingProfile' :: Int -> BaseContext -> BaseContext
    modifyTimingProfile' n bc@(BaseContext {..}) = bc {
      baseContextTestTimerProfile = baseContextTestTimerProfile <> [i|-#{runTreeLabel}-#{runTreeId}-#{leftPadWithZeros n}|]
      }

markAllChildrenWithResult :: (MonadIO m, HasBaseContext context') => [RunNode context] -> context' -> Result -> m ()
markAllChildrenWithResult children baseContext status = do
  now <- liftIO getCurrentTime
  forM_ (L.filter (shouldRunChild' baseContext) $ concatMap getCommons children) $ \child ->
    liftIO $ atomically $ modifyTVar (runTreeStatus child) $ \case
      Running {..} -> Done now statusSetupFinishTime statusTeardownStartTime now status
      done@(Done {}) -> done { statusResult = status }
      _ -> Done now Nothing Nothing now status

cancelAllChildrenWith :: [RunNode context] -> SomeAsyncException -> IO ()
cancelAllChildrenWith children e = do
  forM_ children $ \node ->
    readTVarIO (runTreeStatus $ runNodeCommon node) >>= \case
      Running {..} -> cancelWith statusAsync e
      NotStarted -> do
        now <- getCurrentTime
        let reason = GotAsyncException Nothing Nothing (SomeAsyncExceptionWithEq e)
        atomically $ writeTVar (runTreeStatus $ runNodeCommon node) (Done now Nothing Nothing now (Failure reason))
      _ -> return ()

shouldRunChild :: (HasBaseContext ctx) => ctx -> RunNodeWithStatus context s l t -> Bool
shouldRunChild ctx node = shouldRunChild' ctx (runNodeCommon node)

shouldRunChild' :: (HasBaseContext ctx) => ctx -> RunNodeCommonWithStatus s l t -> Bool
shouldRunChild' ctx common = case baseContextOnlyRunIds $ getBaseContext ctx of
  Nothing -> True
  Just ids -> (runTreeId common) `S.member` ids

-- * Running examples

runExampleM :: HasBaseContext r => ExampleM r () -> r -> TVar (Seq LogEntry) -> Maybe String -> IO Result
runExampleM ex ctx logs exceptionMessage = runExampleM' ex ctx logs exceptionMessage >>= \case
  Left err -> return $ Failure err
  Right () -> return Success

runExampleM' :: HasBaseContext r => ExampleM r a -> r -> TVar (Seq LogEntry) -> Maybe String -> IO (Either FailureReason a)
runExampleM' ex ctx logs exceptionMessage = do
  maybeTestDirectory <- getTestDirectory ctx
  let options = baseContextOptions $ getBaseContext ctx

  -- We want our handleAny call to be *inside* the withLogFn call, because
  -- withFile will catch IOException and fill in its own information, making the
  -- resulting error confusing
  withLogFn maybeTestDirectory options $ \logFn ->
    handleAny (wrapInFailureReasonIfNecessary exceptionMessage)
      (Right <$> (runLoggingT (runReaderT (unExampleT ex) ctx) logFn))

  where
    withLogFn :: Maybe FilePath -> Options -> (LogFn -> IO a) -> IO a
    withLogFn Nothing (Options {..}) action = action (logToMemory optionsSavedLogLevel logs)
    withLogFn (Just logPath) (Options {..}) action = withFile (logPath </> "test_logs.txt") AppendMode $ \h -> do
      hSetBuffering h LineBuffering
      action (logToMemoryAndFile optionsMemoryLogLevel optionsSavedLogLevel optionsLogFormatter logs h)

    getTestDirectory :: (HasBaseContext a) => a -> IO (Maybe FilePath)
    getTestDirectory (getBaseContext -> (BaseContext {..})) = case baseContextPath of
      Nothing -> return Nothing
      Just dir -> do
        createDirectoryIfMissing True dir
        return $ Just dir

    wrapInFailureReasonIfNecessary :: Maybe String -> SomeException -> IO (Either FailureReason a)
    wrapInFailureReasonIfNecessary msg e = return $ Left $ case fromException e of
      Just (x :: FailureReason) -> x
      _ -> case fromException e of
        Just (SomeExceptionWithCallStack e' cs) -> GotException (Just cs) msg (SomeExceptionWithEq (SomeException e'))
        _ -> GotException Nothing msg (SomeExceptionWithEq e)

addSetupFinishTimeToStatus :: (MonadIO m) => TVar Status -> UTCTime -> m ()
addSetupFinishTimeToStatus statusVar setupFinishTime = atomically $ modifyTVar statusVar $ \case
  status@(Running {}) -> status { statusSetupFinishTime = Just setupFinishTime }
  status@(Done {}) -> status { statusSetupFinishTime = Just setupFinishTime }
  x -> x

addTeardownStartTimeToStatus :: (MonadIO m) => TVar Status -> UTCTime -> m ()
addTeardownStartTimeToStatus statusVar t = atomically $ modifyTVar statusVar $ \case
  status@(Running {}) -> status { statusTeardownStartTime = Just t }
  status@(Done {}) -> status { statusTeardownStartTime = Just t }
  x -> x

recordExceptionInStatus :: (MonadIO m) => TVar Status -> SomeException -> m ()
recordExceptionInStatus status e = do
  endTime <- liftIO getCurrentTime
  let ret = case fromException e of
        Just (e' :: SomeAsyncException) -> Failure (GotAsyncException Nothing Nothing (SomeAsyncExceptionWithEq e'))
        _ -> case fromException e of
          Just (e' :: FailureReason) -> Failure e'
          _ -> Failure (GotException Nothing Nothing (SomeExceptionWithEq e))
  liftIO $ atomically $ modifyTVar status $ \case
    Running {..} -> Done statusStartTime statusSetupFinishTime statusTeardownStartTime endTime ret
    _ -> Done endTime Nothing Nothing endTime ret

timed :: (MonadUnliftIO m) => Bool -> BaseContext -> String -> m a -> m (a, UTCTime, UTCTime)
timed recordTime bc@(BaseContext {..}) label action = do
  let timerFn = if recordTime then timeAction' (getTestTimer bc) baseContextTestTimerProfile (T.pack label) else id

  startTime <- liftIO getCurrentTime
  ret <- timerFn action
  endTime <- liftIO getCurrentTime
  pure (ret, startTime, endTime)