add test in fastapi case

Author: bisgaard-itis

packages/service-library/tests/aiohttp/conftest.py

@@ -3,7 +3,6 @@
 
 
 from collections.abc import Iterator
-from unittest.mock import patch
 
 import pytest
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
@@ -15,7 +14,7 @@
 def mock_otel_collector(mocker: MockerFixture) -> Iterator[InMemorySpanExporter]:
     memory_exporter = InMemorySpanExporter()
     span_processor = SimpleSpanProcessor(memory_exporter)
-    with patch(
+    mocker.patch(
         "servicelib.aiohttp.tracing._create_span_processor", return_value=span_processor
-    ):
-        yield memory_exporter
+    )
+    yield memory_exporter

packages/service-library/tests/fastapi/conftest.py

@@ -3,15 +3,18 @@
 # pylint: disable=unused-variable
 
 import socket
-from collections.abc import AsyncIterator, Callable
+from collections.abc import AsyncIterator, Callable, Iterator
 from typing import cast
 
 import arrow
 import pytest
 from fastapi import APIRouter, FastAPI
 from fastapi.params import Query
 from httpx import ASGITransport, AsyncClient
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from pydantic.types import PositiveFloat
+from pytest_mock import MockerFixture
 
 
 @pytest.fixture
@@ -55,3 +58,13 @@ def go() -> int:
             return cast(int, s.getsockname()[1])
 
     return go
+
+
+@pytest.fixture
+def mock_otel_collector(mocker: MockerFixture) -> Iterator[InMemorySpanExporter]:
+    memory_exporter = InMemorySpanExporter()
+    span_processor = SimpleSpanProcessor(memory_exporter)
+    mocker.patch(
+        "servicelib.fastapi.tracing._create_span_processor", return_value=span_processor
+    )
+    yield memory_exporter

packages/service-library/tests/fastapi/test_tracing.py

@@ -10,10 +10,16 @@
 import pip
 import pytest
 from fastapi import FastAPI
+from fastapi.exceptions import HTTPException
+from fastapi.responses import PlainTextResponse
+from fastapi.testclient import TestClient
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 from pydantic import ValidationError
 from servicelib.fastapi.tracing import (
     get_tracing_instrumentation_lifespan,
+    initialize_fastapi_app_tracing,
 )
+from servicelib.tracing import _OSPARC_TRACE_ID_HEADER
 from settings_library.tracing import TracingSettings
 
 
@@ -167,3 +173,45 @@ async def test_tracing_setup_package_detection(
             service_name="Mock-Openetlemetry-Pytest",
         )(app=mocked_app):
             pass
+
+
+@pytest.mark.parametrize(
+    "tracing_settings_in",
+    [
+        ("http://opentelemetry-collector", 4318),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize(
+    "server_response",
+    [
+        PlainTextResponse("ok"),
+        HTTPException(status_code=400, detail="error"),
+    ],
+)
+async def test_trace_id_in_response_header(
+    mock_otel_collector: InMemorySpanExporter,
+    mocked_app: FastAPI,
+    set_and_clean_settings_env_vars: Callable,
+    tracing_settings_in: Callable,
+    uninstrument_opentelemetry: Iterator[None],
+    server_response: PlainTextResponse | HTTPException,
+) -> None:
+    tracing_settings = TracingSettings()
+
+    @mocked_app.get("/")
+    async def handler():
+        if isinstance(server_response, HTTPException):
+            raise server_response
+        return server_response
+
+    async for _ in get_tracing_instrumentation_lifespan(
+        tracing_settings=tracing_settings,
+        service_name="Mock-OpenTelemetry-Pytest",
+    )(app=mocked_app):
+        initialize_fastapi_app_tracing(mocked_app, add_response_trace_id_header=True)
+        client = TestClient(mocked_app)
+        response = client.get("/")
+        assert _OSPARC_TRACE_ID_HEADER in response.headers
+        trace_id = response.headers[_OSPARC_TRACE_ID_HEADER]
+        assert len(trace_id) == 32  # Ensure trace ID is a 32-character hex string