PropagationTest.hs

module Kafka.Effectful.OpenTelemetry.PropagationTest (tests) where

import Data.ByteString (ByteString)
import Data.ByteString.Char8 qualified as BSC
import Data.CaseInsensitive qualified as CI
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text.Encoding
import Kafka.Consumer.Types (
    ConsumerRecord (..),
    Offset (..),
    Timestamp (NoTimestamp),
 )
import Kafka.Effectful.OpenTelemetry.Propagation (
    extractTraceContextFromRecord,
    injectTraceContextIntoRecord,
    kafkaHeadersToRequestHeaders,
    kafkaHeadersToTextMap,
    requestHeadersToKafkaHeaders,
    textMapToKafkaHeaders,
 )
import Kafka.Producer.Types (
    ProducePartition (UnassignedPartition),
    ProducerRecord (..),
 )
import Kafka.Types (
    PartitionId (..),
    TopicName (..),
    headersFromList,
    headersToList,
 )
import OpenTelemetry.Context qualified as Context
import OpenTelemetry.Propagator (textMapLookup)
import OpenTelemetry.Trace (initializeGlobalTracerProvider)
import OpenTelemetry.Trace.Core (
    SpanContext (..),
    defaultTraceFlags,
    getSpanContext,
    wrapSpanContext,
 )
import OpenTelemetry.Trace.Id (
    Base (Base16),
    SpanId,
    TraceId,
    baseEncodedToSpanId,
    baseEncodedToTraceId,
    traceIdBaseEncodedText,
 )
import OpenTelemetry.Trace.TraceState qualified as TraceState
import Test.Tasty (TestTree, testGroup, withResource)
import Test.Tasty.HUnit (assertBool, assertEqual, assertFailure, testCase)

{- | Known traceparent value derived from the W3C trace-context
specification\'s example. Trace ID
@0af7651916cd43dd8448eb211c80319c@, span ID
@b7ad6b7169203331@, sampled (flag @01@).
-}
sampleTraceparent :: ByteString
sampleTraceparent =
    "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"

sampleTraceIdHex :: Text
sampleTraceIdHex = "0af7651916cd43dd8448eb211c80319c"

sampleSpanIdHex :: Text
sampleSpanIdHex = "b7ad6b7169203331"

tests :: TestTree
tests =
    -- Initialize the global tracer provider exactly once for the whole
    -- test group. The SDK\'s default propagator pipeline includes
    -- @w3cTraceContextPropagator@, which is what
    -- @injectTraceContextIntoRecord@ \/ @extractTraceContextFromRecord@
    -- consult under the hood.
    withResource initializeGlobalTracerProvider (\_ -> pure ()) $ \_ ->
        testGroup
            "Propagation"
            [ testCase "round-trip kafka headers <-> TextMap" $ do
                let original =
                        headersFromList
                            [ ("traceparent", sampleTraceparent)
                            , ("Custom-Header", "value")
                            ]
                    roundTripped =
                        textMapToKafkaHeaders
                            (kafkaHeadersToTextMap original)
                lookup "traceparent" (headersToList roundTripped)
                    `shouldBeJust` sampleTraceparent
                lookup "Custom-Header" (headersToList roundTripped)
                    `shouldBeJust` "value"
            , testCase "kafkaHeadersToTextMap provides case-insensitive lookup" $ do
                let h = headersFromList [("TraceParent", sampleTraceparent)]
                    tm = kafkaHeadersToTextMap h
                textMapLookup "traceparent" tm
                    `shouldBeJust` Text.Encoding.decodeUtf8 sampleTraceparent
            , testCase "request header compatibility helpers still case-fold" $ do
                let original =
                        headersFromList
                            [ ("traceparent", sampleTraceparent)
                            , ("Custom-Header", "value")
                            ]
                    roundTripped =
                        requestHeadersToKafkaHeaders
                            (kafkaHeadersToRequestHeaders original)
                lookup "traceparent" (headersToList roundTripped)
                    `shouldBeJust` sampleTraceparent
                lookup "custom-header" (headersToList roundTripped)
                    `shouldBeJust` "value"
            , testCase "kafkaHeadersToRequestHeaders preserves traceparent value" $ do
                let h = headersFromList [("traceparent", sampleTraceparent)]
                    rh = kafkaHeadersToRequestHeaders h
                lookup (CI.mk "traceparent") rh
                    `shouldBeJust` sampleTraceparent
            , testCase "injectTraceContextIntoRecord adds W3C traceparent" $ do
                tid <- decodeHexTraceId sampleTraceIdHex
                sid <- decodeHexSpanId sampleSpanIdHex
                let ctx =
                        Context.insertSpan
                            (wrapSpanContext (frozenContextWith tid sid))
                            Context.empty
                injected <-
                    injectTraceContextIntoRecord
                        ctx
                        emptyProducerRecord
                let injectedHeaders = headersToList (prHeaders injected)
                case lookup "traceparent" injectedHeaders of
                    Nothing ->
                        assertFailure
                            "expected the producer record to carry a traceparent header"
                    Just header ->
                        assertBool
                            ( "traceparent did not contain the trace-id "
                                <> Text.unpack sampleTraceIdHex
                                <> ", got: "
                                <> BSC.unpack header
                            )
                            (BSC.pack (Text.unpack sampleTraceIdHex) `BSC.isInfixOf` header)
            , testCase "extractTraceContextFromRecord recovers parent context" $ do
                let cr = consumerRecordWithHeaders sampleTraceparent
                ctx <-
                    extractTraceContextFromRecord cr Context.empty
                case Context.lookupSpan ctx of
                    Nothing ->
                        assertFailure
                            "expected the extracted context to carry a span"
                    Just span_ -> do
                        sc <- getSpanContext span_
                        let recovered = traceIdBaseEncodedText Base16 (traceId sc)
                        assertEqual
                            "recovered trace-id should match the inbound traceparent"
                            sampleTraceIdHex
                            recovered
            ]

shouldBeJust :: (Eq a, Show a) => Maybe a -> a -> IO ()
shouldBeJust Nothing expected =
    assertFailure ("expected Just " <> show expected <> ", got Nothing")
shouldBeJust (Just actual) expected =
    assertEqual "values differ" expected actual

decodeHexTraceId :: Text -> IO TraceId
decodeHexTraceId hex =
    case baseEncodedToTraceId Base16 (BSC.pack (Text.unpack hex)) of
        Right tid -> pure tid
        Left err -> assertFailure ("invalid hex trace id: " <> err) >> error "unreachable"

decodeHexSpanId :: Text -> IO SpanId
decodeHexSpanId hex =
    case baseEncodedToSpanId Base16 (BSC.pack (Text.unpack hex)) of
        Right sid -> pure sid
        Left err -> assertFailure ("invalid hex span id: " <> err) >> error "unreachable"

{- | Build a 'SpanContext' that carries the supplied trace and span
IDs. Trace flags are 'defaultTraceFlags' (unsampled); the W3C
propagator preserves the bytes anyway.
-}
frozenContextWith :: TraceId -> SpanId -> SpanContext
frozenContextWith tid sid =
    SpanContext
        { traceId = tid
        , spanId = sid
        , traceFlags = defaultTraceFlags
        , isRemote = False
        , traceState = TraceState.empty
        }

emptyProducerRecord :: ProducerRecord
emptyProducerRecord =
    ProducerRecord
        { prTopic = TopicName "demo"
        , prPartition = UnassignedPartition
        , prKey = Nothing
        , prValue = Just "value"
        , prHeaders = headersFromList []
        }

consumerRecordWithHeaders ::
    ByteString ->
    ConsumerRecord (Maybe ByteString) (Maybe ByteString)
consumerRecordWithHeaders traceparent =
    ConsumerRecord
        { crTopic = TopicName "demo"
        , crPartition = PartitionId 0
        , crOffset = Offset 0
        , crTimestamp = NoTimestamp
        , crHeaders = headersFromList [("traceparent", traceparent)]
        , crKey = Nothing
        , crValue = Just "value"
        }