SemanticTest.hs

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

import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.HashMap.Lazy qualified as HashMap
import Data.Int (Int64)
import Data.Text (Text)
import Kafka.Consumer.ConsumerProperties qualified as ConsumerProperties
import Kafka.Consumer.Types (
    ConsumerGroupId (..),
    ConsumerRecord (..),
    Offset (..),
    Timestamp (NoTimestamp),
 )
import Kafka.Effectful.OpenTelemetry.Semantic (
    consumerRecordAttributes,
    consumerRecordAttributesWith,
    producerRecordAttributes,
    producerRecordAttributesWith,
 )
import Kafka.Producer.Types (
    ProducePartition (SpecifiedPartition, UnassignedPartition),
    ProducerRecord (..),
 )
import Kafka.Types (
    ClientId (..),
    PartitionId (..),
    TopicName (..),
    headersFromList,
 )
import OpenTelemetry.Attributes.Attribute (
    Attribute (AttributeValue),
    PrimitiveAttribute (IntAttribute, TextAttribute),
 )
import OpenTelemetry.SemanticsConfig (StabilityOpt (..))
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertFailure, testCase, (@?=))

tests :: TestTree
tests =
    testGroup
        "Semantic"
        [ producerTests
        , consumerTests
        ]

producerTests :: TestTree
producerTests =
    testGroup
        "producerRecordAttributes"
        [ testCase "legacy helper keeps old operation key" $ do
            let attrs = producerRecordAttributes (sampleProducerRecord (Just "k1"))
            HashMap.lookup "messaging.system" attrs @?= Just (textAttr "kafka")
            HashMap.lookup "messaging.destination.name" attrs @?= Just (textAttr "orders")
            HashMap.lookup "messaging.operation" attrs @?= Just (textAttr "send")
            assertAbsent "messaging.operation.name" attrs
            assertAbsent "messaging.operation.type" attrs
        , testCase "stable mode uses stable operation keys" $ do
            let attrs = producerRecordAttributesWith Stable (sampleProducerRecord (Just "k1"))
            HashMap.lookup "messaging.operation.name" attrs @?= Just (textAttr "send")
            HashMap.lookup "messaging.operation.type" attrs @?= Just (textAttr "send")
            assertAbsent "messaging.operation" attrs
        , testCase "duplicate mode emits old and stable operation keys" $ do
            let attrs = producerRecordAttributesWith StableAndOld (sampleProducerRecord (Just "k1"))
            HashMap.lookup "messaging.operation" attrs @?= Just (textAttr "send")
            HashMap.lookup "messaging.operation.name" attrs @?= Just (textAttr "send")
            HashMap.lookup "messaging.operation.type" attrs @?= Just (textAttr "send")
        , testCase "adds partition, key, and body size when available" $ do
            let attrs = producerRecordAttributesWith Old (sampleProducerRecord (Just "hello"))
            HashMap.lookup "messaging.kafka.destination.partition" attrs
                @?= Just (intAttr 42)
            HashMap.lookup "messaging.kafka.message.key" attrs
                @?= Just (textAttr "hello")
            HashMap.lookup "messaging.message.body.size" attrs
                @?= Just (intAttr 5)
        , testCase "omits partition for UnassignedPartition" $ do
            let attrs =
                    producerRecordAttributesWith
                        Old
                        (sampleProducerRecord (Just "k1"))
                            { prPartition = UnassignedPartition
                            }
            assertAbsent "messaging.kafka.destination.partition" attrs
        , testCase "omits invalid or absent keys" $ do
            let invalidAttrs =
                    producerRecordAttributesWith
                        Old
                        (sampleProducerRecord (Just (BS.pack [0xC3, 0x28])))
                missingAttrs =
                    producerRecordAttributesWith Old (sampleProducerRecord Nothing)
            assertAbsent "messaging.kafka.message.key" invalidAttrs
            assertAbsent "messaging.kafka.message.key" missingAttrs
        ]

consumerTests :: TestTree
consumerTests =
    testGroup
        "consumerRecordAttributes"
        [ testCase "legacy helper keeps old operation key" $ do
            let attrs = consumerRecordAttributes (sampleConsumerRecord (Just "k1"))
            HashMap.lookup "messaging.system" attrs @?= Just (textAttr "kafka")
            HashMap.lookup "messaging.destination.name" attrs @?= Just (textAttr "orders")
            HashMap.lookup "messaging.operation" attrs @?= Just (textAttr "process")
            assertAbsent "messaging.operation.name" attrs
            assertAbsent "messaging.operation.type" attrs
        , testCase "old mode uses legacy consumer group key" $ do
            let attrs = consumerRecordAttributesWith Old sampleConsumerProps (sampleConsumerRecord (Just "k1"))
            HashMap.lookup "messaging.kafka.consumer.group" attrs
                @?= Just (textAttr "orders-group")
            assertAbsent "messaging.consumer.group.name" attrs
        , testCase "stable mode uses stable operation and consumer group keys" $ do
            let attrs = consumerRecordAttributesWith Stable sampleConsumerProps (sampleConsumerRecord (Just "k1"))
            HashMap.lookup "messaging.operation.name" attrs @?= Just (textAttr "process")
            HashMap.lookup "messaging.operation.type" attrs @?= Just (textAttr "process")
            HashMap.lookup "messaging.consumer.group.name" attrs
                @?= Just (textAttr "orders-group")
            assertAbsent "messaging.operation" attrs
            assertAbsent "messaging.kafka.consumer.group" attrs
        , testCase "duplicate mode emits old and stable operation and group keys" $ do
            let attrs = consumerRecordAttributesWith StableAndOld sampleConsumerProps (sampleConsumerRecord (Just "k1"))
            HashMap.lookup "messaging.operation" attrs @?= Just (textAttr "process")
            HashMap.lookup "messaging.operation.name" attrs @?= Just (textAttr "process")
            HashMap.lookup "messaging.operation.type" attrs @?= Just (textAttr "process")
            HashMap.lookup "messaging.kafka.consumer.group" attrs
                @?= Just (textAttr "orders-group")
            HashMap.lookup "messaging.consumer.group.name" attrs
                @?= Just (textAttr "orders-group")
        , testCase "adds client id, partition, offset, key, and body size" $ do
            let attrs = consumerRecordAttributesWith Old sampleConsumerProps (sampleConsumerRecord (Just "hello"))
            HashMap.lookup "messaging.client.id" attrs
                @?= Just (textAttr "orders-client")
            HashMap.lookup "messaging.kafka.destination.partition" attrs
                @?= Just (intAttr 7)
            HashMap.lookup "messaging.kafka.message.offset" attrs
                @?= Just (intAttr 42)
            HashMap.lookup "messaging.kafka.message.key" attrs
                @?= Just (textAttr "hello")
            HashMap.lookup "messaging.message.body.size" attrs
                @?= Just (intAttr 5)
        , testCase "omits invalid or absent keys" $ do
            let invalidAttrs =
                    consumerRecordAttributesWith
                        Old
                        sampleConsumerProps
                        (sampleConsumerRecord (Just (BS.pack [0xC3, 0x28])))
                missingAttrs =
                    consumerRecordAttributesWith Old sampleConsumerProps (sampleConsumerRecord Nothing)
            assertAbsent "messaging.kafka.message.key" invalidAttrs
            assertAbsent "messaging.kafka.message.key" missingAttrs
        ]

sampleProducerRecord :: Maybe ByteString -> ProducerRecord
sampleProducerRecord key =
    ProducerRecord
        { prTopic = TopicName "orders"
        , prPartition = SpecifiedPartition 42
        , prKey = key
        , prValue = Just "value"
        , prHeaders = headersFromList []
        }

sampleConsumerRecord ::
    Maybe ByteString ->
    ConsumerRecord (Maybe ByteString) (Maybe ByteString)
sampleConsumerRecord key =
    ConsumerRecord
        { crTopic = TopicName "orders"
        , crPartition = PartitionId 7
        , crOffset = Offset 42
        , crTimestamp = NoTimestamp
        , crHeaders = headersFromList []
        , crKey = key
        , crValue = Just "value"
        }

sampleConsumerProps :: ConsumerProperties.ConsumerProperties
sampleConsumerProps =
    ConsumerProperties.groupId (ConsumerGroupId "orders-group")
        <> ConsumerProperties.clientId (ClientId "orders-client")

textAttr :: Text -> Attribute
textAttr = AttributeValue . TextAttribute

intAttr :: Int64 -> Attribute
intAttr = AttributeValue . IntAttribute

assertAbsent :: Text -> HashMap.HashMap Text Attribute -> IO ()
assertAbsent key attrs =
    case HashMap.lookup key attrs of
        Nothing -> pure ()
        Just attr ->
            assertFailure $
                "expected "
                    <> show key
                    <> " to be absent, got "
                    <> show attr