improve test

Author: bisgaard-itis

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

@@ -55,18 +55,24 @@ def set_and_clean_settings_env_vars(
         monkeypatch.setenv(
             "TRACING_OPENTELEMETRY_COLLECTOR_PORT", f"{tracing_settings_in[1]}"
         )
+    sampling_probability_mocked = False
+    if tracing_settings_in[2]:
+        sampling_probability_mocked = True
+        monkeypatch.setenv("TRACING_SAMPLING_PROBABILITY", f"{tracing_settings_in[2]}")
     yield
     if endpoint_mocked:
         monkeypatch.delenv("TRACING_OPENTELEMETRY_COLLECTOR_ENDPOINT")
     if port_mocked:
         monkeypatch.delenv("TRACING_OPENTELEMETRY_COLLECTOR_PORT")
+    if sampling_probability_mocked:
+        monkeypatch.delenv("TRACING_SAMPLING_PROBABILITY")
 
 
 @pytest.mark.parametrize(
     "tracing_settings_in",
     [
-        ("http://opentelemetry-collector", 4318),
-        ("http://opentelemetry-collector", "4318"),
+        ("http://opentelemetry-collector", 4318, 1.0),
+        ("http://opentelemetry-collector", "4318", 1.0),
     ],
     indirect=True,
 )
@@ -91,15 +97,16 @@ async def test_valid_tracing_settings(
 @pytest.mark.parametrize(
     "tracing_settings_in",
     [
-        ("http://opentelemetry-collector", 80),
-        ("http://opentelemetry-collector", 1238712936),
-        ("opentelemetry-collector", 4318),
-        ("httsdasp://ot@##el-collector", 4318),
-        (" !@#$%^&*()[]{};:,<>?\\|`~+=/'\"", 4318),
+        ("http://opentelemetry-collector", 80, 0.5),
+        ("http://opentelemetry-collector", 1238712936, 0.5),
+        ("opentelemetry-collector", 4318, 0.5),
+        ("httsdasp://ot@##el-collector", 4318, 0.5),
+        (" !@#$%^&*()[]{};:,<>?\\|`~+=/'\"", 4318, 0.5),
         # The following exceeds max DNS name length
         (
             "".join(random.choice(string.ascii_letters) for _ in range(300)),
             "1238712936",
+            0.5,
         ),  # noqa: S311
     ],
     indirect=True,
@@ -143,14 +150,14 @@ def manage_package(request):
     "tracing_settings_in, manage_package",
     [
         (
-            ("http://opentelemetry-collector", 4318),
+            ("http://opentelemetry-collector", 4318, 1.0),
             (
                 "opentelemetry-instrumentation-botocore",
                 "opentelemetry.instrumentation.botocore",
             ),
         ),
         (
-            ("http://opentelemetry-collector", "4318"),
+            ("http://opentelemetry-collector", "4318", 1.0),
             (
                 "opentelemetry-instrumentation-aiopg",
                 "opentelemetry.instrumentation.aiopg",
@@ -184,7 +191,7 @@ async def test_tracing_setup_package_detection(
 @pytest.mark.parametrize(
     "tracing_settings_in",
     [
-        ("http://opentelemetry-collector", 4318),
+        ("http://opentelemetry-collector", 4318, 1.0),
     ],
     indirect=True,
 )
@@ -233,7 +240,7 @@ async def handler(handler_data: dict):
 @pytest.mark.parametrize(
     "tracing_settings_in",
     [
-        ("http://opentelemetry-collector", 4318),
+        ("http://opentelemetry-collector", 4318, 1.0),
     ],
     indirect=True,
 )
@@ -284,3 +291,57 @@ async def handler(handler_data: dict):
             for span in spans
             if span.context is not None and span.attributes is not None
         )
+
+
+@pytest.mark.parametrize(
+    "tracing_settings_in",
+    [
+        ("http://opentelemetry-collector", 4318, 0.05),
+    ],
+    indirect=True,
+)
+async def test_tracing_sampling_probability_effective(
+    mock_otel_collector: InMemorySpanExporter,
+    mocked_app: FastAPI,
+    set_and_clean_settings_env_vars: Callable[[], None],
+    tracing_settings_in: Callable[[], dict[str, Any]],
+):
+    """
+    This test checks that the TRACING_SAMPLING_PROBABILITY setting in TracingSettings
+    is effective by sending 1000 requests and verifying that the number of collected traces
+    is close to 0.05 * 1000 (with some tolerance).
+    """
+    n_requests = 1000
+    tolerance_probability = 0.5
+
+    tracing_settings = TracingSettings()
+
+    async def handler():
+        return PlainTextResponse("ok")
+
+    mocked_app.get("/")(handler)
+
+    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)
+        for _ in range(n_requests):
+            client.get("/")
+        spans = mock_otel_collector.get_finished_spans()
+        trace_ids = set()
+        for span in spans:
+            if span.context is not None:
+                trace_ids.add(span.context.trace_id)
+        num_traces = len(trace_ids)
+        expected_num_traces = int(
+            tracing_settings.TRACING_SAMPLING_PROBABILITY * n_requests
+        )
+        # Allow a 50% tolerance due to randomness
+        tolerance = int(tolerance_probability * expected_num_traces)
+        assert (
+            expected_num_traces - tolerance
+            <= num_traces
+            <= expected_num_traces + tolerance
+        ), f"Expected roughly {expected_num_traces} distinct trace ids, got {num_traces}"