Add optional OpenTelemetry traces, metrics and OTLP logs
This commit is contained in:
@@ -0,0 +1,200 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user