test: add static analysis to prevent O(n) linear scans in router

Add AST-based test to detect 'for ... in self.model_list' anti-pattern.
Enforces use of index maps (model_id_to_deployment_index_map and
model_name_to_deployment_indices) for O(1) lookups instead of O(n) iteration.
This commit is contained in:
AlexsanderHamir
2025-10-16 14:41:58 -07:00
parent 254f50a655
commit 35de05dc8b
@@ -1,6 +1,7 @@
import sys
import os
import pytest
import ast
sys.path.insert(
0, os.path.abspath("../..")
@@ -177,3 +178,90 @@ class TestRouterIndexManagement:
# Verify: New entry is added
assert "claude-3" in router.model_name_to_deployment_indices
assert router.model_name_to_deployment_indices["claude-3"] == [0]
def test_no_linear_scans_in_router(self):
"""
Static analysis test to ensure Router doesn't use O(n) linear scans.
Scans router.py for 'in self.model_list' pattern which indicates
inefficient O(n) iteration instead of using index-based O(1) lookups.
Methods should use:
- model_id_to_deployment_index_map for O(1) model_id lookups
- model_name_to_deployment_indices for O(1) + O(k) model_name lookups
"""
# Methods that are allowed to iterate through self.model_list
ALLOWED_METHODS = [
"_get_deployment_by_litellm_model", # Edge case: lookup by litellm_params.model (not indexed)
]
# Get path to router.py
router_file = os.path.join(
os.path.dirname(os.path.dirname(os.path.dirname(__file__))),
"litellm",
"router.py"
)
# Read the file
with open(router_file, 'r') as f:
content = f.read()
# Parse with AST
tree = ast.parse(content)
# Find violations
violations = []
ignore_methods = set(ALLOWED_METHODS)
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
method_name = node.name
# Skip ignored methods
if method_name in ignore_methods:
continue
# Get source for this method
try:
method_source = ast.get_source_segment(content, node)
if not method_source:
continue
# Check for the anti-pattern: "in self.model_list"
# This catches: for x in self.model_list, if x in self.model_list, etc.
if "in self.model_list" in method_source:
# Extract the specific line for better error reporting
lines = method_source.split('\n')
pattern_line = None
for line in lines:
if "in self.model_list" in line:
pattern_line = line.strip()
break
violations.append({
"method": method_name,
"line": node.lineno,
"pattern": pattern_line or "in self.model_list"
})
except Exception:
# Skip if we can't get source segment
pass
# Assert no violations
if violations:
error_msg = "\n".join([
f" - {v['method']}() at line {v['line']}: {v['pattern']}"
for v in violations
])
pytest.fail(
f"\n{'='*70}\n"
f"Found O(n) linear scan pattern in router.py:\n\n"
f"{error_msg}\n\n"
f"These methods should use index maps instead:\n"
f" - model_id_to_deployment_index_map (for model_id lookups)\n"
f" - model_name_to_deployment_indices (for model_name lookups)\n\n"
f"If a method legitimately needs O(n) iteration, add it to\n"
f"ALLOWED_METHODS in this test method.\n"
f"{'='*70}\n"
)