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