{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

-- | Instrumentation for Hspec test suites
module OpenTelemetry.Instrumentation.Hspec (
  wrapSpec,
  wrapExampleInSpan,
  instrumentSpec,
) where

import Control.Monad (void)
import Control.Monad.IO.Class
import Control.Monad.Reader
import qualified Data.List as List
import Data.Text (Text)
import qualified Data.Text as T
import OpenTelemetry.Attributes (Attributes)
import OpenTelemetry.Context
import OpenTelemetry.Context.ThreadLocal (adjustContext, attachContext, getContext)
import OpenTelemetry.Trace.Core
import Test.Hspec.Core.Spec (ActionWith, Item (..), Spec, SpecWith, Tree (..), mapSpecForest, mapSpecItem_)


{- | Creates a wrapper function that you can pass a spec into.

   This function will wrap each @it@ test case with a span with the name of
   the test case.

   The context in which this is called determines the parent span of all of
   the spec items.
-}
wrapSpec :: (MonadIO m) => m (SpecWith a -> SpecWith a)
wrapSpec :: forall (m :: * -> *) a. MonadIO m => m (SpecWith a -> SpecWith a)
wrapSpec = do
  TracerProvider
tp <- m TracerProvider
forall (m :: * -> *). MonadIO m => m TracerProvider
getGlobalTracerProvider
  let tracer :: Tracer
tracer = TracerProvider -> InstrumentationLibrary -> TracerOptions -> Tracer
makeTracer TracerProvider
tp $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
  Context
context <- m Context
forall (m :: * -> *). MonadIO m => m Context
getContext

  -- FIXME: this kind of just dumps everything flat into one span per `it`. We
  -- could possibly do better, e.g. finding the `describe`s and making them into
  -- spans but I am not sure how that would be achieved.
  (SpecWith a -> SpecWith a) -> m (SpecWith a -> SpecWith a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((SpecWith a -> SpecWith a) -> m (SpecWith a -> SpecWith a))
-> (SpecWith a -> SpecWith a) -> m (SpecWith a -> SpecWith a)
forall a b. (a -> b) -> a -> b
$ \SpecWith a
spec -> (Item a -> Item a) -> SpecWith a -> SpecWith a
forall a b. (Item a -> Item b) -> SpecWith a -> SpecWith b
mapSpecItem_ (Tracer -> Context -> Item a -> Item a
forall a. Tracer -> Context -> Item a -> Item a
wrapExampleInSpan Tracer
tracer Context
context) SpecWith a
spec


{- | Wraps one example in a span parented by the specified context, and ensures
   the thread running the spec item will have a context available.
-}
wrapExampleInSpan :: Tracer -> Context -> Item a -> Item a
wrapExampleInSpan :: forall a. Tracer -> Context -> Item a -> Item a
wrapExampleInSpan Tracer
tp Context
traceContext item :: Item a
item@Item {itemExample :: forall a.
Item a
-> Params
-> (ActionWith a -> IO ())
-> ProgressCallback
-> IO Result
itemExample = Params -> (ActionWith a -> IO ()) -> ProgressCallback -> IO Result
ex, itemRequirement :: forall a. Item a -> String
itemRequirement = String
req} =
  Item a
item
    { itemExample = \Params
params ActionWith a -> IO ()
aroundAction ProgressCallback
pcb -> do
        let aroundAction' :: ActionWith a -> IO ()
aroundAction' ActionWith a
a = do
              -- we need to reattach the context, since we are on a forked thread
              IO (Maybe Context) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe Context) -> IO ()) -> IO (Maybe Context) -> IO ()
forall a b. (a -> b) -> a -> b
$ Context -> IO (Maybe Context)
forall (m :: * -> *). MonadIO m => Context -> m (Maybe Context)
attachContext Context
traceContext
              Tracer -> Text -> SpanArguments -> IO () -> IO ()
forall (m :: * -> *) a.
(MonadUnliftIO m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> m a -> m a
inSpan Tracer
tp (String -> Text
T.pack String
req) SpanArguments
defaultSpanArguments (ActionWith a -> IO ()
aroundAction ActionWith a
a)

        Params -> (ActionWith a -> IO ()) -> ProgressCallback -> IO Result
ex Params
params ActionWith a -> IO ()
aroundAction' ProgressCallback
pcb
    }


{- | Instrument each test case. Each 'describe' and friends will add
 a span, and the final test will be in a span described by 'it'.
-}
instrumentSpec :: Tracer -> Context -> SpecWith a -> SpecWith a
instrumentSpec :: forall a. Tracer -> Context -> SpecWith a -> SpecWith a
instrumentSpec Tracer
tracer Context
traceContext SpecWith a
spec = do
  ([SpecTree a] -> [SpecTree a]) -> SpecWith a -> SpecWith a
forall a b r.
([SpecTree a] -> [SpecTree b]) -> SpecM a r -> SpecM b r
mapSpecForest ((SpecTree a -> SpecTree a) -> [SpecTree a] -> [SpecTree a]
forall a b. (a -> b) -> [a] -> [b]
map ([String] -> SpecTree a -> SpecTree a
forall {c} {a}. [String] -> Tree c (Item a) -> Tree c (Item a)
go [])) SpecWith a
spec
  where
    go :: [String] -> Tree c (Item a) -> Tree c (Item a)
go [String]
spans Tree c (Item a)
t = case Tree c (Item a)
t of
      Node String
str [Tree c (Item a)]
rest ->
        String -> [Tree c (Item a)] -> Tree c (Item a)
forall c a. String -> [Tree c a] -> Tree c a
Node String
str ((Tree c (Item a) -> Tree c (Item a))
-> [Tree c (Item a)] -> [Tree c (Item a)]
forall a b. (a -> b) -> [a] -> [b]
map ([String] -> Tree c (Item a) -> Tree c (Item a)
go (String
str String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
spans)) [Tree c (Item a)]
rest)
      NodeWithCleanup Maybe (String, Location)
mloc c
c [Tree c (Item a)]
rest ->
        Maybe (String, Location)
-> c -> [Tree c (Item a)] -> Tree c (Item a)
forall c a. Maybe (String, Location) -> c -> [Tree c a] -> Tree c a
NodeWithCleanup Maybe (String, Location)
mloc c
c ((Tree c (Item a) -> Tree c (Item a))
-> [Tree c (Item a)] -> [Tree c (Item a)]
forall a b. (a -> b) -> [a] -> [b]
map ([String] -> Tree c (Item a) -> Tree c (Item a)
go [String]
spans) [Tree c (Item a)]
rest)
      Leaf Item a
item ->
        Item a -> Tree c (Item a)
forall c a. a -> Tree c a
Leaf
          Item a
item
            { itemExample = \Params
params ActionWith a -> IO ()
aroundAction ProgressCallback
pcb -> do
                let aroundAction' :: ActionWith a -> IO ()
aroundAction' ActionWith a
a = do
                      -- we need to reattach the context, since we are on a forked thread
                      IO (Maybe Context) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (Maybe Context) -> IO ()) -> IO (Maybe Context) -> IO ()
forall a b. (a -> b) -> a -> b
$ Context -> IO (Maybe Context)
forall (m :: * -> *). MonadIO m => Context -> m (Maybe Context)
attachContext Context
traceContext
                      [String] -> IO () -> IO ()
forall {t :: * -> *} {m :: * -> *} {a}.
(Foldable t, MonadUnliftIO m) =>
t String -> m a -> m a
addSpans [String]
spans (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Tracer -> Text -> SpanArguments -> IO () -> IO ()
forall (m :: * -> *) a.
(MonadUnliftIO m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> m a -> m a
inSpan Tracer
tracer (String -> Text
T.pack (Item a -> String
forall a. Item a -> String
itemRequirement Item a
item)) SpanArguments
defaultSpanArguments (ActionWith a -> IO ()
aroundAction ActionWith a
a)

                Item a
-> Params
-> (ActionWith a -> IO ())
-> ProgressCallback
-> IO Result
forall a.
Item a
-> Params
-> (ActionWith a -> IO ())
-> ProgressCallback
-> IO Result
itemExample Item a
item Params
params ActionWith a -> IO ()
aroundAction' ProgressCallback
pcb
            }

    addSpans :: t String -> m a -> m a
addSpans t String
spans m a
k =
      (m a -> String -> m a) -> m a -> t String -> m a
forall b a. (b -> a -> b) -> b -> t a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (\m a
acc String
x -> Tracer -> Text -> SpanArguments -> m a -> m a
forall (m :: * -> *) a.
(MonadUnliftIO m, HasCallStack) =>
Tracer -> Text -> SpanArguments -> m a -> m a
inSpan Tracer
tracer (String -> Text
T.pack String
x) SpanArguments
defaultSpanArguments m a
acc) m a
k t String
spans