module Convex.Tasty.Streaming (
  streamingJsonReporter,
  listTestsJsonIngredient,
  streamingIngredients,
  defaultMainStreaming,
) where

import Control.Concurrent.Async (forConcurrently_)
import Control.Concurrent.MVar (MVar, newMVar, withMVar)
import Control.Concurrent.STM
import Control.Monad (when)
import Convex.Tasty.Streaming.TMSummary (
  TMRecorder,
  TMStoreOption (..),
  TraceRecorder (..),
  lookupThreatModelSummary,
  newTMStore,
  storeRecorder,
 )
import Convex.Tasty.Streaming.TreeMap (buildTestMap)
import Convex.Tasty.Streaming.Types
import Data.Aeson (encode)
import Data.ByteString.Lazy.Char8 qualified as BL8
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.IntMap.Strict (IntMap)
import Data.IntMap.Strict qualified as IntMap
import Data.Proxy (Proxy (..))
import Data.Tagged (Tagged (..))
import Data.Text qualified as Text
import Data.Typeable (Typeable)
import System.IO (BufferMode (..), hFlush, hSetBuffering, stdout)
import Test.Tasty (TestTree, defaultMainWithIngredients, localOption)
import Test.Tasty.Ingredients (Ingredient (..))
import Test.Tasty.Ingredients.ConsoleReporter (consoleTestReporter)
import Test.Tasty.Options (IsOption (..), OptionDescription (..), lookupOption, mkFlagCLParser, safeRead)
import Test.Tasty.Runners (
  FailureReason (..),
  Outcome (..),
  Progress (..),
  Result (..),
  Status (..),
  listingTests,
 )

-- | Command-line option to enable streaming JSON output
newtype StreamingJson = StreamingJson Bool
  deriving (StreamingJson -> StreamingJson -> Bool
(StreamingJson -> StreamingJson -> Bool)
-> (StreamingJson -> StreamingJson -> Bool) -> Eq StreamingJson
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: StreamingJson -> StreamingJson -> Bool
== :: StreamingJson -> StreamingJson -> Bool
$c/= :: StreamingJson -> StreamingJson -> Bool
/= :: StreamingJson -> StreamingJson -> Bool
Eq, Eq StreamingJson
Eq StreamingJson =>
(StreamingJson -> StreamingJson -> Ordering)
-> (StreamingJson -> StreamingJson -> Bool)
-> (StreamingJson -> StreamingJson -> Bool)
-> (StreamingJson -> StreamingJson -> Bool)
-> (StreamingJson -> StreamingJson -> Bool)
-> (StreamingJson -> StreamingJson -> StreamingJson)
-> (StreamingJson -> StreamingJson -> StreamingJson)
-> Ord StreamingJson
StreamingJson -> StreamingJson -> Bool
StreamingJson -> StreamingJson -> Ordering
StreamingJson -> StreamingJson -> StreamingJson
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: StreamingJson -> StreamingJson -> Ordering
compare :: StreamingJson -> StreamingJson -> Ordering
$c< :: StreamingJson -> StreamingJson -> Bool
< :: StreamingJson -> StreamingJson -> Bool
$c<= :: StreamingJson -> StreamingJson -> Bool
<= :: StreamingJson -> StreamingJson -> Bool
$c> :: StreamingJson -> StreamingJson -> Bool
> :: StreamingJson -> StreamingJson -> Bool
$c>= :: StreamingJson -> StreamingJson -> Bool
>= :: StreamingJson -> StreamingJson -> Bool
$cmax :: StreamingJson -> StreamingJson -> StreamingJson
max :: StreamingJson -> StreamingJson -> StreamingJson
$cmin :: StreamingJson -> StreamingJson -> StreamingJson
min :: StreamingJson -> StreamingJson -> StreamingJson
Ord, Typeable)

instance IsOption StreamingJson where
  defaultValue :: StreamingJson
defaultValue = Bool -> StreamingJson
StreamingJson Bool
False
  parseValue :: [Char] -> Maybe StreamingJson
parseValue = (Bool -> StreamingJson) -> Maybe Bool -> Maybe StreamingJson
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> StreamingJson
StreamingJson (Maybe Bool -> Maybe StreamingJson)
-> ([Char] -> Maybe Bool) -> [Char] -> Maybe StreamingJson
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Maybe Bool
forall a. Read a => [Char] -> Maybe a
safeRead
  optionName :: Tagged StreamingJson [Char]
optionName = [Char] -> Tagged StreamingJson [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"streaming-json"
  optionHelp :: Tagged StreamingJson [Char]
optionHelp = [Char] -> Tagged StreamingJson [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"Enable streaming NDJSON test output to stdout"
  optionCLParser :: Parser StreamingJson
optionCLParser = Mod FlagFields StreamingJson
-> StreamingJson -> Parser StreamingJson
forall v. IsOption v => Mod FlagFields v -> v -> Parser v
mkFlagCLParser Mod FlagFields StreamingJson
forall a. Monoid a => a
mempty (Bool -> StreamingJson
StreamingJson Bool
True)

-- | Command-line option to disable iteration trace collection
newtype NoTrace = NoTrace Bool
  deriving (NoTrace -> NoTrace -> Bool
(NoTrace -> NoTrace -> Bool)
-> (NoTrace -> NoTrace -> Bool) -> Eq NoTrace
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NoTrace -> NoTrace -> Bool
== :: NoTrace -> NoTrace -> Bool
$c/= :: NoTrace -> NoTrace -> Bool
/= :: NoTrace -> NoTrace -> Bool
Eq, Eq NoTrace
Eq NoTrace =>
(NoTrace -> NoTrace -> Ordering)
-> (NoTrace -> NoTrace -> Bool)
-> (NoTrace -> NoTrace -> Bool)
-> (NoTrace -> NoTrace -> Bool)
-> (NoTrace -> NoTrace -> Bool)
-> (NoTrace -> NoTrace -> NoTrace)
-> (NoTrace -> NoTrace -> NoTrace)
-> Ord NoTrace
NoTrace -> NoTrace -> Bool
NoTrace -> NoTrace -> Ordering
NoTrace -> NoTrace -> NoTrace
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: NoTrace -> NoTrace -> Ordering
compare :: NoTrace -> NoTrace -> Ordering
$c< :: NoTrace -> NoTrace -> Bool
< :: NoTrace -> NoTrace -> Bool
$c<= :: NoTrace -> NoTrace -> Bool
<= :: NoTrace -> NoTrace -> Bool
$c> :: NoTrace -> NoTrace -> Bool
> :: NoTrace -> NoTrace -> Bool
$c>= :: NoTrace -> NoTrace -> Bool
>= :: NoTrace -> NoTrace -> Bool
$cmax :: NoTrace -> NoTrace -> NoTrace
max :: NoTrace -> NoTrace -> NoTrace
$cmin :: NoTrace -> NoTrace -> NoTrace
min :: NoTrace -> NoTrace -> NoTrace
Ord, Typeable)

instance IsOption NoTrace where
  defaultValue :: NoTrace
defaultValue = Bool -> NoTrace
NoTrace Bool
False
  parseValue :: [Char] -> Maybe NoTrace
parseValue = (Bool -> NoTrace) -> Maybe Bool -> Maybe NoTrace
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> NoTrace
NoTrace (Maybe Bool -> Maybe NoTrace)
-> ([Char] -> Maybe Bool) -> [Char] -> Maybe NoTrace
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Maybe Bool
forall a. Read a => [Char] -> Maybe a
safeRead
  optionName :: Tagged NoTrace [Char]
optionName = [Char] -> Tagged NoTrace [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"no-trace"
  optionHelp :: Tagged NoTrace [Char]
optionHelp = [Char] -> Tagged NoTrace [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"Disable iteration trace collection (only effective with --streaming-json)"
  optionCLParser :: Parser NoTrace
optionCLParser = Mod FlagFields NoTrace -> NoTrace -> Parser NoTrace
forall v. IsOption v => Mod FlagFields v -> v -> Parser v
mkFlagCLParser Mod FlagFields NoTrace
forall a. Monoid a => a
mempty (Bool -> NoTrace
NoTrace Bool
True)

-- | Command-line option to list tests as JSON without running them
newtype ListTestsJson = ListTestsJson Bool
  deriving (ListTestsJson -> ListTestsJson -> Bool
(ListTestsJson -> ListTestsJson -> Bool)
-> (ListTestsJson -> ListTestsJson -> Bool) -> Eq ListTestsJson
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ListTestsJson -> ListTestsJson -> Bool
== :: ListTestsJson -> ListTestsJson -> Bool
$c/= :: ListTestsJson -> ListTestsJson -> Bool
/= :: ListTestsJson -> ListTestsJson -> Bool
Eq, Eq ListTestsJson
Eq ListTestsJson =>
(ListTestsJson -> ListTestsJson -> Ordering)
-> (ListTestsJson -> ListTestsJson -> Bool)
-> (ListTestsJson -> ListTestsJson -> Bool)
-> (ListTestsJson -> ListTestsJson -> Bool)
-> (ListTestsJson -> ListTestsJson -> Bool)
-> (ListTestsJson -> ListTestsJson -> ListTestsJson)
-> (ListTestsJson -> ListTestsJson -> ListTestsJson)
-> Ord ListTestsJson
ListTestsJson -> ListTestsJson -> Bool
ListTestsJson -> ListTestsJson -> Ordering
ListTestsJson -> ListTestsJson -> ListTestsJson
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: ListTestsJson -> ListTestsJson -> Ordering
compare :: ListTestsJson -> ListTestsJson -> Ordering
$c< :: ListTestsJson -> ListTestsJson -> Bool
< :: ListTestsJson -> ListTestsJson -> Bool
$c<= :: ListTestsJson -> ListTestsJson -> Bool
<= :: ListTestsJson -> ListTestsJson -> Bool
$c> :: ListTestsJson -> ListTestsJson -> Bool
> :: ListTestsJson -> ListTestsJson -> Bool
$c>= :: ListTestsJson -> ListTestsJson -> Bool
>= :: ListTestsJson -> ListTestsJson -> Bool
$cmax :: ListTestsJson -> ListTestsJson -> ListTestsJson
max :: ListTestsJson -> ListTestsJson -> ListTestsJson
$cmin :: ListTestsJson -> ListTestsJson -> ListTestsJson
min :: ListTestsJson -> ListTestsJson -> ListTestsJson
Ord, Typeable)

instance IsOption ListTestsJson where
  defaultValue :: ListTestsJson
defaultValue = Bool -> ListTestsJson
ListTestsJson Bool
False
  parseValue :: [Char] -> Maybe ListTestsJson
parseValue = (Bool -> ListTestsJson) -> Maybe Bool -> Maybe ListTestsJson
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> ListTestsJson
ListTestsJson (Maybe Bool -> Maybe ListTestsJson)
-> ([Char] -> Maybe Bool) -> [Char] -> Maybe ListTestsJson
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Char] -> Maybe Bool
forall a. Read a => [Char] -> Maybe a
safeRead
  optionName :: Tagged ListTestsJson [Char]
optionName = [Char] -> Tagged ListTestsJson [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"list-tests-json"
  optionHelp :: Tagged ListTestsJson [Char]
optionHelp = [Char] -> Tagged ListTestsJson [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"List all tests as a JSON object and exit without running"
  optionCLParser :: Parser ListTestsJson
optionCLParser = Mod FlagFields ListTestsJson
-> ListTestsJson -> Parser ListTestsJson
forall v. IsOption v => Mod FlagFields v -> v -> Parser v
mkFlagCLParser Mod FlagFields ListTestsJson
forall a. Monoid a => a
mempty (Bool -> ListTestsJson
ListTestsJson Bool
True)

{- | Internal option carrying a shared 'IORef Bool' that is set to 'True' by
the streaming reporter when @--streaming-json@ is active.  The
'TraceRecorder' callback checks this before emitting any events.
-}
newtype StreamingEnabledRef = StreamingEnabledRef (Maybe (IORef Bool))

instance IsOption StreamingEnabledRef where
  defaultValue :: StreamingEnabledRef
defaultValue = Maybe (IORef Bool) -> StreamingEnabledRef
StreamingEnabledRef Maybe (IORef Bool)
forall a. Maybe a
Nothing
  parseValue :: [Char] -> Maybe StreamingEnabledRef
parseValue = Maybe StreamingEnabledRef -> [Char] -> Maybe StreamingEnabledRef
forall a b. a -> b -> a
const Maybe StreamingEnabledRef
forall a. Maybe a
Nothing
  optionName :: Tagged StreamingEnabledRef [Char]
optionName = [Char] -> Tagged StreamingEnabledRef [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"streaming-enabled-ref"
  optionHelp :: Tagged StreamingEnabledRef [Char]
optionHelp = [Char] -> Tagged StreamingEnabledRef [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"internal: streaming enabled flag"

{- | Internal option carrying a shared 'MVar ()' so the reporter and the
'TraceRecorder' use the same output lock, preventing interleaved NDJSON lines.
-}
newtype OutputLockRef = OutputLockRef (Maybe (MVar ()))

instance IsOption OutputLockRef where
  defaultValue :: OutputLockRef
defaultValue = Maybe (MVar ()) -> OutputLockRef
OutputLockRef Maybe (MVar ())
forall a. Maybe a
Nothing
  parseValue :: [Char] -> Maybe OutputLockRef
parseValue = Maybe OutputLockRef -> [Char] -> Maybe OutputLockRef
forall a b. a -> b -> a
const Maybe OutputLockRef
forall a. Maybe a
Nothing
  optionName :: Tagged OutputLockRef [Char]
optionName = [Char] -> Tagged OutputLockRef [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"output-lock-ref"
  optionHelp :: Tagged OutputLockRef [Char]
optionHelp = [Char] -> Tagged OutputLockRef [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"internal: shared output lock"

{- | Internal option carrying a shared 'IORef' so the reporter can publish the
test map and the 'TraceRecorder' can read it back to resolve test IDs.
-}
newtype TestMapRef = TestMapRef (Maybe (IORef (IntMap TestInfo)))

instance IsOption TestMapRef where
  defaultValue :: TestMapRef
defaultValue = Maybe (IORef (IntMap TestInfo)) -> TestMapRef
TestMapRef Maybe (IORef (IntMap TestInfo))
forall a. Maybe a
Nothing
  parseValue :: [Char] -> Maybe TestMapRef
parseValue = Maybe TestMapRef -> [Char] -> Maybe TestMapRef
forall a b. a -> b -> a
const Maybe TestMapRef
forall a. Maybe a
Nothing
  optionName :: Tagged TestMapRef [Char]
optionName = [Char] -> Tagged TestMapRef [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"test-map-ref"
  optionHelp :: Tagged TestMapRef [Char]
optionHelp = [Char] -> Tagged TestMapRef [Char]
forall {k} (s :: k) b. b -> Tagged s b
Tagged [Char]
"internal: shared test map reference"

{- | The streaming JSON reporter ingredient.

When activated via @--streaming-json@, replaces console output with
newline-delimited JSON events streamed to stdout.
-}
streamingJsonReporter :: Ingredient
streamingJsonReporter :: Ingredient
streamingJsonReporter = [OptionDescription]
-> (OptionSet
    -> TestTree -> Maybe (StatusMap -> IO (Time -> IO Bool)))
-> Ingredient
TestReporter
  [ Proxy StreamingJson -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy StreamingJson
forall {k} (t :: k). Proxy t
Proxy :: Proxy StreamingJson)
  , Proxy NoTrace -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy NoTrace
forall {k} (t :: k). Proxy t
Proxy :: Proxy NoTrace)
  , Proxy TMStoreOption -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy TMStoreOption
forall {k} (t :: k). Proxy t
Proxy :: Proxy TMStoreOption)
  , Proxy TMRecorder -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy TMRecorder
forall {k} (t :: k). Proxy t
Proxy :: Proxy TMRecorder)
  , Proxy TraceRecorder -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy TraceRecorder
forall {k} (t :: k). Proxy t
Proxy :: Proxy TraceRecorder)
  , Proxy TestMapRef -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy TestMapRef
forall {k} (t :: k). Proxy t
Proxy :: Proxy TestMapRef)
  , Proxy StreamingEnabledRef -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy StreamingEnabledRef
forall {k} (t :: k). Proxy t
Proxy :: Proxy StreamingEnabledRef)
  , Proxy OutputLockRef -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy OutputLockRef
forall {k} (t :: k). Proxy t
Proxy :: Proxy OutputLockRef)
  ]
  ((OptionSet
  -> TestTree -> Maybe (StatusMap -> IO (Time -> IO Bool)))
 -> Ingredient)
-> (OptionSet
    -> TestTree -> Maybe (StatusMap -> IO (Time -> IO Bool)))
-> Ingredient
forall a b. (a -> b) -> a -> b
$ \OptionSet
opts TestTree
tree -> do
    let StreamingJson Bool
enabled = OptionSet -> StreamingJson
forall v. IsOption v => OptionSet -> v
lookupOption OptionSet
opts
    if Bool -> Bool
not Bool
enabled
      then Maybe (StatusMap -> IO (Time -> IO Bool))
forall a. Maybe a
Nothing
      else (StatusMap -> IO (Time -> IO Bool))
-> Maybe (StatusMap -> IO (Time -> IO Bool))
forall a. a -> Maybe a
Just ((StatusMap -> IO (Time -> IO Bool))
 -> Maybe (StatusMap -> IO (Time -> IO Bool)))
-> (StatusMap -> IO (Time -> IO Bool))
-> Maybe (StatusMap -> IO (Time -> IO Bool))
forall a b. (a -> b) -> a -> b
$ \StatusMap
statusMap -> do
        let TMStoreOption Maybe TMStore
mStore = OptionSet -> TMStoreOption
forall v. IsOption v => OptionSet -> v
lookupOption OptionSet
opts
            TestMapRef Maybe (IORef (IntMap TestInfo))
mTestMapRef = OptionSet -> TestMapRef
forall v. IsOption v => OptionSet -> v
lookupOption OptionSet
opts
            StreamingEnabledRef Maybe (IORef Bool)
mEnabledRef = OptionSet -> StreamingEnabledRef
forall v. IsOption v => OptionSet -> v
lookupOption OptionSet
opts

        -- Signal that streaming is active so the TraceRecorder callback
        -- (which checks the same IORef) actually emits events.
        -- When --no-trace is passed, leave the ref as False so that both
        -- trEnabled and recordIteration remain no-ops.
        let NoTrace Bool
noTrace = OptionSet -> NoTrace
forall v. IsOption v => OptionSet -> v
lookupOption OptionSet
opts
        case Maybe (IORef Bool)
mEnabledRef of
          Just IORef Bool
ref -> IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Bool
ref (Bool -> Bool
not Bool
noTrace)
          Maybe (IORef Bool)
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

        -- Set line buffering for streaming
        Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
LineBuffering

        -- Use the shared output lock if provided, otherwise create a new one
        -- (backward compatibility when the reporter is used without
        -- defaultMainStreaming).
        let OutputLockRef Maybe (MVar ())
mSharedLock = OptionSet -> OutputLockRef
forall v. IsOption v => OptionSet -> v
lookupOption OptionSet
opts
        MVar ()
outputLock <- IO (MVar ())
-> (MVar () -> IO (MVar ())) -> Maybe (MVar ()) -> IO (MVar ())
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ()) MVar () -> IO (MVar ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (MVar ())
mSharedLock
        let emit :: Event -> IO ()
emit Event
evt = MVar () -> (() -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar ()
outputLock ((() -> IO ()) -> IO ()) -> (() -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \()
_ -> Event -> IO ()
emitEvent Event
evt

        -- Build the test index -> metadata map
        IntMap TestInfo
testMap <- OptionSet -> TestTree -> IO (IntMap TestInfo)
buildTestMap OptionSet
opts TestTree
tree

        -- Populate the shared test map ref so TraceRecorder can resolve IDs
        case Maybe (IORef (IntMap TestInfo))
mTestMapRef of
          Just IORef (IntMap TestInfo)
ref -> IORef (IntMap TestInfo) -> IntMap TestInfo -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef (IntMap TestInfo)
ref IntMap TestInfo
testMap
          Maybe (IORef (IntMap TestInfo))
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

        -- Emit suite_started with full test list
        let testInfos :: [TestInfo]
testInfos = ((Int, TestInfo) -> TestInfo) -> [(Int, TestInfo)] -> [TestInfo]
forall a b. (a -> b) -> [a] -> [b]
map (Int, TestInfo) -> TestInfo
forall a b. (a, b) -> b
snd ([(Int, TestInfo)] -> [TestInfo])
-> [(Int, TestInfo)] -> [TestInfo]
forall a b. (a -> b) -> a -> b
$ IntMap TestInfo -> [(Int, TestInfo)]
forall a. IntMap a -> [(Int, a)]
IntMap.toAscList IntMap TestInfo
testMap
        Event -> IO ()
emit (Event -> IO ()) -> Event -> IO ()
forall a b. (a -> b) -> a -> b
$ [TestInfo] -> Event
SuiteStarted [TestInfo]
testInfos

        -- Track results for final summary
        TVar [(Int, Result)]
resultsVar <- [(Int, Result)] -> IO (TVar [(Int, Result)])
forall a. a -> IO (TVar a)
newTVarIO ([] :: [(Int, Result)])

        -- Watch each test concurrently
        [(Int, TVar Status)] -> ((Int, TVar Status) -> IO ()) -> IO ()
forall (f :: * -> *) a b. Foldable f => f a -> (a -> IO b) -> IO ()
forConcurrently_ (StatusMap -> [(Int, TVar Status)]
forall a. IntMap a -> [(Int, a)]
IntMap.toAscList StatusMap
statusMap) (((Int, TVar Status) -> IO ()) -> IO ())
-> ((Int, TVar Status) -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Int
idx, TVar Status
statusTVar) -> do
          -- Wait until the test starts
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            Status
status <- TVar Status -> STM Status
forall a. TVar a -> STM a
readTVar TVar Status
statusTVar
            case Status
status of
              Status
NotStarted -> STM ()
forall a. STM a
retry
              Status
_ -> () -> STM ()
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

          -- Emit test_started
          Event -> IO ()
emit (Event -> IO ()) -> Event -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Event
TestStarted Int
idx

          -- Wait for completion, emitting progress events along the way
          let waitLoop :: Maybe ([Char], Float) -> IO Result
waitLoop Maybe ([Char], Float)
lastSeen = do
                Either Progress Result
next <- STM (Either Progress Result) -> IO (Either Progress Result)
forall a. STM a -> IO a
atomically (STM (Either Progress Result) -> IO (Either Progress Result))
-> STM (Either Progress Result) -> IO (Either Progress Result)
forall a b. (a -> b) -> a -> b
$ do
                  Status
status <- TVar Status -> STM Status
forall a. TVar a -> STM a
readTVar TVar Status
statusTVar
                  case Status
status of
                    Status
NotStarted -> STM (Either Progress Result)
forall a. STM a
retry
                    Executing Progress
p ->
                      let cur :: ([Char], Float)
cur = (Progress -> [Char]
progressText Progress
p, Progress -> Float
progressPercent Progress
p)
                       in if ([Char], Float) -> Maybe ([Char], Float)
forall a. a -> Maybe a
Just ([Char], Float)
cur Maybe ([Char], Float) -> Maybe ([Char], Float) -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe ([Char], Float)
lastSeen
                            then STM (Either Progress Result)
forall a. STM a
retry
                            else Either Progress Result -> STM (Either Progress Result)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Progress -> Either Progress Result
forall a b. a -> Either a b
Left Progress
p)
                    Done Result
r -> Either Progress Result -> STM (Either Progress Result)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result -> Either Progress Result
forall a b. b -> Either a b
Right Result
r)
                case Either Progress Result
next of
                  Left Progress
p -> do
                    Event -> IO ()
emit (Event -> IO ()) -> Event -> IO ()
forall a b. (a -> b) -> a -> b
$
                      TestProgress
                        { epId :: Int
epId = Int
idx
                        , epMessage :: Text
epMessage = [Char] -> Text
Text.pack (Progress -> [Char]
progressText Progress
p)
                        , epPercent :: Float
epPercent = Progress -> Float
progressPercent Progress
p
                        }
                    Maybe ([Char], Float) -> IO Result
waitLoop (([Char], Float) -> Maybe ([Char], Float)
forall a. a -> Maybe a
Just (Progress -> [Char]
progressText Progress
p, Progress -> Float
progressPercent Progress
p))
                  Right Result
r -> Result -> IO Result
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
r
          Result
result <- Maybe ([Char], Float) -> IO Result
waitLoop Maybe ([Char], Float)
forall a. Maybe a
Nothing

          -- Record result
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar [(Int, Result)]
-> ([(Int, Result)] -> [(Int, Result)]) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar [(Int, Result)]
resultsVar ((Int
idx, Result
result) (Int, Result) -> [(Int, Result)] -> [(Int, Result)]
forall a. a -> [a] -> [a]
:)

          -- Look up structured threat-model summary by "<group>/<name>" key
          let testInfo :: Maybe TestInfo
testInfo = Int -> IntMap TestInfo -> Maybe TestInfo
forall a. Int -> IntMap a -> Maybe a
IntMap.lookup Int
idx IntMap TestInfo
testMap
              key :: [Char]
key = case Maybe TestInfo
testInfo of
                Just TestInfo
ti
                  | (Text
parent : [Text]
_) <- [Text] -> [Text]
forall a. [a] -> [a]
reverse (TestInfo -> [Text]
tiPath TestInfo
ti) ->
                      Text -> [Char]
Text.unpack Text
parent [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
"/" [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> Text -> [Char]
Text.unpack (TestInfo -> Text
tiName TestInfo
ti)
                Just TestInfo
ti -> Text -> [Char]
Text.unpack (TestInfo -> Text
tiName TestInfo
ti)
                Maybe TestInfo
Nothing -> [Char]
""
          Maybe ThreatModelSummary
mSummary <- case Maybe TMStore
mStore of
            Just TMStore
store -> TMStore -> [Char] -> IO (Maybe ThreatModelSummary)
lookupThreatModelSummary TMStore
store [Char]
key
            Maybe TMStore
Nothing -> Maybe ThreatModelSummary -> IO (Maybe ThreatModelSummary)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ThreatModelSummary
forall a. Maybe a
Nothing

          -- Emit test_done
          let outcome :: TestOutcome
outcome = case Result -> Outcome
resultOutcome Result
result of
                Outcome
Success -> TestOutcome
TestSuccess
                Failure FailureReason
reason ->
                  FailureInfo -> TestOutcome
TestFailure (FailureInfo -> TestOutcome) -> FailureInfo -> TestOutcome
forall a b. (a -> b) -> a -> b
$
                    FailureInfo
                      { fiReason :: Text
fiReason = [Char] -> Text
Text.pack ([Char] -> Text) -> [Char] -> Text
forall a b. (a -> b) -> a -> b
$ FailureReason -> [Char]
showFailureReason FailureReason
reason
                      , fiMessage :: Text
fiMessage = [Char] -> Text
Text.pack (Result -> [Char]
resultDescription Result
result)
                      }
          Event -> IO ()
emit (Event -> IO ()) -> Event -> IO ()
forall a b. (a -> b) -> a -> b
$
            TestDone
              { edId :: Int
edId = Int
idx
              , edOutcome :: TestOutcome
edOutcome = TestOutcome
outcome
              , edDuration :: Time
edDuration = Result -> Time
resultTime Result
result
              , edDescription :: Text
edDescription = [Char] -> Text
Text.pack (Result -> [Char]
resultDescription Result
result)
              , edThreatModel :: Maybe ThreatModelSummary
edThreatModel = Maybe ThreatModelSummary
mSummary
              }

        -- Emit suite_done summary
        [(Int, Result)]
allResults <- TVar [(Int, Result)] -> IO [(Int, Result)]
forall a. TVar a -> IO a
readTVarIO TVar [(Int, Result)]
resultsVar
        let passed :: Int
passed = [()] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [() | (Int
_, Result
r) <- [(Int, Result)]
allResults, Result -> Bool
isSuccess Result
r]
        let failed :: Int
failed = [(Int, Result)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Result)]
allResults Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
passed

        -- Return the "finalize" callback
        (Time -> IO Bool) -> IO (Time -> IO Bool)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Time -> IO Bool) -> IO (Time -> IO Bool))
-> (Time -> IO Bool) -> IO (Time -> IO Bool)
forall a b. (a -> b) -> a -> b
$ \Time
totalTime -> do
          Event -> IO ()
emit (Event -> IO ()) -> Event -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Time -> Event
SuiteDone Int
passed Int
failed Time
totalTime
          Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
failed Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0)

-- | Emit a single NDJSON event line to stdout
emitEvent :: Event -> IO ()
emitEvent :: Event -> IO ()
emitEvent Event
evt = do
  ByteString -> IO ()
BL8.putStrLn (Event -> ByteString
forall a. ToJSON a => a -> ByteString
encode Event
evt)
  Handle -> IO ()
hFlush Handle
stdout

-- | Check if a Result is a success
isSuccess :: Result -> Bool
isSuccess :: Result -> Bool
isSuccess Result
r = case Result -> Outcome
resultOutcome Result
r of
  Outcome
Success -> Bool
True
  Outcome
_ -> Bool
False

-- | Show a FailureReason as text
showFailureReason :: FailureReason -> String
showFailureReason :: FailureReason -> [Char]
showFailureReason FailureReason
TestFailed = [Char]
"TestFailed"
showFailureReason (TestThrewException SomeException
e) = [Char]
"TestThrewException: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e
showFailureReason (TestTimedOut Integer
n) = [Char]
"TestTimedOut: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Integer -> [Char]
forall a. Show a => a -> [Char]
show Integer
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"μs"
showFailureReason FailureReason
TestDepFailed = [Char]
"TestDepFailed"

{- | Find the Tasty test ID for a test identified by group name and category.
Searches the test map for a 'TestInfo' whose path contains the group name
and whose name matches the category (e.g. \"Positive tests\", \"Negative tests\").
Returns @-1@ as a fallback when the test is not found.
-}
findTestId :: IntMap TestInfo -> String -> String -> Int
findTestId :: IntMap TestInfo -> [Char] -> [Char] -> Int
findTestId IntMap TestInfo
testMap [Char]
group [Char]
category =
  let categoryName :: [Char]
categoryName = case [Char]
category of
        [Char]
"positive" -> [Char]
"Positive tests"
        [Char]
"negative" -> [Char]
"Negative tests"
        [Char]
other -> [Char]
other
      matches :: [(Int, TestInfo)]
matches =
        IntMap TestInfo -> [(Int, TestInfo)]
forall a. IntMap a -> [(Int, a)]
IntMap.toList (IntMap TestInfo -> [(Int, TestInfo)])
-> IntMap TestInfo -> [(Int, TestInfo)]
forall a b. (a -> b) -> a -> b
$
          (TestInfo -> Bool) -> IntMap TestInfo -> IntMap TestInfo
forall a. (a -> Bool) -> IntMap a -> IntMap a
IntMap.filter
            ( \TestInfo
ti ->
                [Char] -> Text
Text.pack [Char]
group Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` TestInfo -> [Text]
tiPath TestInfo
ti
                  Bool -> Bool -> Bool
&& TestInfo -> Text
tiName TestInfo
ti Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== [Char] -> Text
Text.pack [Char]
categoryName
            )
            IntMap TestInfo
testMap
   in case [(Int, TestInfo)]
matches of
        ((Int
testId, TestInfo
_) : [(Int, TestInfo)]
_) -> Int
testId
        [] -> -Int
1 -- fallback: test not found

{- | Ingredient that lists the test tree as JSON and exits without running tests.

Activated via @--list-tests-json@.
-}
listTestsJsonIngredient :: Ingredient
listTestsJsonIngredient :: Ingredient
listTestsJsonIngredient = [OptionDescription]
-> (OptionSet -> TestTree -> Maybe (IO Bool)) -> Ingredient
TestManager
  [Proxy ListTestsJson -> OptionDescription
forall v. IsOption v => Proxy v -> OptionDescription
Option (Proxy ListTestsJson
forall {k} (t :: k). Proxy t
Proxy :: Proxy ListTestsJson)]
  ((OptionSet -> TestTree -> Maybe (IO Bool)) -> Ingredient)
-> (OptionSet -> TestTree -> Maybe (IO Bool)) -> Ingredient
forall a b. (a -> b) -> a -> b
$ \OptionSet
opts TestTree
tree -> do
    let ListTestsJson Bool
enabled = OptionSet -> ListTestsJson
forall v. IsOption v => OptionSet -> v
lookupOption OptionSet
opts
    if Bool -> Bool
not Bool
enabled
      then Maybe (IO Bool)
forall a. Maybe a
Nothing
      else IO Bool -> Maybe (IO Bool)
forall a. a -> Maybe a
Just (IO Bool -> Maybe (IO Bool)) -> IO Bool -> Maybe (IO Bool)
forall a b. (a -> b) -> a -> b
$ do
        Handle -> BufferMode -> IO ()
hSetBuffering Handle
stdout BufferMode
LineBuffering
        IntMap TestInfo
testMap <- OptionSet -> TestTree -> IO (IntMap TestInfo)
buildTestMap OptionSet
opts TestTree
tree
        let testInfos :: [TestInfo]
testInfos = ((Int, TestInfo) -> TestInfo) -> [(Int, TestInfo)] -> [TestInfo]
forall a b. (a -> b) -> [a] -> [b]
map (Int, TestInfo) -> TestInfo
forall a b. (a, b) -> b
snd ([(Int, TestInfo)] -> [TestInfo])
-> [(Int, TestInfo)] -> [TestInfo]
forall a b. (a -> b) -> a -> b
$ IntMap TestInfo -> [(Int, TestInfo)]
forall a. IntMap a -> [(Int, a)]
IntMap.toAscList IntMap TestInfo
testMap
        Event -> IO ()
emitEvent (Event -> IO ()) -> Event -> IO ()
forall a b. (a -> b) -> a -> b
$ [TestInfo] -> Event
SuiteStarted [TestInfo]
testInfos
        Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

-- | Default ingredients with streaming reporter added
streamingIngredients :: [Ingredient]
streamingIngredients :: [Ingredient]
streamingIngredients = [Ingredient
listingTests, Ingredient
listTestsJsonIngredient, Ingredient
streamingJsonReporter, Ingredient
consoleTestReporter]

{- | Drop-in replacement for 'defaultMain' that supports @--streaming-json@.

If you bypass this entry point and wire 'streamingIngredients' manually,
threat-model summaries will not appear in the JSON output unless you
also call
@'localOption' ('TMStoreOption' (Just store)) . 'localOption' ('storeRecorder' store)@
on your tree (with a freshly-allocated store from 'newTMStore').
-}
defaultMainStreaming :: TestTree -> IO ()
defaultMainStreaming :: TestTree -> IO ()
defaultMainStreaming TestTree
tree = do
  TMStore
store <- IO TMStore
newTMStore
  IORef (IntMap TestInfo)
testMapRef <- IntMap TestInfo -> IO (IORef (IntMap TestInfo))
forall a. a -> IO (IORef a)
newIORef IntMap TestInfo
forall a. IntMap a
IntMap.empty
  IORef Bool
enabledRef <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False -- set to True by the reporter when --streaming-json is active
  -- Create a single shared output lock used by both the streaming reporter
  -- and the TraceRecorder so their NDJSON lines never interleave.
  MVar ()
outputLock <- () -> IO (MVar ())
forall a. a -> IO (MVar a)
newMVar ()
  -- Create a trace recorder that emits TestTrace events as NDJSON to stdout.
  -- The recorder reads the shared testMapRef (populated by the reporter at
  -- startup) to resolve the numeric Tasty test ID for each trace event.
  --
  -- Both 'trEnabled' and 'recordIteration' read the shared 'enabledRef',
  -- so when --streaming-json is NOT passed (or --no-trace is passed) the
  -- ref stays False: test bodies take the fast path and no NDJSON lines
  -- go to stdout.
  let traceRec :: TraceRecorder
traceRec =
        TraceRecorder
          { trEnabled :: IO Bool
trEnabled = IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
enabledRef
          , recordIteration :: [Char] -> [Char] -> Value -> IO ()
recordIteration = \[Char]
group [Char]
category Value
iterationJson -> do
              Bool
enabled <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
enabledRef
              Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
enabled (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                IntMap TestInfo
testMap <- IORef (IntMap TestInfo) -> IO (IntMap TestInfo)
forall a. IORef a -> IO a
readIORef IORef (IntMap TestInfo)
testMapRef
                let testId :: Int
testId = IntMap TestInfo -> [Char] -> [Char] -> Int
findTestId IntMap TestInfo
testMap [Char]
group [Char]
category
                MVar () -> (() -> IO ()) -> IO ()
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar ()
outputLock ((() -> IO ()) -> IO ()) -> (() -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \()
_ ->
                  Event -> IO ()
emitEvent (Event -> IO ()) -> Event -> IO ()
forall a b. (a -> b) -> a -> b
$
                    TestTrace
                      { ettTestId :: Int
ettTestId = Int
testId
                      , ettCategory :: Text
ettCategory = [Char] -> Text
Text.pack [Char]
category
                      , ettTrace :: Value
ettTrace = Value
iterationJson
                      }
          }
  let tree' :: TestTree
tree' =
        TMStoreOption -> TestTree -> TestTree
forall v. IsOption v => v -> TestTree -> TestTree
localOption (Maybe TMStore -> TMStoreOption
TMStoreOption (TMStore -> Maybe TMStore
forall a. a -> Maybe a
Just TMStore
store)) (TestTree -> TestTree) -> TestTree -> TestTree
forall a b. (a -> b) -> a -> b
$
          TMRecorder -> TestTree -> TestTree
forall v. IsOption v => v -> TestTree -> TestTree
localOption (TMStore -> TMRecorder
storeRecorder TMStore
store) (TestTree -> TestTree) -> TestTree -> TestTree
forall a b. (a -> b) -> a -> b
$
            TestMapRef -> TestTree -> TestTree
forall v. IsOption v => v -> TestTree -> TestTree
localOption (Maybe (IORef (IntMap TestInfo)) -> TestMapRef
TestMapRef (IORef (IntMap TestInfo) -> Maybe (IORef (IntMap TestInfo))
forall a. a -> Maybe a
Just IORef (IntMap TestInfo)
testMapRef)) (TestTree -> TestTree) -> TestTree -> TestTree
forall a b. (a -> b) -> a -> b
$
              StreamingEnabledRef -> TestTree -> TestTree
forall v. IsOption v => v -> TestTree -> TestTree
localOption (Maybe (IORef Bool) -> StreamingEnabledRef
StreamingEnabledRef (IORef Bool -> Maybe (IORef Bool)
forall a. a -> Maybe a
Just IORef Bool
enabledRef)) (TestTree -> TestTree) -> TestTree -> TestTree
forall a b. (a -> b) -> a -> b
$
                OutputLockRef -> TestTree -> TestTree
forall v. IsOption v => v -> TestTree -> TestTree
localOption (Maybe (MVar ()) -> OutputLockRef
OutputLockRef (MVar () -> Maybe (MVar ())
forall a. a -> Maybe a
Just MVar ()
outputLock)) (TestTree -> TestTree) -> TestTree -> TestTree
forall a b. (a -> b) -> a -> b
$
                  TraceRecorder -> TestTree -> TestTree
forall v. IsOption v => v -> TestTree -> TestTree
localOption TraceRecorder
traceRec TestTree
tree
  [Ingredient] -> TestTree -> IO ()
defaultMainWithIngredients [Ingredient]
streamingIngredients TestTree
tree'