litellm/tests/code_coverage_tests/router_code_coverage.py

import ast
import os


def get_function_names_from_file(file_path):
    """
    Extracts all function names from a given Python file.
    """
    with open(file_path, "r", encoding="utf-8") as file:
        tree = ast.parse(file.read())

    function_names = []

    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            # Top-level functions
            function_names.append(node.name)
        elif isinstance(node, ast.ClassDef):
            # Functions inside classes
            for class_node in node.body:
                if isinstance(class_node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                    function_names.append(class_node.name)

    return function_names


def get_all_functions_called_in_tests(base_dir):
    """
    Returns a set of function names that are called in test functions
    inside the unit test directories,
    specifically in files containing the word 'router'.
    """
    called_functions = set()
    test_dirs = ["unit"]

    for test_dir in test_dirs:
        dir_path = os.path.join(base_dir, test_dir)
        if not os.path.exists(dir_path):
            print(f"Warning: Directory {dir_path} does not exist.")
            continue

        print("dir_path: ", dir_path)
        for root, _, files in os.walk(dir_path):
            for file in files:
                if file.endswith(".py") and ("router" in file.lower() or test_dir == "unit"):
                    print("file: ", file)
                    file_path = os.path.join(root, file)
                    with open(file_path, "r", encoding="utf-8") as f:
                        try:
                            tree = ast.parse(f.read())
                        except SyntaxError:
                            print(f"Warning: Syntax error in file {file_path}")
                            continue
                    if file == "test_router_validate_fallbacks.py":
                        print(f"tree: {tree}")
                    for node in ast.walk(tree):
                        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
                            called_functions.add(node.func.id)
                        elif isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
                            called_functions.add(node.func.attr)

    return called_functions


def get_functions_from_router(file_path):
    """
    Extracts all functions defined in router.py.
    """
    return get_function_names_from_file(file_path)


ignored_function_names = [
    "_acancel_batch",
    "_acreate_batch",
    "_acreate_file",
    "__init__",
    "avector_store_create",  # Tested via proxy vector_store_endpoints (files lack "router" in name)
    "_override_vector_store_methods_for_router",  # No-op placeholder, called during Router init
    "_merge_tools_from_deployment",  # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name)
    "_invalidate_access_groups_cache",  # Tested indirectly via set_model_list, upsert_model etc. (test files lack "router" in name)
    "has_buffered_provider_output",  # Property, so its reads in test_router.py are never an ast.Call
    "chunks",  # Property on FallbackAwareAnthropicMessagesStream, so its reads in tests are never an ast.Call
    "messages",  # Property on FallbackAwareAnthropicMessagesStream, so its reads in tests are never an ast.Call
    "model",  # Property on FallbackAwareAnthropicMessagesStream, so its reads in tests are never an ast.Call
    "_routing_groups",  # Property getter and setter reads are never ast.Call nodes
    "_request_header",  # Tested through Claude Code session routing in test_router.py
    "_claude_code_session_router_cache_key",  # Tested through Claude Code session routing in test_router.py
    "_delete_claude_code_session_router_binding",  # Tested through Redis cleanup failure in test_router.py
    "_resolve_claude_code_session_router",  # Tested through Claude Code session routing in test_router.py
    "_get_claude_code_session_router_binding",  # Tested through the two-worker session routing test in test_router.py
    "_apply_updated_routing_strategy_args",  # Tested via update_settings in test_lowest_latency.py (file lacks "router" in name)
    "_wait_for_scheduler_turn",  # Tested through prioritized acompletion and atext_completion in test_router.py
    "arm_routing_read_prefetch",  # Tested in tests/unit/caching/test_request_redis_batch_pre_call.py (file lacks "router" in name)
    "_configured_model_info",  # Tested through get_configured_service_tiers in test_router.py
    "_routable_deployments",  # Tested through get_configured_service_tiers and get_routable_upstream_model in test_router.py
    "_async_get_available_deployment",  # Body of the `route {model}` phase wrapper, exercised through async_get_available_deployment in test_router.py
    "_async_get_available_deployment_for_pass_through",  # Same, through async_get_available_deployment_for_pass_through in test_router.py
    "_embedding",
    "_aembedding",
    "_stream_pre_content_error",  # Tested through the non-retriable retry error tests in test_router.py
    "_deployment_num_retries",  # Tested through the deployment num_retries mid-stream budget test in test_router.py
    "_request_fallback_list",  # Tested through every mid-stream retry test in test_router.py
    "_request_model_group",  # Tested through test_anthropic_messages_retry_budget_precedence_direct_call
    "_mid_stream_retry_trigger",  # Tested through the retry policy mid-stream budget test in test_router.py
    "_anthropic_messages_group_retry_policy",  # Tested through the retry budget precedence test in test_router.py
    "_anthropic_messages_resolved_retry_policy",  # Tested through the malformed retry policy tests in test_router.py
    "_anthropic_messages_plain_retry_budget",  # Tested through the retry budget precedence test in test_router.py
    "_stream_should_retry",  # Tested through every mid-stream retry test in test_router.py
    "_astream_retry_same_group",  # Tested through the dropped-before-content retry tests in test_router.py
    "_aanthropic_messages_yield_recovered",  # Tested through every mid-stream retry and fallback test in test_router.py
    "_anthropic_messages_policy_retries",  # Tested through the retry budget precedence test in test_router.py
    "_get_wildcard_deployments",  # Tested through the get_model_list_of_routed_group wildcard test in test_router.py
    "_is_fallback_hop",  # Tested through the order fallback hop tests in test_router_order_fallback.py
    "_deployment_that_just_failed",  # Tested through the same-boundary hop test in test_router_order_fallback.py
    "_collect_ptu_share_teams",  # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
    "_resolve_team_ptu_ceiling",  # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py
]


def main():
    router_file = [
        "./litellm/router.py",
        "./litellm/router_utils/batch_utils.py",
        "./litellm/router_utils/pattern_match_deployments.py",
    ]
    # router_file = [
    #     "../../litellm/router.py",
    #     "../../litellm/router_utils/pattern_match_deployments.py",
    #     "../../litellm/router_utils/batch_utils.py",
    # ]  ## LOCAL TESTING
    tests_dir = "./tests/"  # Update this path if your tests directory is located elsewhere
    # tests_dir = "../../tests/"  # LOCAL TESTING

    router_functions = []
    for file in router_file:
        router_functions.extend(get_functions_from_router(file))
    print("router_functions: ", router_functions)
    called_functions_in_tests = get_all_functions_called_in_tests(tests_dir)
    untested_functions = [fn for fn in router_functions if fn not in called_functions_in_tests]

    if untested_functions:
        all_untested_functions = []
        for func in untested_functions:
            if func not in ignored_function_names:
                all_untested_functions.append(func)
        untested_perc = (len(all_untested_functions)) / len(router_functions)
        print("untested_perc: ", untested_perc)
        if untested_perc > 0:
            print("The following functions in router.py are not tested:")
            raise Exception(
                f"{untested_perc * 100:.2f}% of functions in router.py are not tested: {all_untested_functions}"
            )
    else:
        print("All functions in router.py are covered by tests.")


if __name__ == "__main__":
    main()