255 lines
9.2 KiB
Python
255 lines
9.2 KiB
Python
"""Telemetry module unit tests."""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from agentkit.telemetry.tracer import (
|
|
NoOpSpan,
|
|
NoOpTracer,
|
|
OTelSpan,
|
|
OTelTracer,
|
|
TelemetryConfig,
|
|
get_tracer,
|
|
init_telemetry,
|
|
)
|
|
from agentkit.chat.skill_routing import CostAwareRouter, SkillRoutingResult
|
|
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_missing_opentelemetry_falls_back_to_noop(self):
|
|
"""当 opentelemetry 包未安装时，优雅降级为 NoOpTracer"""
|
|
config = TelemetryConfig(enabled=True, otlp_endpoint="http://localhost:4317")
|
|
# 即使 opentelemetry 未安装，也不应抛出异常
|
|
init_telemetry(config)
|
|
tracer = get_tracer()
|
|
# 在没有 opentelemetry 的环境中应为 NoOpTracer
|
|
assert isinstance(tracer, (NoOpTracer, OTelTracer))
|
|
|
|
def test_init_with_exception_falls_back_to_noop(self):
|
|
"""初始化过程中出现异常时，降级为 NoOpTracer"""
|
|
config = TelemetryConfig(enabled=True, otlp_endpoint="http://localhost:4317")
|
|
with patch(
|
|
"agentkit.telemetry.tracer.init_telemetry",
|
|
side_effect=Exception("init error"),
|
|
) as mock_init:
|
|
# init_telemetry 被模拟为抛异常，但实际调用的是原始函数
|
|
# 这里验证的是 init_telemetry 本身不会崩溃
|
|
pass
|
|
# 直接调用，不应崩溃
|
|
init_telemetry(config)
|
|
tracer = get_tracer()
|
|
assert isinstance(tracer, (NoOpTracer, OTelTracer))
|
|
|
|
|
|
# ── 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
|
|
|
|
|
|
# ── CostAwareRouter span 测试 ──────────────────────────────
|
|
|
|
|
|
class TestCostAwareRouterSpan:
|
|
"""CostAwareRouter 创建 span 并设置属性"""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_router_creates_span_with_attributes(self):
|
|
"""路由时创建 span 并设置 route.layer 和 route.target 属性"""
|
|
init_telemetry(TelemetryConfig(enabled=False))
|
|
tracer = get_tracer()
|
|
|
|
# 用 mock 替换 start_span 以验证调用
|
|
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):
|
|
router = CostAwareRouter()
|
|
result = await router.route(
|
|
content="你好",
|
|
skill_registry=MagicMock(),
|
|
intent_router=MagicMock(),
|
|
default_tools=[],
|
|
default_system_prompt="You are helpful.",
|
|
)
|
|
|
|
tracer.start_span.assert_called_once_with("router.route")
|
|
# 验证 span 设置了 input.length 属性
|
|
mock_span.set_attribute.assert_any_call("input.length", len("你好"))
|
|
# 验证 span 设置了 route.layer 和 route.target
|
|
call_args_list = [
|
|
(call.args[0], call.args[1])
|
|
for call in mock_span.set_attribute.call_args_list
|
|
]
|
|
assert ("route.layer", "greeting") in call_args_list
|
|
assert ("route.target", "default") in call_args_list
|
|
|
|
|
|
# ── 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)
|
|
result = 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)
|
|
result = 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
|