diff --git a/tests/unit/proxy/policy_engine/test_attachment_registry.py b/tests/unit/proxy/policy_engine/test_attachment_registry.py index faa8d67fe3a..aefb73eea70 100644 --- a/tests/unit/proxy/policy_engine/test_attachment_registry.py +++ b/tests/unit/proxy/policy_engine/test_attachment_registry.py @@ -4,8 +4,11 @@ Unit tests for AttachmentRegistry - tests policy attachment matching. Tests the main entry point: get_attached_policies() """ -import time +import sys +from collections.abc import Sequence from datetime import datetime, timezone +from types import FrameType +from typing import Final, Protocol from unittest.mock import AsyncMock, MagicMock import pytest @@ -18,6 +21,28 @@ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher from litellm.types.proxy.policy_engine import Policy, PolicyCondition, PolicyGuardrails, PolicyMatchContext +class _TraceFunction(Protocol): + def __call__(self, frame: FrameType, event: str, arg: object) -> "_TraceFunction | None": ... + + +def _lines_run_resolving(registry: AttachmentRegistry, context: PolicyMatchContext) -> tuple[Sequence[str], int]: + lines_run = 0 # rebind-ok: the trace hook counts line events into this closure cell + + def trace(frame: FrameType, event: str, arg: object) -> _TraceFunction | None: + nonlocal lines_run + if event == "line": + lines_run += 1 # rebind-ok: the trace hook counts line events into this closure cell + return trace + + previous_tracer: Final = sys.gettrace() + sys.settrace(trace) + try: + attached: Final = registry.get_attached_policies(context) + finally: + sys.settrace(previous_tracer) + return attached, lines_run + + class TestGetAttachedPolicies: """Test get_attached_policies - the main entry point.""" @@ -267,17 +292,28 @@ class TestGetAttachedPolicies: assert attached.count("multi-policy") == 1 def test_many_distinct_policies_resolve_in_linear_time(self): - policy_count = 20_000 - registry = AttachmentRegistry() - registry.load_attachments([{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)]) - context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") + small_policy_count: Final = 1_000 + small_registry: Final = AttachmentRegistry() + small_registry.load_attachments( + [{"policy": f"policy-{index}", "scope": "*"} for index in range(small_policy_count)] + ) + small_context: Final = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") + small_attached, small_lines = _lines_run_resolving(small_registry, small_context) - started = time.perf_counter() - attached = registry.get_attached_policies(context) - elapsed = time.perf_counter() - started + large_policy_count: Final = 4_000 + large_registry: Final = AttachmentRegistry() + large_registry.load_attachments( + [{"policy": f"policy-{index}", "scope": "*"} for index in range(large_policy_count)] + ) + large_context: Final = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") + large_attached, large_lines = _lines_run_resolving(large_registry, large_context) - assert attached == [f"policy-{index}" for index in range(policy_count)] - assert elapsed < 1.0, f"{policy_count} attachments took {elapsed:.2f}s, dedup is no longer one pass" + assert small_attached == [f"policy-{index}" for index in range(small_policy_count)] + assert large_attached == [f"policy-{index}" for index in range(large_policy_count)] + assert large_lines < 8 * small_lines, ( + f"{large_lines} line events for {large_policy_count} policies vs {small_lines} for {small_policy_count}, " + "dedup is no longer one pass" + ) def test_no_attachments_returns_empty(self): """Test empty attachments returns empty list.""" diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index 47f17a10969..c1f2b9c876f 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -16,7 +16,7 @@ import re from collections.abc import Mapping from contextlib import nullcontext from dataclasses import dataclass -from datetime import datetime +from datetime import datetime, timedelta from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Any, Dict, Final @@ -4278,6 +4278,7 @@ async def test_ProxyConfig__update_general_settings_leaves_first_registration_to async def test_ProxyConfig__update_general_settings_runtime_interval_job_carries_the_stagger_offset(monkeypatch): """Once the scheduler is running the sync owns registration and the job it adds is staggered.""" from apscheduler.schedulers.asyncio import AsyncIOScheduler + from apscheduler.triggers.interval import IntervalTrigger from litellm.proxy.common_utils.scheduled_job_stagger import _OffsetTrigger @@ -4286,12 +4287,17 @@ async def test_ProxyConfig__update_general_settings_runtime_interval_job_carries monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) pc = ProxyConfig() + pc.settings.load_yaml({"scheduled_job_stagger": {"offsets": {"spend_log_cleanup_job": 120}}}) monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings) try: await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"}) jobs = real_scheduler.get_jobs() assert [job.id for job in jobs] == ["spend_log_cleanup_job"] - assert isinstance(jobs[0].trigger, _OffsetTrigger), repr(jobs[0].trigger) + trigger: Final = jobs[0].trigger + assert isinstance(trigger, _OffsetTrigger), repr(trigger) + assert trigger.offset == timedelta(seconds=120) + assert isinstance(trigger.base, IntervalTrigger) + assert trigger.base.interval == timedelta(days=1) finally: real_scheduler.shutdown(wait=False)