ContextSpec.hs

module Freckle.App.OpenTelemetry.ContextSpec
  ( spec
  ) where

import Prelude

import AppExample
import Blammo.Logging.LogSettings (defaultLogSettings)
import Blammo.Logging.Logger (Logger, newTestLogger)
import Blammo.Logging.Setup (HasLogger (..))
import Control.Lens (lens)
import Control.Monad (void)
import Control.Monad.IO.Class (MonadIO)
import Data.List qualified as List
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as T
import Freckle.App.OpenTelemetry
import Freckle.App.OpenTelemetry.Context
import GHC.Stack (HasCallStack)
import Network.HTTP.Types.Header (Header)
import OpenTelemetry.Exporter.Span (ExportResult (..), SpanExporter (..))
import OpenTelemetry.Registry (registerSpanExporterFactoryIfAbsent)
import OpenTelemetry.Trace.Core qualified as Trace
import Test.Hspec (Spec, describe, it)
import Test.Hspec.Expectations.Lifted
  ( shouldBe
  , shouldNotBe
  , shouldReturn
  )
import Test.Hspec.Expectations.Lifted qualified as Hspec (expectationFailure)

data App = App
  { appLogger :: Logger
  , appTracer :: Tracer
  }

instance HasLogger App where
  loggerL = lens appLogger $ \x y -> x {appLogger = y}

instance HasTracer App where
  tracerL = lens appTracer $ \x y -> x {appTracer = y}

loadApp :: (App -> IO ()) -> IO ()
loadApp f = do
  appLogger <- newTestLogger defaultLogSettings

  -- A TracerProvider with no SpanProcessor never attaches its spans to
  -- context, so getCurrentSpanContext wouldn't see them; register a no-op
  -- exporter so one exists without actually exporting anywhere
  void
    $ registerSpanExporterFactoryIfAbsent "freckle-otel-test-noop"
    $ pure noopSpanExporter

  withTracerProvider $ \tracerProvider -> do
    let appTracer = makeTracer tracerProvider "app" tracerOptions
    f App {..}

noopSpanExporter :: SpanExporter
noopSpanExporter =
  SpanExporter
    { spanExporterExport = const $ pure Success
    , spanExporterShutdown = pure Trace.ShutdownSuccess
    , spanExporterForceFlush = pure Trace.FlushSuccess
    }

spec :: Spec
spec = withApp loadApp $ do
  describe "injectContext" $ do
    it "sets request headers from existing context" $ appExample @App $ do
      inSpan "example" defaultSpanArguments $ do
        spanContext <- assertCurrentSpanContext
        headers <- injectContext ([] :: [Header])

        let headerTraceParent = do
              bs <- List.lookup "traceparent" headers
              fromTraceParent $ T.decodeUtf8 bs

        fmap fst headerTraceParent
          `shouldBe` Just (traceIdToHex $ Trace.traceId spanContext)
        fmap snd headerTraceParent
          `shouldBe` Just (spanIdToHex $ Trace.spanId spanContext)
        List.lookup "tracestate" headers `shouldBe` Just ""

  describe "extractContext" $ do
    it "sets the context from the headers" $ appExample @App $ do
      let
        traceId = "fba7cd19bff3ef866a599d6f6a85baef"
        spanId = "5747fd6b144f4009"

        headers :: [Header]
        headers =
          [ ("traceparent", T.encodeUtf8 $ toTraceParent traceId spanId)
          , ("tracestate", "")
          ]

      extractContext headers

      spanContext <- assertCurrentSpanContext
      traceIdToHex (Trace.traceId spanContext) `shouldBe` traceId
      spanIdToHex (Trace.spanId spanContext) `shouldBe` spanId

    it "does nothing with no headers" $ appExample @App $ do
      getCurrentSpanContext `shouldReturn` Nothing

      extractContext ([] :: [Header])

      getCurrentSpanContext `shouldReturn` Nothing

  describe "processWithContext" $ do
    it "creates a child-span context" $ appExample @App $ do
      inSpan "example" defaultSpanArguments $ do
        outerSpanContext <- assertCurrentSpanContext

        let
          traceId = traceIdToHex $ Trace.traceId outerSpanContext
          spanId = spanIdToHex $ Trace.spanId outerSpanContext

          headers :: [Header]
          headers =
            [ ("traceparent", T.encodeUtf8 $ toTraceParent traceId spanId)
            , ("tracestate", "")
            ]

        processWithContext "loop" defaultSpanArguments headers $ \_ -> do
          innerSpanContext <- assertCurrentSpanContext

          -- Ideally we'd assert the inner span's parent is the outer span, but
          -- the library doesn't expose the necessary functionality. All we can
          -- do is assert it's a fresh span of the same trace.
          Trace.traceId innerSpanContext `shouldBe` Trace.traceId outerSpanContext
          Trace.spanId innerSpanContext `shouldNotBe` Trace.spanId outerSpanContext

    it "sets the current context back into the headers" $ appExample @App $ do
      inSpan "example" defaultSpanArguments $ do
        let
          traceId = "fba7cd19bff3ef866a599d6f6a85baef"
          spanId = "5747fd6b144f4009"

          headers :: [Header]
          headers =
            [ ("traceparent", T.encodeUtf8 $ toTraceParent traceId spanId)
            , ("tracestate", "")
            ]

        processWithContext "loop" defaultSpanArguments headers $ \headers' -> do
          spanContext <- assertCurrentSpanContext

          let headerTraceParent = do
                bs <- List.lookup "traceparent" headers'
                fromTraceParent $ T.decodeUtf8 bs

          fmap fst headerTraceParent `shouldBe` Just traceId
          fmap snd headerTraceParent
            `shouldBe` Just (spanIdToHex $ Trace.spanId spanContext)

    it "sets a fresh context back into the headers" $ appExample @App $ do
      inSpan "example" defaultSpanArguments $ do
        let
          headers :: [Header]
          headers = []

        processWithContext "loop" defaultSpanArguments headers $ \headers' -> do
          spanContext <- assertCurrentSpanContext

          let headerTraceParent = do
                bs <- List.lookup "traceparent" headers'
                fromTraceParent $ T.decodeUtf8 bs

          fmap fst headerTraceParent
            `shouldBe` Just (traceIdToHex $ Trace.traceId spanContext)
          fmap snd headerTraceParent
            `shouldBe` Just (spanIdToHex $ Trace.spanId spanContext)

assertCurrentSpanContext :: (HasCallStack, MonadIO m) => m Trace.SpanContext
assertCurrentSpanContext =
  getCurrentSpanContext >>= \case
    Nothing -> expectationFailure "Expect there to be a SpanContext"
    Just spanContext -> pure spanContext

toTraceParent :: Text -> Text -> Text
toTraceParent traceId spanId = "00-" <> traceId <> "-" <> spanId <> "-00"

fromTraceParent :: Text -> Maybe (Text, Text)
fromTraceParent a = case T.splitOn "-" a of
  [_version, traceId, spanId, _flags] -> Just (traceId, spanId)
  _ -> Nothing

expectationFailure :: (HasCallStack, MonadIO m) => String -> m a
expectationFailure msg = Hspec.expectationFailure msg >> error "unreachable"