254 lines
9.1 KiB
Python
254 lines
9.1 KiB
Python
"""Telemetry module unit tests."""
|
||
|
||
from unittest.mock import MagicMock, patch
|
||
|
||
import pytest
|
||
|
||
from agentkit.telemetry.tracer import (
|
||
NoOpSpan,
|
||
NoOpTracer,
|
||
OTelTracer,
|
||
TelemetryConfig,
|
||
get_tracer,
|
||
init_telemetry,
|
||
)
|
||
from agentkit.quality.alignment import AlignmentGuard, AlignmentConfig
|
||
|
||
|
||
# ── NoOpSpan 测试 ──────────────────────────────────────────
|
||
|
||
|
||
class TestNoOpSpan:
|
||
"""NoOpSpan no-op 行为测试"""
|
||
|
||
def test_is_recording_returns_false(self):
|
||
span = NoOpSpan()
|
||
assert span.is_recording() is False
|
||
|
||
def test_set_attribute_does_nothing(self):
|
||
span = NoOpSpan()
|
||
# 不应抛出异常
|
||
span.set_attribute("key", "value")
|
||
|
||
def test_add_event_does_nothing(self):
|
||
span = NoOpSpan()
|
||
span.add_event("event_name", {"attr": "val"})
|
||
|
||
def test_record_exception_does_nothing(self):
|
||
span = NoOpSpan()
|
||
span.record_exception(ValueError("test"))
|
||
|
||
def test_context_manager(self):
|
||
span = NoOpSpan()
|
||
with span as s:
|
||
assert s is span
|
||
|
||
|
||
# ── NoOpTracer 测试 ────────────────────────────────────────
|
||
|
||
|
||
class TestNoOpTracer:
|
||
"""NoOpTracer no-op 行为测试"""
|
||
|
||
def test_start_span_returns_noop_span(self):
|
||
tracer = NoOpTracer()
|
||
span = tracer.start_span("test.span")
|
||
assert isinstance(span, NoOpSpan)
|
||
|
||
def test_start_as_current_span_returns_noop_span(self):
|
||
tracer = NoOpTracer()
|
||
span = tracer.start_as_current_span("test.span")
|
||
assert isinstance(span, NoOpSpan)
|
||
|
||
def test_start_span_with_attributes(self):
|
||
tracer = NoOpTracer()
|
||
span = tracer.start_span("test.span", attributes={"key": "value"})
|
||
assert isinstance(span, NoOpSpan)
|
||
|
||
def test_start_span_as_context_manager(self):
|
||
tracer = NoOpTracer()
|
||
with tracer.start_span("test.span") as span:
|
||
assert isinstance(span, NoOpSpan)
|
||
|
||
|
||
# ── TelemetryConfig 测试 ───────────────────────────────────
|
||
|
||
|
||
class TestTelemetryConfig:
|
||
"""TelemetryConfig 默认值测试"""
|
||
|
||
def test_default_values(self):
|
||
config = TelemetryConfig()
|
||
assert config.enabled is False
|
||
assert config.otlp_endpoint == ""
|
||
assert config.service_name == "agentkit"
|
||
assert config.sample_rate == 1.0
|
||
|
||
def test_custom_values(self):
|
||
config = TelemetryConfig(
|
||
enabled=True,
|
||
otlp_endpoint="http://localhost:4317",
|
||
service_name="my-service",
|
||
sample_rate=0.5,
|
||
)
|
||
assert config.enabled is True
|
||
assert config.otlp_endpoint == "http://localhost:4317"
|
||
assert config.service_name == "my-service"
|
||
assert config.sample_rate == 0.5
|
||
|
||
|
||
# ── init_telemetry 测试 ────────────────────────────────────
|
||
|
||
|
||
class TestInitTelemetry:
|
||
"""init_telemetry 初始化测试"""
|
||
|
||
def test_disabled_returns_noop_tracer(self):
|
||
config = TelemetryConfig(enabled=False)
|
||
init_telemetry(config)
|
||
tracer = get_tracer()
|
||
assert isinstance(tracer, NoOpTracer)
|
||
|
||
def test_enabled_initializes_otel_tracer(self):
|
||
"""U1: enabled=True 初始化 OTelTracer（OTel 现为 required 依赖，不再 ImportError fallback）"""
|
||
config = TelemetryConfig(enabled=True, otlp_endpoint="http://localhost:4317")
|
||
init_telemetry(config)
|
||
tracer = get_tracer()
|
||
# OTel 已安装，应返回 OTelTracer 实例（而非 NoOpTracer）
|
||
assert isinstance(tracer, OTelTracer)
|
||
|
||
def test_init_with_exception_falls_back_to_noop(self):
|
||
"""初始化过程中出现运行时异常时，降级为 NoOpTracer（保留运行时容错）"""
|
||
config = TelemetryConfig(enabled=True, otlp_endpoint="http://localhost:4317")
|
||
# OTLPSpanExporter 在 init_telemetry 函数内部 import，patch 源模块符号
|
||
# opentelemetry.sdk.trace.TracerProvider 是函数内 import 的真实符号
|
||
with patch(
|
||
"opentelemetry.sdk.trace.TracerProvider",
|
||
side_effect=RuntimeError("provider init failed"),
|
||
):
|
||
init_telemetry(config)
|
||
tracer = get_tracer()
|
||
assert isinstance(tracer, NoOpTracer)
|
||
|
||
|
||
# ── get_tracer 测试 ────────────────────────────────────────
|
||
|
||
|
||
class TestGetTracer:
|
||
"""get_tracer 全局实例测试"""
|
||
|
||
def test_returns_global_tracer(self):
|
||
# 确保先重置为 NoOpTracer
|
||
init_telemetry(TelemetryConfig(enabled=False))
|
||
tracer = get_tracer()
|
||
assert isinstance(tracer, NoOpTracer)
|
||
|
||
def test_get_tracer_returns_same_instance(self):
|
||
init_telemetry(TelemetryConfig(enabled=False))
|
||
tracer1 = get_tracer()
|
||
tracer2 = get_tracer()
|
||
assert tracer1 is tracer2
|
||
|
||
|
||
# ── AlignmentGuard span 测试 ──────────────────────────────
|
||
|
||
|
||
class TestAlignmentGuardSpan:
|
||
"""AlignmentGuard 创建 span 并设置属性"""
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_guard_creates_span_with_attributes(self):
|
||
"""对齐检查时创建 span 并设置 guard.passed 和 guard.checked_by 属性"""
|
||
init_telemetry(TelemetryConfig(enabled=False))
|
||
tracer = get_tracer()
|
||
|
||
mock_span = MagicMock()
|
||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||
mock_span.__exit__ = MagicMock(return_value=False)
|
||
mock_span.set_attribute = MagicMock()
|
||
|
||
with patch.object(tracer, "start_span", return_value=mock_span):
|
||
config = AlignmentConfig(constraints=["forbidden_word"])
|
||
guard = AlignmentGuard(config)
|
||
await guard.check_output({"content": "This contains forbidden_word"})
|
||
|
||
tracer.start_span.assert_called_once_with("guard.check")
|
||
call_args_list = [
|
||
(call.args[0], call.args[1]) for call in mock_span.set_attribute.call_args_list
|
||
]
|
||
assert ("guard.constraints_count", 1) in call_args_list
|
||
assert ("guard.passed", False) in call_args_list
|
||
assert ("guard.checked_by", "rule") in call_args_list
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_guard_span_on_pass(self):
|
||
"""对齐检查通过时 span 设置 guard.passed=True"""
|
||
init_telemetry(TelemetryConfig(enabled=False))
|
||
tracer = get_tracer()
|
||
|
||
mock_span = MagicMock()
|
||
mock_span.__enter__ = MagicMock(return_value=mock_span)
|
||
mock_span.__exit__ = MagicMock(return_value=False)
|
||
mock_span.set_attribute = MagicMock()
|
||
|
||
with patch.object(tracer, "start_span", return_value=mock_span):
|
||
config = AlignmentConfig(constraints=["forbidden_word"])
|
||
guard = AlignmentGuard(config)
|
||
await guard.check_output({"content": "This is clean text"})
|
||
|
||
call_args_list = [
|
||
(call.args[0], call.args[1]) for call in mock_span.set_attribute.call_args_list
|
||
]
|
||
assert ("guard.passed", True) in call_args_list
|
||
assert ("guard.checked_by", "rule") in call_args_list
|
||
|
||
|
||
# ── setup_telemetry smoke test ───────────────────────────
|
||
|
||
|
||
class TestSetupTelemetry:
|
||
"""setup_telemetry FastAPI 装配 smoke test（U1）"""
|
||
|
||
def test_disabled_is_noop(self):
|
||
"""enabled=False 时不做任何装配"""
|
||
from fastapi import FastAPI
|
||
|
||
from agentkit.telemetry.setup import setup_telemetry
|
||
|
||
# spec=FastAPI 让 getattr 返回真实的默认值（而非 MagicMock 自动属性）
|
||
app = MagicMock(spec=FastAPI)
|
||
setup_telemetry(app, {"enabled": False})
|
||
# app 不应被 instrument（spec 让未设属性返回 AttributeError→default False）
|
||
assert not getattr(app, "_is_instrumented_by_opentelemetry", False)
|
||
|
||
def test_enabled_instruments_fastapi(self):
|
||
"""enabled=True 时 FastAPIInstrumentor 装配 app（不崩溃即可）"""
|
||
from fastapi import FastAPI
|
||
|
||
from agentkit.telemetry.setup import setup_telemetry
|
||
|
||
app = FastAPI()
|
||
# 使用一个无效 endpoint — exporter 会失败但不阻塞 instrument
|
||
setup_telemetry(
|
||
app,
|
||
{
|
||
"enabled": True,
|
||
"otlp_endpoint": "http://127.0.0.1:4317",
|
||
"service_name": "test-agentkit",
|
||
"export_traces": True,
|
||
"export_metrics": True,
|
||
},
|
||
)
|
||
# FastAPIInstrumentor 装配后会设置此属性
|
||
assert getattr(app, "_is_instrumented_by_opentelemetry", False) is True
|
||
|
||
def test_none_config_is_noop(self):
|
||
"""config=None 不崩溃"""
|
||
from fastapi import FastAPI
|
||
|
||
from agentkit.telemetry.setup import setup_telemetry
|
||
|
||
app = MagicMock(spec=FastAPI)
|
||
setup_telemetry(app, None)
|
||
assert not getattr(app, "_is_instrumented_by_opentelemetry", False)
|