simple-nanoshare/tests/test_tracing_configuration.py

"""Exporter configuration, sampling and ASGI lifecycle regression tests."""
import asyncio

import pytest
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.sdk.trace.sampling import ParentBased, TraceIdRatioBased
from opentelemetry.trace import StatusCode
from quart import Quart, g

from my_modules.observability import RequestTracingMiddleware, register_request_tracing
from quart_common.web.wide_event import register_wide_event_logging

@pytest.mark.parametrize("rate", ["-0.1", "1.1", "nan", "invalid"])
def test_invalid_sample_ratio_rejected(monkeypatch, rate):
  monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://collector.example:4318")
  monkeypatch.setenv("OTEL_TRACES_SAMPLER_ARG", rate)
  with pytest.raises(ValueError):
    register_request_tracing(Quart(__name__))

@pytest.mark.parametrize("service,legacy,expected", [
  (None, None, "nanoshare"),
  (None, "legacy-name", "legacy-name"),
  ("otel-name", "legacy-name", "otel-name"),
])
def test_exporter_configuration_and_shutdown(monkeypatch, service, legacy, expected):
  from opentelemetry.exporter.otlp.proto.http import trace_exporter
  from my_modules import observability

  monkeypatch.setenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", "http://collector.example:4318/v1/traces")
  monkeypatch.setenv("OTEL_TRACES_SAMPLER_ARG", "0.25")
  for key, value in [("OTEL_SERVICE_NAME", service), ("SERVICE_NAME", legacy)]:
    if value is None:
      monkeypatch.delenv(key, raising=False)
    else:
      monkeypatch.setenv(key, value)
  exporter = InMemorySpanExporter()
  timeouts = []

  def make_exporter(**kwargs):
    timeouts.append(kwargs["timeout"])
    return exporter

  providers = []
  monkeypatch.setattr(trace_exporter, "OTLPSpanExporter", make_exporter)
  monkeypatch.setattr(observability.trace, "set_tracer_provider", providers.append)
  app = Quart(__name__)
  provider = register_request_tracing(app)
  try:
    assert providers == [provider]
    assert provider.resource.attributes["service.name"] == expected
    assert provider.sampler._root.rate == 0.25
    assert timeouts == [2]
    shutdowns = []
    monkeypatch.setattr(provider, "shutdown", lambda: shutdowns.append(True))
    asyncio.run(app.after_serving_funcs[0]())
    assert shutdowns == [True]
  finally:
    # Restore the real shutdown before releasing the batch processor thread.
    monkeypatch.undo()
    provider.shutdown()

def test_wide_event_and_response_use_the_active_trace(monkeypatch):
  monkeypatch.setenv("WIDE_EVENT_LOGGING", "true")
  monkeypatch.setenv("WIDE_EVENT_SAMPLE_RATE", "1")
  provider = TracerProvider()
  exporter = InMemorySpanExporter()
  provider.add_span_processor(SimpleSpanProcessor(exporter))
  events = []

  class Logger:
    async def info(self, message):
      events.append(message)

    error = info
    warning = info

  async def run():
    app = Quart(__name__)
    register_wide_event_logging(app, Logger())
    register_request_tracing(app, provider=provider)

    @app.get("/files/<file_id>")
    async def download(file_id):
      return {"trace_id": g.wide_event["trace_id"]}

    result = await app.test_client().get("/files/private-id", headers={"X-Trace-ID": "untrusted-id"})
    trace_id = (await result.get_json())["trace_id"]
    assert result.headers["X-Trace-ID"] == trace_id
    assert len(trace_id) == 32
    assert result.headers["X-Request-ID"]
    return trace_id

  try:
    trace_id = asyncio.run(run())
    assert events
    assert trace_id in events[0]
    span = exporter.get_finished_spans()[0]
    assert span.attributes["request.id"]
    assert span.name == "GET /files/<file_id>"
  finally:
    provider.shutdown()

@pytest.mark.parametrize("parent_sampled,expected_count", [(False, 0), (True, 1)])
def test_upstream_sampling_is_respected(parent_sampled, expected_count):
  provider = TracerProvider(sampler=ParentBased(TraceIdRatioBased(1 if not parent_sampled else 0)))
  exporter = InMemorySpanExporter()
  provider.add_span_processor(SimpleSpanProcessor(exporter))

  async def run():
    app = Quart(__name__)
    register_request_tracing(app, provider=provider)

    @app.get("/")
    async def home():
      return "ok"

    flag = "01" if parent_sampled else "00"
    result = await app.test_client().get("/", headers={"traceparent": f"00-{'1' * 32}-{'2' * 16}-{flag}"})
    assert result.headers["X-Trace-ID"] == "1" * 32

  try:
    asyncio.run(run())
    assert len(exporter.get_finished_spans()) == expected_count
  finally:
    provider.shutdown()

def test_asgi_error_does_not_export_exception_text():
  provider = TracerProvider()
  exporter = InMemorySpanExporter()
  provider.add_span_processor(SimpleSpanProcessor(exporter))

  async def broken(scope, receive, send):
    raise RuntimeError("credential=never-export")

  async def run():
    middleware = RequestTracingMiddleware(broken, provider.get_tracer("test"))
    with pytest.raises(RuntimeError):
      await middleware({"type": "http", "method": "GET"}, None, None)

  try:
    asyncio.run(run())
    span = exporter.get_finished_spans()[0]
    assert span.status.status_code == StatusCode.ERROR
    assert span.attributes["error.type"] == "RuntimeError"
    assert span.events == ()
    assert span.status.description is None
  finally:
    provider.shutdown()

def test_non_http_scopes_pass_through():
  provider = TracerProvider()
  exporter = InMemorySpanExporter()
  provider.add_span_processor(SimpleSpanProcessor(exporter))
  seen = []

  async def app(scope, receive, send):
    seen.append(scope)

  async def run():
    middleware = RequestTracingMiddleware(app, provider.get_tracer("test"))
    for kind in ("websocket", "lifespan"):
      await middleware({"type": kind}, None, None)

  try:
    asyncio.run(run())
    assert [scope["type"] for scope in seen] == ["websocket", "lifespan"]
    assert not exporter.get_finished_spans()
  finally:
    provider.shutdown()