{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

module OpenTelemetry.Instrumentation.Tasty (instrumentTestTree, instrumentTestTreeWithTracer) where

import Control.Exception (bracket)
import Data.Tagged (Tagged, retag)
import Data.Text qualified as T
import OpenTelemetry.Context (insertSpan, lookupSpan, removeSpan)
import OpenTelemetry.Context.ThreadLocal (adjustContext, getContext)
import OpenTelemetry.Trace.Core (Span, SpanStatus (Error, Ok), Tracer, addAttribute, createSpan, defaultSpanArguments, detectInstrumentationLibrary, endSpan, getGlobalTracerProvider, inSpan, makeTracer, setStatus, tracerOptions)
import Test.Tasty (TestTree, withResource)
import Test.Tasty.Options (OptionDescription)
import Test.Tasty.Providers (IsTest (run, testOptions))
import Test.Tasty.Runners (Outcome (Failure, Success), ResourceSpec (ResourceSpec), Result (Result, resultDescription, resultOutcome), TestTree (After, AskOptions, PlusTestOptions, SingleTest, TestGroup, WithResource))


{- | A test case with a wrapper function that can do some IO around the
test. We use the wrapper to set up spans appropriately.
-}
data WrappedTest t = WrappedTest
  {forall t. WrappedTest t -> forall a. IO a -> IO a
wrapper :: forall a. IO a -> IO a, forall t. WrappedTest t -> t
innerTest :: t}


instance IsTest t => IsTest (WrappedTest t) where
  run :: OptionSet -> WrappedTest t -> (Progress -> IO ()) -> IO Result
run OptionSet
opts (WrappedTest {forall a. IO a -> IO a
wrapper :: forall t. WrappedTest t -> forall a. IO a -> IO a
wrapper :: forall a. IO a -> IO a
wrapper, t
innerTest :: forall t. WrappedTest t -> t
innerTest :: t
innerTest}) Progress -> IO ()
progress =
    IO Result -> IO Result
forall a. IO a -> IO a
wrapper (IO Result -> IO Result) -> IO Result -> IO Result
forall a b. (a -> b) -> a -> b
$ do
      Context
ctx <- IO Context
forall (m :: * -> *). MonadIO m => m Context
getContext
      let mspan :: Maybe Span
mspan = Context -> Maybe Span
lookupSpan Context
ctx
      res :: Result
res@Result {Outcome
resultOutcome :: Result -> Outcome
resultOutcome :: Outcome
resultOutcome, String
resultDescription :: Result -> String
resultDescription :: String
resultDescription} <- OptionSet -> t -> (Progress -> IO ()) -> IO Result
forall t.
IsTest t =>
OptionSet -> t -> (Progress -> IO ()) -> IO Result
run OptionSet
opts t
innerTest Progress -> IO ()
progress
      case Maybe Span
mspan of
        Just Span
s -> do
          Span -> Text -> Text -> IO ()
forall (m :: * -> *) a.
(MonadIO m, ToAttribute a) =>
Span -> Text -> a -> m ()
addAttribute Span
s Text
"result.description" (String -> Text
T.pack String
resultDescription)
          case Outcome
resultOutcome of
            Outcome
Success -> do
              Span -> SpanStatus -> IO ()
forall (m :: * -> *). MonadIO m => Span -> SpanStatus -> m ()
setStatus Span
s (SpanStatus -> IO ()) -> SpanStatus -> IO ()
forall a b. (a -> b) -> a -> b
$ SpanStatus
Ok
            Failure FailureReason
reason -> do
              Span -> SpanStatus -> IO ()
forall (m :: * -> *). MonadIO m => Span -> SpanStatus -> m ()
setStatus Span
s (SpanStatus -> IO ()) -> SpanStatus -> IO ()
forall a b. (a -> b) -> a -> b
$ Text -> SpanStatus
Error (Text -> SpanStatus) -> Text -> SpanStatus
forall a b. (a -> b) -> a -> b
$ String -> Text
T.pack (String -> Text) -> String -> Text
forall a b. (a -> b) -> a -> b
$ FailureReason -> String
forall a. Show a => a -> String
show FailureReason
reason
        Maybe Span
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

      Result -> IO Result
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Result
res
  testOptions :: Tagged (WrappedTest t) [OptionDescription]
testOptions = Tagged t [OptionDescription]
-> Tagged (WrappedTest 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])


-- | Transform a 'TestTree' into one that emits spans around tests and test groups.
instrumentTestTree :: TestTree -> IO TestTree
instrumentTestTree :: TestTree -> IO TestTree
instrumentTestTree TestTree
t = do
  TracerProvider
provider <- IO TracerProvider
forall (m :: * -> *). MonadIO m => m TracerProvider
getGlobalTracerProvider
  let tracer :: Tracer
tracer = TracerProvider -> InstrumentationLibrary -> TracerOptions -> Tracer
makeTracer TracerProvider
provider $Addr#
Int
HashMap Text Attribute
Addr# -> Int -> Text
Text -> Text -> Text -> Attributes -> InstrumentationLibrary
HashMap Text Attribute -> Int -> Int -> Attributes
forall k v. HashMap k v
detectInstrumentationLibrary TracerOptions
tracerOptions
  TestTree -> IO TestTree
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TestTree -> IO TestTree) -> TestTree -> IO TestTree
forall a b. (a -> b) -> a -> b
$ Tracer -> TestTree -> TestTree
instrumentTestTreeWithTracer Tracer
tracer TestTree
t


-- | See 'instrumentTestTree'.
instrumentTestTreeWithTracer :: Tracer -> TestTree -> TestTree
instrumentTestTreeWithTracer :: Tracer -> TestTree -> TestTree
instrumentTestTreeWithTracer Tracer
tracer = Tracer -> IO (Maybe Span) -> TestTree -> TestTree
instrumentTestTree' Tracer
tracer (Maybe Span -> IO (Maybe Span)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Span
forall a. Maybe a
Nothing)


instrumentTestTree'
  :: Tracer
  -> IO (Maybe Span)
  -> TestTree
  -> TestTree
instrumentTestTree' :: Tracer -> IO (Maybe Span) -> TestTree -> TestTree
instrumentTestTree' Tracer
tracer = IO (Maybe Span) -> TestTree -> TestTree
go
  where
    -- See Note [Test parallelism] for why we pass around 'getParentSpan'
    go :: IO (Maybe Span) -> TestTree -> TestTree
    go :: IO (Maybe Span) -> TestTree -> TestTree
go IO (Maybe Span)
getParentSpan = \case
      TestGroup String
name [TestTree]
tests ->
        -- We use 'withResource' to associate the creation and destruction of the
        -- group span with the beginning and end of the group itself. This way
        -- 'tasty' manages the lifetime of the span for us.
        IO Span -> (Span -> IO ()) -> (IO Span -> TestTree) -> TestTree
forall a. IO a -> (a -> IO ()) -> (IO a -> TestTree) -> TestTree
withResource
          (Text -> IO Span
mkSpan (String -> Text
T.pack String
name))
          (\Span
s -> Span -> Maybe Timestamp -> IO ()
forall (m :: * -> *). MonadIO m => Span -> Maybe Timestamp -> m ()
endSpan Span
s Maybe Timestamp
forall a. Maybe a
Nothing)
          ((IO Span -> TestTree) -> TestTree)
-> (IO Span -> TestTree) -> TestTree
forall a b. (a -> b) -> a -> b
$ \IO Span
getGroupSpan ->
            let getParentSpan' :: IO (Maybe Span)
getParentSpan' = Span -> Maybe Span
forall a. a -> Maybe a
Just (Span -> Maybe Span) -> IO Span -> IO (Maybe Span)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Span
getGroupSpan
            in String -> [TestTree] -> TestTree
TestGroup String
name ((TestTree -> TestTree) -> [TestTree] -> [TestTree]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (IO (Maybe Span) -> TestTree -> TestTree
go IO (Maybe Span)
getParentSpan') [TestTree]
tests)
      SingleTest String
name t
t ->
        String -> WrappedTest t -> TestTree
forall t. IsTest t => String -> t -> TestTree
SingleTest String
name (WrappedTest t -> TestTree) -> WrappedTest t -> TestTree
forall a b. (a -> b) -> a -> b
$
          WrappedTest {wrapper :: forall a. IO a -> IO a
wrapper = String -> forall a. IO a -> IO a
withNamedSpan String
name, innerTest :: t
innerTest = t
t}
      WithResource (ResourceSpec IO a
acquire a -> IO ()
release) IO a -> TestTree
f ->
        -- Add spans for resource acquisition and release
        -- Nit: currently we don't create a span for the top-level itself, so
        -- if you acquire outside a test group then the spans will be detached.
        -- We could add a span for the top level, although it's maybe a little odd.
        let newResourceSpec :: ResourceSpec a
newResourceSpec = IO a -> (a -> IO ()) -> ResourceSpec a
forall a. IO a -> (a -> IO ()) -> ResourceSpec a
ResourceSpec (String -> forall a. IO a -> IO a
withNamedSpan String
"acquire" IO a
acquire) ((IO () -> IO ()) -> (a -> IO ()) -> a -> IO ()
forall a b. (a -> b) -> (a -> a) -> a -> b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (String -> forall a. IO a -> IO a
withNamedSpan String
"release") a -> IO ()
release)
        in ResourceSpec a -> (IO a -> TestTree) -> TestTree
forall a. ResourceSpec a -> (IO a -> TestTree) -> TestTree
WithResource ResourceSpec a
newResourceSpec (TestTree -> TestTree
go' (TestTree -> TestTree) -> (IO a -> TestTree) -> IO a -> TestTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO a -> TestTree
f)
      PlusTestOptions OptionSet -> OptionSet
modifier TestTree
t -> (OptionSet -> OptionSet) -> TestTree -> TestTree
PlusTestOptions OptionSet -> OptionSet
modifier (TestTree -> TestTree
go' TestTree
t)
      AskOptions OptionSet -> TestTree
f -> (OptionSet -> TestTree) -> TestTree
AskOptions (TestTree -> TestTree
go' (TestTree -> TestTree)
-> (OptionSet -> TestTree) -> OptionSet -> TestTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OptionSet -> TestTree
f)
      After DependencyType
d Expr
e TestTree
t -> DependencyType -> Expr -> TestTree -> TestTree
After DependencyType
d Expr
e (TestTree -> TestTree) -> TestTree -> TestTree
forall a b. (a -> b) -> a -> b
$ TestTree -> TestTree
go' TestTree
t
      where
        go' :: TestTree -> TestTree
go' = IO (Maybe Span) -> TestTree -> TestTree
go IO (Maybe Span)
getParentSpan
        mkSpan :: Text -> IO Span
mkSpan Text
name = do
          Context
ctx <- IO Context
forall (m :: * -> *). MonadIO m => m Context
getContext
          -- See Note [Test parallelism]
          Maybe Span
parentSpan <- IO (Maybe Span)
getParentSpan
          -- This does not modify the thread-local context, just locally
          Context
ctx' <- case Maybe Span
parentSpan of
            Just Span
s -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> IO Context) -> Context -> IO Context
forall a b. (a -> b) -> a -> b
$ Span -> Context -> Context
insertSpan Span
s Context
ctx
            Maybe Span
Nothing -> Context -> IO Context
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Context
ctx
          Tracer -> Context -> Text -> SpanArguments -> IO Span
forall (m :: * -> *).
(MonadIO m, HasCallStack) =>
Tracer -> Context -> Text -> SpanArguments -> m Span
createSpan Tracer
tracer Context
ctx' Text
name SpanArguments
defaultSpanArguments
        withNamedSpan :: String -> (forall a. IO a -> IO a)
        withNamedSpan :: String -> forall a. IO a -> IO a
withNamedSpan String
name IO a
act = do
          -- See Note [Test parallelism]
          Maybe Span
parentSpan <- IO (Maybe Span)
getParentSpan
          let wrapper :: IO a -> IO a
wrapper = case Maybe Span
parentSpan of
                Just Span
ps -> Span -> forall a. IO a -> IO a
withParentSpan Span
ps
                Maybe Span
Nothing -> IO a -> IO a
forall a. a -> a
id
          IO a -> IO a
forall a. IO a -> IO a
wrapper (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ Tracer -> Text -> SpanArguments -> IO a -> IO a
forall (m :: * -> *) a.
(MonadUnliftIO m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> m a -> m a
inSpan Tracer
tracer (String -> Text
T.pack String
name) SpanArguments
defaultSpanArguments IO a
act


-- Possibly should upstream this to the SDK?

{- | Given a span, produces a wrapper function that sets the given span
as the installed span in the context.
-}
withParentSpan :: Span -> (forall a. IO a -> IO a)
withParentSpan :: Span -> forall a. IO a -> IO a
withParentSpan Span
parentSpan IO a
act =
  IO (Maybe Span, Context)
-> ((Maybe Span, Context) -> IO ())
-> ((Maybe Span, Context) -> IO a)
-> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO (Maybe Span, Context)
setup (Maybe Span, Context) -> IO ()
forall {m :: * -> *} {b}. MonadIO m => (Maybe Span, b) -> m ()
teardown (((Maybe Span, Context) -> IO a) -> IO a)
-> ((Maybe Span, Context) -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \(Maybe Span, Context)
_ -> IO a
act
  where
    setup :: IO (Maybe Span, Context)
setup = do
      Context
ctx <- IO Context
forall (m :: * -> *). MonadIO m => m Context
getContext
      (Context -> Context) -> IO ()
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext (Span -> Context -> Context
insertSpan Span
parentSpan)
      (Maybe Span, Context) -> IO (Maybe Span, Context)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Context -> Maybe Span
lookupSpan Context
ctx, Context
ctx)
    teardown :: (Maybe Span, b) -> m ()
teardown (Maybe Span
originalParentSpan, b
_ctx) = do
      (Context -> Context) -> m ()
forall (m :: * -> *). MonadIO m => (Context -> Context) -> m ()
adjustContext ((Context -> Context) -> m ()) -> (Context -> Context) -> m ()
forall a b. (a -> b) -> a -> b
$ \Context
ctx -> Context -> (Span -> Context) -> Maybe Span -> Context
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Context -> Context
removeSpan Context
ctx) (Span -> Context -> Context
`insertSpan` Context
ctx) Maybe Span
originalParentSpan

{- Note [Test parallelism]
Tasty runs tests in parallel by default, and we don't want to disturb that.
However, that means that any individual test case may be running on a random thread at tasty's
discretion, and so we can't rely on the thread-local context to link up spans.

Our solution is just to track (an action to access) the parent span manually as we traverse
the tree, so we can connect them up manually.
-}