"""Telemetry unit tests: app gauges, process metrics, log export and counts. The OpenTelemetry packages are an optional extra (``tavolo-app[otel]``), so the provider-dependent tests skip when they are not installed. """ from __future__ import annotations import hashlib import logging import os import sys import unittest from types import SimpleNamespace from typing import Optional from unittest.mock import patch from kaya.session import Session from tavolo.telemetry import ( CountingInMemorySessionStore, CountingRedisSessionStore, TelemetryMixin, _NoOtelLogs, session_tag_hook, telemetry_from_settings, ) from tavolo.config import Settings from tests.helpers import async_test try: from opentelemetry.sdk._logs import LoggingHandler, LoggerProvider from opentelemetry.sdk._logs.export import ( InMemoryLogRecordExporter, SimpleLogRecordProcessor, ) from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import InMemoryMetricReader 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, ) HAVE_OTEL = True except ImportError: # pragma: no cover - installs without the otel extra HAVE_OTEL = False def _settings(env: dict) -> Settings: with patch.dict(os.environ, env, clear=True): return Settings.from_env() class _FakeSessions: def __init__(self, count: int = 0) -> None: self.value: Optional[int] = count self.error: Optional[Exception] = None async def count(self) -> int: if self.error is not None: raise self.error assert self.value is not None return self.value class _FakeGames: def __init__(self, count: int = 0) -> None: self.value: Optional[int] = count async def count_ongoing(self) -> int: assert self.value is not None return self.value class _FakeRedis: """Minimal async Redis double covering the session-count scan.""" def __init__(self, keys: list) -> None: self.keys = keys async def scan_iter(self, match: str = "", count: int = 0): prefix = match.rstrip("*") for key in self.keys: if key.startswith(prefix): yield key class TelemetrySetupTest(unittest.TestCase): def test_disabled_returns_none(self) -> None: self.assertIsNone( telemetry_from_settings( _settings({}), _FakeSessions(), _FakeGames() ) ) def test_missing_otel_dependencies_raise(self) -> None: settings = _settings({"OTEL_ENABLED": "true"}) missing = "opentelemetry.exporter.otlp.proto.http._log_exporter" with patch.dict(sys.modules, {missing: None}): with self.assertRaises(RuntimeError): telemetry_from_settings(settings, _FakeSessions(), _FakeGames()) class SessionCountTest(unittest.TestCase): @async_test async def test_in_memory_counts_live_sessions(self) -> None: now = [1000.0] store = CountingInMemorySessionStore(clock=lambda: now[0]) await store.save("a", Session("a", {"user": "alice"}), max_age=10) await store.save("b", Session("b", {"user": "bob"}), max_age=10) await store.save("c", Session("c", {"user": "carol"})) # no expiry self.assertEqual(3, await store.count()) now[0] += 11 # a and b have expired self.assertEqual(1, await store.count()) await store.delete("c") self.assertEqual(0, await store.count()) @async_test async def test_redis_counts_only_session_keys(self) -> None: store = CountingRedisSessionStore( _FakeRedis(["kaya:session:a", "kaya:session:b", "tavolo:code:X"]) ) self.assertEqual(2, await store.count()) class NoOtelLogsFilterTest(unittest.TestCase): def test_filters_opentelemetry_records_only(self) -> None: filter_ = _NoOtelLogs() self.assertFalse( filter_.filter(logging.makeLogRecord({"name": "opentelemetry.sdk"})) ) self.assertTrue( filter_.filter(logging.makeLogRecord({"name": "tavolo.app"})) ) self.assertTrue( filter_.filter(logging.makeLogRecord({"name": "granian.server"})) ) class _RecordedSpan: def __init__(self) -> None: self.attributes: dict = {} def set_attribute(self, key: str, value) -> None: self.attributes[key] = value class SessionTagHookTest(unittest.TestCase): def test_tags_authenticated_session(self) -> None: span = _RecordedSpan() carrier = SimpleNamespace( session=Session("sid", {"oidc_user": {"sub": "alice"}}) ) session_tag_hook(span, carrier) self.assertEqual( hashlib.sha256(b"sid").hexdigest()[:16], span.attributes["session.id"], ) self.assertEqual("alice", span.attributes["user.id"]) def test_tags_session_without_user(self) -> None: span = _RecordedSpan() session_tag_hook(span, SimpleNamespace(session=Session("sid", {}))) self.assertIn("session.id", span.attributes) self.assertNotIn("user.id", span.attributes) def test_tags_user_when_session_id_not_yet_minted(self) -> None: span = _RecordedSpan() carrier = SimpleNamespace( session=Session(None, {"oidc_user": {"sub": "alice"}}) ) session_tag_hook(span, carrier) self.assertNotIn("session.id", span.attributes) self.assertEqual("alice", span.attributes["user.id"]) def test_missing_session_is_a_noop(self) -> None: span = _RecordedSpan() session_tag_hook(span, SimpleNamespace()) self.assertEqual({}, span.attributes) def test_malformed_user_is_ignored(self) -> None: span = _RecordedSpan() carrier = SimpleNamespace(session=Session("sid", {"oidc_user": "alice"})) session_tag_hook(span, carrier) self.assertIn("session.id", span.attributes) self.assertNotIn("user.id", span.attributes) @unittest.skipUnless(HAVE_OTEL, "tavolo-app[otel] not installed") class SessionTagIntegrationTest(unittest.TestCase): """The hook is wired as kaya-otel's response hook, so it sees the session after the session mixin persisted and minted it.""" @async_test async def test_spans_carry_hashed_session_and_oidc_sub(self) -> None: from httpx import ASGITransport, AsyncClient from kaya.core import HttpContext, KayaApp from kaya.otel import OTelMixin from kaya.session import InMemorySessionStore, SessionMixin span_exporter = InMemorySpanExporter() tracer_provider = TracerProvider( resource=Resource.create({"service.name": "test"}) ) tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) app = KayaApp( mixins=[ OTelMixin( service_name="test", tracer_provider=tracer_provider, meter_provider=MeterProvider( metric_readers=[InMemoryMetricReader()] ), server_response_hook=session_tag_hook, websocket_close_hook=session_tag_hook, ), SessionMixin(InMemorySessionStore()), ] ) @app.GET("/login") async def login(ctx: HttpContext) -> None: session = ctx.session assert session is not None session["oidc_user"] = {"sub": "alice"} await ctx.send_str(200, "ok") @app.GET("/me") async def me(ctx: HttpContext) -> None: await ctx.send_str(200, "ok") transport = ASGITransport(app=app) # type: ignore[arg-type] async with AsyncClient( transport=transport, base_url="http://127.0.0.1" ) as client: first = await client.get("/login") second = await client.get("/me") self.assertEqual(200, first.status_code) self.assertEqual(200, second.status_code) # The cookie value is the session id; only its digest is exported. cookie = first.cookies["session_id"] digest = hashlib.sha256(cookie.encode()).hexdigest()[:16] spans = {span.name: span for span in span_exporter.get_finished_spans()} self.assertEqual({"GET /login", "GET /me"}, set(spans)) for name, span in spans.items(): attributes = span.attributes assert attributes is not None self.assertEqual(digest, attributes["session.id"], name) self.assertEqual("alice", attributes["user.id"], name) def _metric_points(reader) -> dict: """Map metric name -> list of (attributes, value) data points.""" data = reader.get_metrics_data() assert data is not None points = {} for resource_metrics in data.resource_metrics: for scope_metrics in resource_metrics.scope_metrics: for metric in scope_metrics.metrics: points[metric.name] = [ (dict(point.attributes), point.value) for point in metric.data.data_points ] return points @unittest.skipUnless(HAVE_OTEL, "tavolo-app[otel] not installed") class MetricExportTest(unittest.TestCase): # The instrumentation packages are process-wide singletons whose # callbacks cannot be re-registered, so the providers and mixin are # built once for the whole class (as the application does). sessions = _FakeSessions(0) games = _FakeGames(0) @classmethod def setUpClass(cls) -> None: resource = Resource.create({"service.name": "test"}) cls.reader = InMemoryMetricReader() cls.provider = MeterProvider(resource=resource, metric_readers=[cls.reader]) cls.span_exporter = InMemorySpanExporter() cls.tracer_provider = TracerProvider(resource=resource) cls.tracer_provider.add_span_processor( SimpleSpanProcessor(cls.span_exporter) ) cls.mixin = TelemetryMixin( meter_provider=cls.provider, tracer_provider=cls.tracer_provider, logger_provider=None, log_handler=None, session_store=cls.sessions, game_store=cls.games, poll_interval_seconds=1, ) cls.mixin.apply(None) # type: ignore[arg-type] @classmethod def tearDownClass(cls) -> None: cls.mixin.shutdown(None) # type: ignore[arg-type] @async_test async def test_app_gauges_report_polled_counts(self) -> None: self.sessions.value = 3 self.sessions.error = None self.games.value = 2 await self.mixin.refresh_counts() points = _metric_points(self.reader) self.assertEqual([({}, 3)], points["tavolo.active_users"]) self.assertEqual([({}, 2)], points["tavolo.ongoing_games"]) @async_test async def test_failed_refresh_keeps_last_value(self) -> None: self.sessions.value = 5 self.sessions.error = None await self.mixin.refresh_counts() self.sessions.error = RuntimeError("redis down") await self.mixin.refresh_counts() points = _metric_points(self.reader) self.assertEqual([({}, 5)], points["tavolo.active_users"]) @async_test async def test_process_and_gc_metrics_are_exported(self) -> None: await self.mixin.refresh_counts() points = _metric_points(self.reader) for name in ( "process.cpu.time", "process.cpu.utilization", "process.memory.usage", "process.memory.virtual", "cpython.gc.collections", "cpython.gc.collected_objects", "cpython.gc.uncollectable_objects", ): self.assertIn(name, points) self.assertGreater(points["process.memory.usage"][0][1], 0) modes = {attrs["type"] for attrs, _ in points["process.cpu.time"]} self.assertEqual({"user", "system"}, modes) generations = { attrs["cpython.gc.generation"] for attrs, _ in points["cpython.gc.collections"] } self.assertEqual({0, 1, 2}, generations) # Host-level metrics are deliberately not configured. self.assertNotIn("system.cpu.time", points) @async_test async def test_database_clients_are_instrumented(self) -> None: import asyncpg.connection import redis.asyncio self.assertTrue(hasattr(redis.asyncio.Redis.execute_command, "__wrapped__")) self.assertTrue(hasattr(asyncpg.connection.Connection.fetch, "__wrapped__")) # Exercise each wrapper through stubs: the downstream call fails (no # server), but the span is recorded before it does. client = object.__new__(redis.asyncio.Redis) client.connection_pool = SimpleNamespace( connection_kwargs={"host": "127.0.0.1", "port": 6379, "db": 0} ) with self.assertRaises(AttributeError): await client.execute_command("GET", "key") connection = object.__new__(asyncpg.connection.Connection) connection._aborted = True # keep __del__ quiet on the stub with self.assertRaises(Exception): await connection.fetch("SELECT 1") spans = { span.name: span for span in self.span_exporter.get_finished_spans() } self.assertEqual("GET ?", spans["GET"].attributes["db.statement"]) self.assertEqual("redis", spans["GET"].attributes["db.system"]) self.assertEqual("SELECT 1", spans["SELECT"].attributes["db.statement"]) self.assertEqual("postgresql", spans["SELECT"].attributes["db.system"]) @unittest.skipUnless(HAVE_OTEL, "tavolo-app[otel] not installed") class LogExportTest(unittest.TestCase): def test_logs_are_exported_and_otel_records_filtered(self) -> None: exporter = InMemoryLogRecordExporter() resource = Resource.create({"service.name": "test"}) provider = LoggerProvider(resource=resource) provider.add_log_record_processor(SimpleLogRecordProcessor(exporter)) handler = LoggingHandler(level=logging.NOTSET, logger_provider=provider) handler.addFilter(_NoOtelLogs()) root = logging.getLogger() previous_level = root.level root.addHandler(handler) root.setLevel(logging.DEBUG) try: logging.getLogger("tavolo.test").info("exported message") logging.getLogger("opentelemetry.test").info("filtered message") bodies = [r.log_record.body for r in exporter.get_finished_logs()] self.assertIn("exported message", bodies) self.assertNotIn("filtered message", bodies) records = exporter.get_finished_logs() self.assertEqual( "test", records[0].resource.attributes["service.name"] ) finally: root.removeHandler(handler) handler.close() root.setLevel(previous_level) provider.shutdown() if __name__ == "__main__": unittest.main()