tavolo/server/tests/test_telemetry.py

"""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, make_principal

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:default:user": {"sub": "alice"}})
        )
        session_tag_hook(span, carrier)
        self.assertEqual(
            hashlib.sha256(b"sid").hexdigest()[:16],
            span.attributes["session.id"],
        )
        self.assertEqual(make_principal("alice"), span.attributes["user.id"])

    def test_tags_user_of_any_provider(self) -> None:
        span = _RecordedSpan()
        carrier = SimpleNamespace(
            session=Session("sid", {"oidc:google:user": {"sub": "alice"}})
        )
        session_tag_hook(span, carrier)
        self.assertEqual(make_principal("alice", "google"), span.attributes["user.id"])

    def test_ignores_legacy_unprefixed_user_key(self) -> None:
        span = _RecordedSpan()
        carrier = SimpleNamespace(
            session=Session("sid", {"oidc_user": {"sub": "alice"}})
        )
        session_tag_hook(span, carrier)
        self.assertIn("session.id", span.attributes)
        self.assertNotIn("user.id", span.attributes)

    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:default:user": {"sub": "alice"}})
        )
        session_tag_hook(span, carrier)
        self.assertNotIn("session.id", span.attributes)
        self.assertEqual(make_principal("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:default: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:default: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(make_principal("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()