test(proxy): make the stagger-offset and linear-dedup guards independent of runner identity and load

tests/unit/proxy/policy_engine/test_attachment_registry.py

@@ -4,8 +4,11 @@
 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.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 @@ def test_same_policy_multiple_attachments_no_duplicates(self):
         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."""

tests/unit/proxy/proxy_server/test_proxy_config.py

@@ -16,7 +16,7 @@
 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)