wai-otel-effectful-1.0.0: test/Main.hs

module Main where

import Data.ByteString (ByteString)
import Data.ByteString qualified as ByteString
import Data.ByteString.Lazy (LazyByteString)
import Effectful
import Effectful.Concurrent (runConcurrent)
import Effectful.Hspec hiding (context)
import Effectful.HttpClient
    ( HttpClient
    , Response
    , httpLbs
    , parseRequest
    , requestHeaders
    , responseHeaders
    , runHttpClient
    )
import Effectful.OpenTelemetry.Protocol.Attributes qualified as Attributes
import Effectful.OpenTelemetry.Tracing (Span (..), Tracing, runInMemoryTracing)
import Effectful.OpenTelemetry.Tracing.Span.Kind qualified as Kind
import Effectful.Wai (Application, responseLBS)
import Effectful.Wai.Handler.Warp (Port, withApplication)
import Effectful.Wai.OpenTelemetry qualified as OpenTelemetry
import Network.HTTP.Types qualified as HTTP
import Prelude

app :: (Tracing :> es) => Application es
app = OpenTelemetry.middleware \_req respond -> respond (responseLBS HTTP.ok200 [] "ok")

get :: (HttpClient :> es) => HTTP.RequestHeaders -> Port -> Eff es (Response LazyByteString)
get headers port = do
    req <- parseRequest ("http://127.0.0.1:" <> show port <> "/hello")
    httpLbs req{requestHeaders = ("Connection", "close") : headers}

inboundTraceparent :: ByteString
inboundTraceparent = "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"

main :: IO ()
main =
    runEff
        . runConcurrent
        . runHttpClient
        . runHspec
        . describe "OpenTelemetry middleware"
        $ do
            it "records a server span" do
                (_, spans) <- runInMemoryTracing $ withApplication (pure app) (get [])
                case spans of
                    [Span{..}] -> do
                        name `shouldBe` "GET /hello"
                        kind `shouldBe` Kind.Server
                        lookup "http.request.method" (Attributes.toList attributes)
                            `shouldBe` Just "GET"
                        lookup "url.path" (Attributes.toList attributes)
                            `shouldBe` Just "/hello"
                    _ -> expectationFailure "expected exactly one span"

            it "returns the trace to a caller that took part in it" do
                (response, _) <-
                    runInMemoryTracing . withApplication (pure app) $
                        get [("traceparent", inboundTraceparent)]
                lookup "traceparent" (responseHeaders response)
                    `shouldSatisfy` maybe
                        False
                        (ByteString.isPrefixOf "00-4bf92f3577b34da6a3ce929d0e0e4736-")

            it "returns no trace to a caller that did not" do
                (response, _) <- runInMemoryTracing $ withApplication (pure app) (get [])
                lookup "traceparent" (responseHeaders response) `shouldBe` Nothing