litellm/tests/unit/proxy/test_caching_routes.py

import json

import pytest
from fastapi.testclient import TestClient
from pydantic import ValidationError



import litellm
from litellm.caching import RedisCache
from litellm.caching.caching import Cache
from litellm.proxy.caching_routes import _extract_cache_params
from litellm.proxy.proxy_server import app

client = TestClient(app)


# Mock successful Redis connection
@pytest.fixture
def mock_redis_success(mocker):
    async def mock_ping():
        return True

    async def mock_add_cache(*args, **kwargs):
        return None

    mock_cache = mocker.MagicMock()
    mock_cache.type = "redis"
    mock_cache.ping = mock_ping
    mock_cache.async_add_cache = mock_add_cache
    mock_cache.cache = RedisCache(
        host="localhost",
        port=6379,
        password="hello",
    )

    mocker.patch.object(litellm, "cache", mock_cache)
    return mock_cache


# Mock failed Redis connection
@pytest.fixture
def mock_redis_failure(mocker):
    async def mock_ping():
        raise Exception("invalid username-password pair")

    mock_cache = mocker.MagicMock()
    mock_cache.type = "redis"
    mock_cache.ping = mock_ping

    mocker.patch.object(litellm, "cache", mock_cache)
    return mock_cache


def test_cache_ping_success(mock_redis_success):
    """Test successful cache ping with regular response"""
    response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-9876"})
    assert response.status_code == 200

    data = response.json()
    assert data["status"] == "healthy"
    assert data["cache_type"] == "redis"
    assert data["ping_response"] is True
    assert data["set_cache_response"] == "success"


def test_cache_ping_with_complex_objects(mock_redis_success, mocker):
    """Test cache ping with non-standard serializable objects"""

    # Mock complex objects in the cache parameters
    class ComplexObject:
        def __str__(self):
            return "complex_object"

    mock_redis_success.cache.complex_attr = ComplexObject()
    mock_redis_success.cache.datetime_attr = mocker.MagicMock()

    response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-9876"})
    assert response.status_code == 200

    # Verify response is JSON serializable
    data = response.json()
    print("data=", json.dumps(data, indent=4))
    assert data["status"] == "healthy"
    assert "litellm_cache_params" in data

    # Verify complex objects were converted to strings
    cache_params = json.loads(data["litellm_cache_params"])
    assert isinstance(cache_params, dict)


def test_cache_ping_with_circular_reference(mock_redis_success):
    """Test cache ping with circular reference in cache parameters"""
    # Create circular reference
    circular_dict = {}
    circular_dict["self"] = circular_dict
    mock_redis_success.cache.circular_ref = circular_dict

    response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-9876"})
    assert response.status_code == 200

    # Verify response is still JSON serializable
    data = response.json()
    assert data["status"] == "healthy"


def test_cache_ping_failure(mock_redis_failure):
    """Test cache ping failure with expected error fields"""
    response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-9876"})
    assert response.status_code == 503

    data = response.json()
    print("data=", json.dumps(data, indent=4, default=str))

    assert "error" in data
    error = data["error"]

    # Verify error contains all expected fields
    assert "message" in error
    error_details = json.loads(error["message"])
    assert "message" in error_details
    assert "litellm_cache_params" in error_details
    assert "health_check_cache_params" in error_details

    # Verify generic static message (exception text must not leak to clients)
    assert error_details["message"] == "Service Unhealthy"


def test_cache_ping_failure_does_not_expose_traceback(mock_redis_failure):
    """CWE-209: Stack trace and exception text must not appear in the HTTP 503 response body."""
    response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-9876"})
    assert response.status_code == 503

    data = response.json()
    error = data.get("error", {})
    raw_body = json.dumps(data)

    # The word "traceback" (case-insensitive) must not appear anywhere in the response
    assert (
        "traceback" not in raw_body.lower()
    ), "CWE-209: Python traceback exposed in HTTP 503 response body"
    # Internal frame paths should not leak either
    assert (
        'File "' not in raw_body
    ), "CWE-209: Python stack frame paths exposed in HTTP 503 response body"
    # Exception text (e.g. Redis hostnames/IPs) must not leak either
    assert (
        "invalid username-password pair" not in raw_body
    ), "CWE-209: Exception message text exposed in HTTP 503 response body"

    # The error message should be the safe static string
    error_details = json.loads(error["message"])
    assert error_details["message"] == "Service Unhealthy"


def test_cache_ping_no_cache_initialized():
    """Test cache ping when no cache is initialized returns 503 with ProxyException envelope.

    Verifies the exact response structure so that regressions in the error format
    (e.g. message moving to a different field, or extra internal details leaking)
    are caught immediately.
    """
    original_cache = litellm.cache
    litellm.cache = None

    try:
        response = client.get(
            "/cache/ping", headers={"Authorization": "Bearer sk-9876"}
        )
        assert response.status_code == 503

        data = response.json()
        print("response data=", json.dumps(data, indent=4))
        # ProxyException is serialised as {"error": {"message": "...", "type": ..., ...}}
        assert "error" in data
        error_details = json.loads(data["error"]["message"])
        assert (
            error_details["message"] == "Cache not initialized. litellm.cache is None"
        )
    finally:
        litellm.cache = original_cache


def test_cache_ping_no_cache_does_not_expose_internals():
    """CWE-209: No-cache 503 must use the ProxyException envelope with no internal details.

    The null-cache path raises ProxyException directly (not HTTPException), so the
    response is {"error": {"message": "...", ...}} — same envelope as other 503s from
    this endpoint — with no tracebacks, source paths, or extra fields leaking.
    """
    original_cache = litellm.cache
    litellm.cache = None

    try:
        response = client.get(
            "/cache/ping", headers={"Authorization": "Bearer sk-9876"}
        )
        assert response.status_code == 503

        raw_body = response.text
        # No Python traceback or source-file paths must appear in the response
        assert "traceback" not in raw_body.lower(), (
            "CWE-209: Python traceback exposed in /cache/ping no-cache response"
        )
        assert 'File "' not in raw_body, (
            "CWE-209: Python stack frame paths exposed in /cache/ping no-cache response"
        )

        data = response.json()
        # Response must use the ProxyException envelope
        assert "error" in data, f"Expected ProxyException envelope, got: {data}"
        error_details = json.loads(data["error"]["message"])
        assert (
            error_details["message"] == "Cache not initialized. litellm.cache is None"
        )
    finally:
        litellm.cache = original_cache


def test_cache_ping_health_check_includes_only_cache_attributes(mock_redis_success):
    """
    Ensure that the /cache/ping endpoint only pulls HealthCheckCacheParams from litellm.cache.cache,
    and not from other attributes on litellm.cache.
    """
    # Add an unrelated field directly to the cache mock; it should NOT appear in health_check_cache_params
    mock_redis_success.some_unrelated_field = "should-not-appear-in-health-check"

    # Add a field on the underlying `cache` object that SHOULD appear
    mock_redis_success.cache.redis_kwargs = {"host": "localhost", "port": 6379}

    response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-9876"})
    assert (
        response.status_code == 200
    ), f"Unexpected status code: {response.status_code}"

    data = response.json()
    print("/cache/ping response data=", json.dumps(data, indent=4))
    health_check_cache_params = data.get("health_check_cache_params", {})
    # The unrelated field we attached at the top-level of litellm.cache should *not* be present
    assert (
        "some_unrelated_field" not in health_check_cache_params
    ), "Found an unexpected field from the mock_redis_success object in health_check_cache_params"

    # The field we attached to 'mock_redis_success.cache' should be present and correctly reported
    assert (
        "redis_kwargs" in health_check_cache_params
    ), "Expected field on `litellm.cache.cache` was not found in health_check_cache_params"
    assert health_check_cache_params["redis_kwargs"] == {
        "host": "localhost",
        "port": 6379,
    }


def test_cache_ping_with_redis_version_float(mock_redis_success):
    """Test cache ping works when redis_version is a float"""
    # Set redis_version as a float
    mock_redis_success.cache.redis_version = 7.2

    response = client.get("/cache/ping", headers={"Authorization": "Bearer sk-9876"})
    assert response.status_code == 200

    data = response.json()
    print("data=", json.dumps(data, indent=4))
    assert data["status"] == "healthy"
    assert data["cache_type"] == "redis"

    cache_params = data["health_check_cache_params"]
    assert isinstance(cache_params, dict)
    assert isinstance(cache_params.get("redis_version"), float)


@pytest.fixture
def mock_redis_client_list_restricted(mocker):
    """Mock Redis cache where CLIENT LIST is restricted (like GCP Redis)"""

    def mock_client_list():
        raise Exception("ERR unknown command 'CLIENT'")

    def mock_info():
        return {
            "redis_version": "6.2.7",
            "used_memory": "1000000",
            "connected_clients": "5",
            "keyspace_hits": "1000",
            "keyspace_misses": "100",
        }

    mock_cache = mocker.MagicMock()
    mock_cache.type = "redis"
    mock_cache.cache = RedisCache(host="localhost", port=6379, password="hello")
    mock_cache.cache.client_list = mock_client_list
    mock_cache.cache.info = mock_info

    mocker.patch.object(litellm, "cache", mock_cache)
    return mock_cache


@pytest.fixture
def mock_redis_client_list_success(mocker):
    """Mock Redis cache where CLIENT LIST works normally"""

    def mock_client_list():
        return [
            {"id": "1", "addr": "127.0.0.1:54321", "name": "client1"},
            {"id": "2", "addr": "127.0.0.1:54322", "name": "client2"},
        ]

    def mock_info():
        return {
            "redis_version": "6.2.7",
            "used_memory": "1000000",
            "connected_clients": "2",
        }

    mock_cache = mocker.MagicMock()
    mock_cache.type = "redis"
    mock_cache.cache = RedisCache(host="localhost", port=6379, password="hello")
    mock_cache.cache.client_list = mock_client_list
    mock_cache.cache.info = mock_info

    mocker.patch.object(litellm, "cache", mock_cache)
    return mock_cache


def test_cache_redis_info_no_cache():
    """Test /cache/redis/info when no cache is initialized"""
    original_cache = litellm.cache
    litellm.cache = None

    response = client.get(
        "/cache/redis/info", headers={"Authorization": "Bearer sk-9876"}
    )
    assert response.status_code == 503

    data = response.json()
    assert "Cache not initialized" in data["detail"]

    # Restore original cache
    litellm.cache = original_cache


def _use_redis_backed_cache(monkeypatch, backend_attributes):
    cache = Cache(type="redis", host="localhost", port=6379, password="hello")
    vars(cache.cache).update(backend_attributes)
    monkeypatch.setattr(litellm, "cache", cache)


@pytest.mark.parametrize(
    ("field", "value"),
    [
        ("host", "redis.internal"),
        ("port", 6379),
        ("port", "6379"),
        ("namespace", "team-a"),
        ("redis_version", 7.2),
        ("redis_version", 7),
        ("redis_version", "7.2.4"),
    ],
)
def test_extract_cache_params_keeps_each_backend_value_and_its_type(monkeypatch, field, value):
    _use_redis_backed_cache(monkeypatch, {field: value})

    params = _extract_cache_params()

    assert sorted(params) == ["host", "namespace", "port", "redis_kwargs", "redis_version"]
    assert repr(params[field]) == repr(value)


def test_extract_cache_params_reads_connection_settings_and_masks_the_password(monkeypatch):
    _use_redis_backed_cache(monkeypatch, {})

    redis_kwargs = _extract_cache_params()["redis_kwargs"]

    assert redis_kwargs["host"] == "localhost"
    assert repr(redis_kwargs["port"]) == "6379"
    assert "hello" not in json.dumps(redis_kwargs)


def test_extract_cache_params_reports_every_malformed_backend_field(monkeypatch):
    _use_redis_backed_cache(monkeypatch, {"host": 7, "namespace": 5, "redis_version": ["7"]})

    with pytest.raises(ValidationError) as exc_info:
        _extract_cache_params()

    assert exc_info.value.title == "HealthCheckCacheParams"
    assert sorted({error["loc"][0] for error in exc_info.value.errors()}) == ["host", "namespace", "redis_version"]


def test_extract_cache_params_is_empty_when_no_cache_is_configured(monkeypatch):
    monkeypatch.setattr(litellm, "cache", None)

    assert _extract_cache_params() == {}


@pytest.mark.parametrize("attribute_name", [1, None, (1, 2), b"host"])
def test_extract_cache_params_is_empty_when_a_backend_attribute_name_is_not_a_string(monkeypatch, attribute_name):
    _use_redis_backed_cache(monkeypatch, {attribute_name: "redis.internal"})

    assert _extract_cache_params() == {}, attribute_name