simple-nanoshare/tests/test_request_tracing.py

import asyncio

import httpx
from opentelemetry.sdk.resources import Resource
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.trace import SpanKind, StatusCode
from quart import Quart, g

from my_helpers.db.convex.ConvexJob import ConvexJob
from my_helpers.db.convex.ConvexWorkerPool import ConvexWorkerPool
from my_helpers.db.convex import telemetry
from my_modules.observability import register_request_tracing

def tracing_provider(monkeypatch):
  exporter = InMemorySpanExporter()
  provider = TracerProvider(resource=Resource.create({"service.name": "test-api"}))
  provider.add_span_processor(SimpleSpanProcessor(exporter))
  monkeypatch.setattr(telemetry, "_TRACER", provider.get_tracer("convex-test"))
  return provider, exporter

def response(value=None, *, error=None):
  body = {"status": "success", "value": value}
  if error:
    body = {"status": "error", "errorMessage": error}
  return httpx.Response(200, json=body, request=httpx.Request("POST", "https://convex.example.com/api/query"))

def test_tracing_is_disabled_without_endpoint(monkeypatch):
  monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False)
  monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", raising=False)
  app = Quart(__name__)
  original = app.asgi_app
  assert register_request_tracing(app) is None
  assert app.asgi_app == original

def test_request_to_worker_to_convex_keeps_trace_and_payload_private(monkeypatch):
  provider, exporter = tracing_provider(monkeypatch)
  captured = []
  events = []

  def post(url, *, json, headers, timeout):
    captured.append({"json": json, "headers": headers})
    return response({"ok": True})

  monkeypatch.setattr(httpx, "post", post)

  async def run():
    app = Quart(__name__)
    pool = ConvexWorkerPool("https://convex.example.com", workers=1)

    @app.before_request
    async def wide_event():
      g.wide_event = {"request_id": "req-test"}

    register_request_tracing(app, provider=provider)

    @app.get("/items/<item_id>")
    async def items(item_id):
      await pool.submit(ConvexJob("job", "query", "items", "get", {"secret": "never-export"}), timeout=2)
      events.append(dict(g.wide_event))
      return {"ok": True}

    await pool.start(health_log_interval=0)
    try:
      result = await app.test_client().get(
        "/items/private-value?token=never-export",
        headers={"traceparent": "00-11111111111111111111111111111111-2222222222222222-01"},
      )
      assert result.status_code == 200
      assert result.headers["x-trace-id"] == "1" * 32
    finally:
      await pool.stop()

  asyncio.run(run())
  spans = exporter.get_finished_spans()
  server = next(s for s in spans if s.kind == SpanKind.SERVER)
  client = next(s for s in spans if s.kind == SpanKind.CLIENT)
  assert server.name == "GET /items/<item_id>"
  assert server.parent.span_id == int("2" * 16, 16)
  assert client.parent.span_id == server.context.span_id
  assert client.context.trace_id == server.context.trace_id == int("1" * 32, 16)
  native_id = f"otel-{'1' * 32}-{client.context.span_id:016x}"
  assert captured[0]["headers"]["convex-request-id"] == native_id
  assert captured[0]["headers"]["traceparent"] == f"00-{'1' * 32}-{client.context.span_id:016x}-01"
  assert client.attributes["convex.request_id"] == native_id
  assert captured[0]["json"]["args"] == {"secret": "never-export"}
  assert events[0]["trace_id"] == "1" * 32
  exported = " ".join(str((s.name, dict(s.attributes), s.events, s.status.description)) for s in spans)
  assert "never-export" not in exported
  assert server.attributes["url.path"] == "/items/private-value"
  assert "token=" not in exported
  provider.shutdown()

def test_concurrent_requests_do_not_inherit_worker_or_other_request_context(monkeypatch):
  provider, exporter = tracing_provider(monkeypatch)
  captured = []

  def post(url, *, json, headers, timeout):
    captured.append((json["args"]["request"], headers["convex-request-id"]))
    return response(None)

  monkeypatch.setattr(httpx, "post", post)

  async def run():
    app = Quart(__name__)
    register_request_tracing(app, provider=provider)
    pool = ConvexWorkerPool("https://convex.example.com", workers=1)

    @app.get("/query/<key>")
    async def query(key):
      await pool.submit(ConvexJob(key, "query", "items", "get", {"request": key}), timeout=2)
      return "ok"

    await pool.start(health_log_interval=0)
    try:
      await asyncio.gather(*[
        app.test_client().get(f"/query/{key}", headers={"traceparent": f"00-{key * 32}-{'3' * 16}-01"})
        for key in ["a", "b"]
      ])
    finally:
      await pool.stop()

  asyncio.run(run())
  assert len(captured) == 2
  for key, native_id in captured:
    assert native_id.startswith(f"otel-{key * 32}-")
  servers = {s.context.trace_id: s for s in exporter.get_finished_spans() if s.kind == SpanKind.SERVER}
  for span in exporter.get_finished_spans():
    if span.kind == SpanKind.CLIENT:
      assert span.parent.span_id == servers[span.context.trace_id].context.span_id
  provider.shutdown()

def test_retry_attempts_get_unique_ids_and_do_not_export_exception_text(monkeypatch):
  provider, exporter = tracing_provider(monkeypatch)
  calls = []

  def post(url, *, json, headers, timeout):
    calls.append(headers["convex-request-id"])
    if len(calls) == 1:
      return response(error="credential=never-export")
    return response("ok")

  monkeypatch.setattr(httpx, "post", post)

  async def run():
    pool = ConvexWorkerPool("https://convex.example.com", max_retries=1, retry_backoff_secs=0)
    await pool.start(health_log_interval=0)
    try:
      with provider.get_tracer("test").start_as_current_span("request"):
        assert await pool.submit(ConvexJob("id", "query", "items", "get", {}), timeout=2) == "ok"
    finally:
      await pool.stop()

  asyncio.run(run())
  clients = [s for s in exporter.get_finished_spans() if s.kind == SpanKind.CLIENT]
  assert len(clients) == 2
  assert calls[0] != calls[1]
  assert clients[0].context.trace_id == clients[1].context.trace_id
  assert clients[0].status.status_code == StatusCode.ERROR
  assert clients[0].attributes["error.type"] == "ConvexExecutionError"
  assert clients[0].events == ()
  assert clients[0].status.description is None
  provider.shutdown()

def test_unmatched_request_records_path_without_query_or_headers(monkeypatch):
  provider, exporter = tracing_provider(monkeypatch)

  async def run():
    app = Quart(__name__)
    register_request_tracing(app, provider=provider)
    result = await app.test_client().get(
      "/missing/%E2%9C%93?token=never-export",
      headers={"Authorization": "Bearer never-export", "Cookie": "secret=never-export"},
    )
    assert result.status_code == 404

  try:
    asyncio.run(run())
    span = exporter.get_finished_spans()[0]
    assert span.attributes["url.path"] == "/missing/✓"
    assert span.attributes["http.response.status_code"] == 404
    assert "http.route" not in span.attributes
    assert span.name == "HTTP GET"
    assert "never-export" not in str(dict(span.attributes))
    assert "url.full" not in span.attributes
    assert "url.query" not in span.attributes
  finally:
    provider.shutdown()

def test_invalid_incoming_traceparent_is_not_used(monkeypatch):
  provider, exporter = tracing_provider(monkeypatch)

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

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

    result = await app.test_client().get("/", headers={"traceparent": f"00-{'0' * 32}-{'0' * 16}-01"})
    assert int(result.headers["x-trace-id"], 16) != 0

  asyncio.run(run())
  assert exporter.get_finished_spans()[0].parent is None
  provider.shutdown()