kca/server/tests/test_telemetry.py

"""Telemetry helper tests.

Provider-dependent tests are skipped unless the optional OpenTelemetry
packages are 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 kca.config import Settings
from kca.models import User
from kca.telemetry import (
    CountingInMemorySessionStore,
    CountingRedisSessionStore,
    _NoOtelLogs,
    _signal_endpoint,
    otel_headers,
    session_tag_hook,
    telemetry_from_settings,
)

from tests.helpers.asynctest import async_test

try:
    from opentelemetry.sdk._logs import LoggingHandler  # noqa: F401
    from opentelemetry.sdk.metrics import MeterProvider  # noqa: F401
    from opentelemetry.sdk.trace import TracerProvider  # noqa: F401

    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

    async def count(self) -> int:
        assert self.value is not None
        return self.value


class _FakeRedis:
    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 _RecordedSpan:
    def __init__(self) -> None:
        self.attributes: dict = {}

    def set_attribute(self, key: str, value) -> None:
        self.attributes[key] = value


class TelemetrySetupTest(unittest.TestCase):
    def test_disabled_returns_none(self) -> None:
        self.assertIsNone(telemetry_from_settings(_settings({}), _FakeSessions()))

    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())


class ProviderSetupTest(unittest.TestCase):
    @unittest.skipUnless(HAVE_OTEL, "OpenTelemetry extra is not installed")
    def test_builds_providers_and_installs_log_handler(self) -> None:
        telemetry = telemetry_from_settings(
            _settings({"OTEL_ENABLED": "1", "OTEL_SERVICE_NAME": "kca-test"}),
            _FakeSessions(),
        )
        assert telemetry is not None
        self.assertIsNotNone(telemetry.meter_provider)
        self.assertIsNotNone(telemetry.tracer_provider)
        self.assertIsNotNone(telemetry.logger_provider)
        self.assertIn(telemetry.log_handler, logging.getLogger().handlers)
        try:
            self.assertIsNotNone(telemetry.mixin)
        finally:
            logging.getLogger().removeHandler(telemetry.log_handler)
            telemetry.log_handler.close()
            telemetry.logger_provider.shutdown()
            telemetry.tracer_provider.shutdown()
            telemetry.meter_provider.shutdown()


class HeaderAndFilterTest(unittest.TestCase):
    def test_otel_headers(self) -> None:
        settings = _settings(
            {"OTEL_EXPORTER_OTLP_HEADERS": "Authorization=Bearer x, X-Tenant=one"}
        )
        self.assertEqual(
            {"Authorization": "Bearer x", "X-Tenant": "one"},
            otel_headers(settings),
        )

    def test_empty_headers(self) -> None:
        self.assertIsNone(otel_headers(_settings({})))

    def test_signal_endpoint(self) -> None:
        self.assertEqual(
            "http://collector:4318/v1/traces",
            _signal_endpoint("http://collector:4318", "traces"),
        )
        self.assertIsNone(_signal_endpoint(None, "metrics"))

    def test_no_otel_logs_filter(self) -> None:
        filter_ = _NoOtelLogs()
        self.assertFalse(filter_.filter(logging.makeLogRecord({"name": "opentelemetry.sdk"})))
        self.assertTrue(filter_.filter(logging.makeLogRecord({"name": "kca.app"})))
        self.assertTrue(filter_.filter(logging.makeLogRecord({"name": "granian.server"})))


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"}))
        self.assertEqual(3, await store.count())
        now[0] += 11
        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", "kca:other"])
        )
        self.assertEqual(2, await store.count())


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("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("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_token_user_cached_on_request(self) -> None:
        span = _RecordedSpan()
        carrier = SimpleNamespace(session=Session("sid", {}))
        setattr(carrier, "_kca_user", User(sub="token-user", display_name="Token"))
        session_tag_hook(span, carrier)
        self.assertEqual("token-user", span.attributes["user.id"])

    def test_ignores_requests_without_session(self) -> None:
        span = _RecordedSpan()
        session_tag_hook(span, SimpleNamespace())
        self.assertEqual({}, span.attributes)


if __name__ == "__main__":
    unittest.main()