litellm/tests/router_unit_tests/test_router_index_management.py

import sys
import os
import pytest
import ast
import ast

sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system path
from litellm import Router


class TestRouterIndexManagement:
    """Test cases for router index management functions"""

    @pytest.fixture
    def router(self):
        """Create a router instance for testing"""
        return Router(model_list=[])

    def test_deletion_updates_model_name_indices(self, router):
        """Test that deleting a deployment updates model_name_to_deployment_indices correctly"""
        router.model_list = [
            {"model_name": "gpt-3.5", "model_info": {"id": "model-1"}},
            {"model_name": "gpt-5.5", "model_info": {"id": "model-2"}},
            {"model_name": "gpt-5.5", "model_info": {"id": "model-3"}},
            {"model_name": "claude", "model_info": {"id": "model-4"}},
        ]
        router.model_id_to_deployment_index_map = {
            "model-1": 0,
            "model-2": 1,
            "model-3": 2,
            "model-4": 3,
        }
        router.model_name_to_deployment_indices = {
            "gpt-3.5": [0],
            "gpt-5.5": [1, 2],
            "claude": [3],
        }

        # Remove one of the duplicate gpt-5.5 deployments
        router._update_deployment_indices_after_removal(
            model_id="model-2", removal_idx=1
        )

        # Verify indices are shifted correctly
        assert router.model_name_to_deployment_indices["gpt-3.5"] == [0]
        assert router.model_name_to_deployment_indices["gpt-5.5"] == [
            1
        ]  # was [1,2], removed 1, shifted 2->1
        assert router.model_name_to_deployment_indices["claude"] == [
            2
        ]  # was [3], shifted to [2]

        # Remove the last gpt-5.5 deployment
        router._update_deployment_indices_after_removal(
            model_id="model-3", removal_idx=1
        )

        # Verify gpt-5.5 is removed from dict when no deployments remain
        assert "gpt-5.5" not in router.model_name_to_deployment_indices
        assert router.model_name_to_deployment_indices["gpt-3.5"] == [0]
        assert router.model_name_to_deployment_indices["claude"] == [1]

    def test_build_model_id_to_deployment_index_map(self, router):
        """Test _build_model_id_to_deployment_index_map function"""
        model_list = [
            {
                "model_name": "gpt-5-mini",
                "litellm_params": {"model": "gpt-5-mini"},
                "model_info": {"id": "model-1"},
            },
            {
                "model_name": "gpt-5.5",
                "litellm_params": {"model": "gpt-5.5"},
                "model_info": {"id": "model-2"},
            },
        ]

        # Test: Build index from model list
        router._build_model_id_to_deployment_index_map(model_list)

        # Verify: model_list is populated
        assert len(router.model_list) == 2
        # Verify: model_id_to_deployment_index_map is correctly built
        assert router.model_id_to_deployment_index_map["model-1"] == 0
        assert router.model_id_to_deployment_index_map["model-2"] == 1

    def test_add_model_to_list_and_index_map_from_model_info(self, router):
        """Test _add_model_to_list_and_index_map extracting model_id from model_info"""
        # Setup: Empty router
        router.model_list = []
        router.model_id_to_deployment_index_map = {}

        # Test: Add model without explicit model_id
        model = {"model": "test-model", "model_info": {"id": "model-info-id"}}
        router._add_model_to_list_and_index_map(model=model)

        # Verify: Model added to list
        assert len(router.model_list) == 1
        assert router.model_list[0] == model

        # Verify: Index map uses model_info.id
        assert router.model_id_to_deployment_index_map["model-info-id"] == 0

    def test_add_model_to_list_and_index_map_multiple_models(self, router):
        """Test _add_model_to_list_and_index_map with multiple models to verify indexing"""
        # Setup: Empty router
        router.model_list = []
        router.model_id_to_deployment_index_map = {}

        # Test: Add multiple models
        model1 = {"model": "model1", "model_info": {"id": "id-1"}}
        model2 = {"model": "model2", "model_info": {"id": "id-2"}}
        model3 = {"model": "model3", "model_info": {"id": "id-3"}}

        router._add_model_to_list_and_index_map(model=model1, model_id="id-1")
        router._add_model_to_list_and_index_map(model=model2, model_id="id-2")
        router._add_model_to_list_and_index_map(model=model3, model_id="id-3")

        # Verify: All models added to list
        assert len(router.model_list) == 3
        assert router.model_list[0] == model1
        assert router.model_list[1] == model2
        assert router.model_list[2] == model3

        # Verify: Correct indices in map
        assert router.model_id_to_deployment_index_map["id-1"] == 0
        assert router.model_id_to_deployment_index_map["id-2"] == 1
        assert router.model_id_to_deployment_index_map["id-3"] == 2

    def test_update_team_model_index(self, router):
        """Test _update_team_model_index updates team_model_to_deployment_indices."""
        model = {
            "model_name": "team-alias",
            "model_info": {
                "id": "dep-1",
                "team_id": "team-abc",
                "team_public_model_name": "gpt-5.5",
            },
        }
        router._update_team_model_index(model, 0)
        assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0]
        router._update_team_model_index(model, 2)
        assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0, 2]

        router._update_team_model_index(
            {"model_name": "x", "model_info": {"id": "dep-2"}}, 5
        )
        assert router.team_model_to_deployment_indices == {
            ("team-abc", "gpt-5.5"): [0, 2],
        }

    def test_has_model_id(self, router):
        """Test has_model_id function for O(1) membership check"""
        # Setup: Add models to router
        router.model_list = [
            {"model": "test1", "model_info": {"id": "model-1"}},
            {"model": "test2", "model_info": {"id": "model-2"}},
            {"model": "test3", "model_info": {"id": "model-3"}},
        ]
        router.model_id_to_deployment_index_map = {
            "model-1": 0,
            "model-2": 1,
            "model-3": 2,
        }

        # Test: Check existing model IDs
        assert router.has_model_id("model-1") == True
        assert router.has_model_id("model-2") == True
        assert router.has_model_id("model-3") == True

        # Test: Check non-existing model IDs
        assert router.has_model_id("non-existent") == False
        assert router.has_model_id("") == False
        assert router.has_model_id("model-4") == False

        # Test: Empty router
        empty_router = Router(model_list=[])
        assert empty_router.has_model_id("any-id") == False

    def test_build_model_name_index(self, router):
        """Test _build_model_name_index function"""
        model_list = [
            {
                "model_name": "gpt-5-mini",
                "litellm_params": {"model": "gpt-5-mini"},
                "model_info": {"id": "model-1"},
            },
            {
                "model_name": "gpt-5.5",
                "litellm_params": {"model": "gpt-5.5"},
                "model_info": {"id": "model-2"},
            },
            {
                "model_name": "gpt-5.5",  # Duplicate model_name, different deployment
                "litellm_params": {"model": "gpt-5.5"},
                "model_info": {"id": "model-3"},
            },
        ]

        # Test: Build index from model list
        router._build_model_name_index(model_list)

        # Verify: model_name_to_deployment_indices is correctly built
        assert "gpt-5-mini" in router.model_name_to_deployment_indices
        assert "gpt-5.5" in router.model_name_to_deployment_indices

        # Verify: gpt-5-mini has single deployment
        assert router.model_name_to_deployment_indices["gpt-5-mini"] == [0]

        # Verify: gpt-5.5 has multiple deployments
        assert router.model_name_to_deployment_indices["gpt-5.5"] == [1, 2]

        # Test: Rebuild index (should clear and rebuild)
        new_model_list = [
            {
                "model_name": "claude-3",
                "litellm_params": {"model": "claude-3"},
                "model_info": {"id": "model-4"},
            },
        ]
        router._build_model_name_index(new_model_list)

        # Verify: Old entries are cleared
        assert "gpt-5-mini" not in router.model_name_to_deployment_indices
        assert "gpt-5.5" not in router.model_name_to_deployment_indices

        # 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)
            "_finalize_adaptive_router_if_configured",  # Init-time prefix scan for "auto_router/adaptive_router" (no index for prefix match)
        ]

        # 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"
            )

    def test_model_names_is_set(self):
        """Verify that model_names uses a set for O(1) lookups, not a list (O(n))"""
        router = Router(model_list=[])

        assert isinstance(
            router.model_names, set
        ), f"model_names should be a set for O(1) lookups, but got {type(router.model_names)}"