mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-28 01:32:17 +00:00
test(redis): separate unit and cluster integration coverage
Keep the mapped test tree mock-only. Run the live cluster cases from the existing local caching integration file, selected by explicit node IDs in the Redis compatibility workflow. Signed-off-by: Silu Panda <31051721+SiluPanda@users.noreply.github.com>
This commit is contained in:
Silu Panda
parent
cda44a021f
commit
c8ff8c1237
3 changed files with 141 additions and 83 deletions
3
.github/workflows/test-redis-compat.yml
vendored
3
.github/workflows/test-redis-compat.yml
vendored
|
|
@ -11,6 +11,7 @@ on:
|
|||
- "litellm/_redis.py"
|
||||
- "litellm/_redis_credential_provider.py"
|
||||
- "tests/test_litellm/test_redis.py"
|
||||
- "tests/local_testing/test_caching.py"
|
||||
- "tests/test_litellm/caching/test_redis_connection_pool.py"
|
||||
- ".github/workflows/test-redis-compat.yml"
|
||||
- "pyproject.toml"
|
||||
|
|
@ -83,6 +84,8 @@ jobs:
|
|||
uv run --no-sync pytest \
|
||||
tests/test_litellm/test_redis.py \
|
||||
tests/test_litellm/caching/test_redis_connection_pool.py \
|
||||
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_azure_credentials \
|
||||
tests/local_testing/test_caching.py::test_sync_cluster_authenticates_with_gcp_credentials \
|
||||
--tb=short -vv \
|
||||
--reruns 2 \
|
||||
--reruns-delay 1 \
|
||||
|
|
|
|||
|
|
@ -1,6 +1,17 @@
|
|||
import os
|
||||
import time
|
||||
import traceback
|
||||
import shutil
|
||||
import subprocess
|
||||
from collections.abc import Callable, Iterator
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
|
||||
import redis
|
||||
|
||||
from litellm._redis import _get_redis_env_kwarg_mapping, get_redis_client
|
||||
from litellm._redis_credential_provider import _token_cache
|
||||
from litellm._uuid import uuid
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
|
@ -1032,6 +1043,102 @@ def test_redis_cache_completion_stream():
|
|||
# test_redis_cache_completion_stream()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clean_cluster_iam_environment(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
||||
for var in ("REDIS_URL", "REDIS_CLUSTER_NODES", "REDIS_SENTINEL_NODES", *_get_redis_env_kwarg_mapping()):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
_token_cache.clear()
|
||||
yield
|
||||
_token_cache.clear()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def authenticated_redis_cluster(tmp_path: Path, unused_tcp_port_factory: Callable[[], int]) -> Iterator[int]:
|
||||
server: Final = shutil.which("redis-server")
|
||||
if server is None:
|
||||
pytest.skip("redis-server is required for the cluster authentication regression tests")
|
||||
port: Final = unused_tcp_port_factory()
|
||||
bus_port: Final = unused_tcp_port_factory()
|
||||
log_path: Final = tmp_path / "redis.log"
|
||||
config: Final = tmp_path / "redis.conf"
|
||||
config.write_text(
|
||||
f"bind 127.0.0.1\nport {port}\ncluster-port {bus_port}\n"
|
||||
f'cluster-enabled yes\ncluster-config-file "{tmp_path / "nodes.conf"}"\n'
|
||||
f'dir "{tmp_path}"\nsave ""\nappendonly no\n'
|
||||
)
|
||||
with log_path.open("w") as log:
|
||||
process: Final = subprocess.Popen((server, str(config)), stdout=log, stderr=subprocess.STDOUT)
|
||||
try:
|
||||
with redis.Redis(host="127.0.0.1", port=port, socket_timeout=1, socket_connect_timeout=1) as admin:
|
||||
for _ in range(100):
|
||||
try:
|
||||
admin.ping()
|
||||
break
|
||||
except redis.ConnectionError:
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
pytest.fail(f"Redis did not start: {log_path.read_text()}")
|
||||
admin.execute_command("CLUSTER", "ADDSLOTS", *range(16384))
|
||||
for _ in range(100):
|
||||
if admin.cluster("INFO")["cluster_state"] == "ok":
|
||||
break
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
pytest.fail(f"Redis cluster did not become ready: {log_path.read_text()}")
|
||||
admin.execute_command(
|
||||
"ACL", "SETUSER", "identity-object-id", "on", ">local-fixture-token", "allcommands", "allkeys"
|
||||
)
|
||||
admin.execute_command("ACL", "SETUSER", "default", "resetpass", ">local-fixture-token")
|
||||
yield port
|
||||
finally:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait()
|
||||
|
||||
|
||||
def test_sync_cluster_authenticates_with_azure_credentials(
|
||||
clean_cluster_iam_environment: None, monkeypatch: pytest.MonkeyPatch, authenticated_redis_cluster: int
|
||||
) -> None:
|
||||
monkeypatch.setenv("REDIS_USERNAME", "identity-object-id")
|
||||
credential: Final = MagicMock()
|
||||
credential.get_token.return_value = SimpleNamespace(token="local-fixture-token")
|
||||
|
||||
with patch("azure.identity.DefaultAzureCredential", return_value=credential):
|
||||
with get_redis_client(
|
||||
startup_nodes=[{"host": "127.0.0.1", "port": authenticated_redis_cluster}],
|
||||
azure_redis_ad_token=True,
|
||||
password="stale-password",
|
||||
socket_timeout=1,
|
||||
socket_connect_timeout=1,
|
||||
) as client:
|
||||
assert client.ping() is True
|
||||
assert client.set("iam-regression", "success") is True
|
||||
assert client.get("iam-regression") == b"success"
|
||||
|
||||
|
||||
def test_sync_cluster_authenticates_with_gcp_credentials(
|
||||
clean_cluster_iam_environment: None, authenticated_redis_cluster: int
|
||||
) -> None:
|
||||
iam_client: Final = MagicMock()
|
||||
iam_client.generate_access_token.return_value = SimpleNamespace(access_token="local-fixture-token")
|
||||
|
||||
with patch("google.cloud.iam_credentials_v1.IAMCredentialsClient", return_value=iam_client):
|
||||
with get_redis_client(
|
||||
startup_nodes=[{"host": "127.0.0.1", "port": authenticated_redis_cluster}],
|
||||
gcp_service_account="projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com",
|
||||
username="stale-user",
|
||||
password="stale-password",
|
||||
socket_timeout=1,
|
||||
socket_connect_timeout=1,
|
||||
) as client:
|
||||
assert client.ping() is True
|
||||
assert client.set("iam-regression", "success") is True
|
||||
assert client.get("iam-regression") == b"success"
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="Local test. Requires running redis cluster locally.")
|
||||
@pytest.mark.asyncio
|
||||
async def test_redis_cache_cluster_init_unit_test():
|
||||
|
|
|
|||
|
|
@ -1,12 +1,6 @@
|
|||
import inspect
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from collections.abc import Callable, Iterator
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -249,91 +243,45 @@ def test_sync_cluster_preserves_credential_provider_identity(clean_redis_environ
|
|||
assert [(node.host, node.port) for node in cluster_kwargs["startup_nodes"]] == [("cluster-node", 6379)]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def authenticated_redis_cluster(tmp_path: Path, unused_tcp_port_factory: Callable[[], int]) -> Iterator[int]:
|
||||
server: Final = shutil.which("redis-server")
|
||||
if server is None:
|
||||
pytest.skip("redis-server is required for the cluster authentication regression tests")
|
||||
port: Final = unused_tcp_port_factory()
|
||||
bus_port: Final = unused_tcp_port_factory()
|
||||
log_path: Final = tmp_path / "redis.log"
|
||||
config: Final = tmp_path / "redis.conf"
|
||||
config.write_text(
|
||||
f"bind 127.0.0.1\nport {port}\ncluster-port {bus_port}\n"
|
||||
f"cluster-enabled yes\ncluster-config-file {tmp_path / 'nodes.conf'}\n"
|
||||
f'dir "{tmp_path}"\nsave ""\nappendonly no\n'
|
||||
)
|
||||
with log_path.open("w") as log:
|
||||
process: Final = subprocess.Popen((server, str(config)), stdout=log, stderr=subprocess.STDOUT)
|
||||
try:
|
||||
with redis.Redis(host="127.0.0.1", port=port, socket_timeout=1, socket_connect_timeout=1) as admin:
|
||||
for _ in range(100):
|
||||
try:
|
||||
admin.ping()
|
||||
break
|
||||
except redis.ConnectionError:
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
pytest.fail(f"Redis did not start: {log_path.read_text()}")
|
||||
admin.execute_command("CLUSTER", "ADDSLOTS", *range(16384))
|
||||
for _ in range(100):
|
||||
if admin.cluster("INFO")["cluster_state"] == "ok":
|
||||
break
|
||||
time.sleep(0.1)
|
||||
else:
|
||||
pytest.fail(f"Redis cluster did not become ready: {log_path.read_text()}")
|
||||
admin.execute_command(
|
||||
"ACL", "SETUSER", "identity-object-id", "on", ">local-fixture-token", "allcommands", "allkeys"
|
||||
)
|
||||
admin.execute_command("ACL", "SETUSER", "default", "resetpass", ">local-fixture-token")
|
||||
yield port
|
||||
finally:
|
||||
process.terminate()
|
||||
try:
|
||||
process.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
process.kill()
|
||||
process.wait()
|
||||
|
||||
|
||||
def test_sync_cluster_authenticates_with_azure_credentials(
|
||||
clean_redis_environment: None, monkeypatch: pytest.MonkeyPatch, authenticated_redis_cluster: int
|
||||
) -> None:
|
||||
def test_sync_cluster_authenticates_with_azure_credentials(clean_redis_environment, monkeypatch):
|
||||
monkeypatch.setenv("REDIS_USERNAME", "identity-object-id")
|
||||
credential: Final = MagicMock()
|
||||
credential.get_token.return_value = SimpleNamespace(token="local-fixture-token")
|
||||
credential = MagicMock()
|
||||
credential.get_token.return_value = SimpleNamespace(token="azure-access-token")
|
||||
|
||||
with patch("azure.identity.DefaultAzureCredential", return_value=credential):
|
||||
with get_redis_client(
|
||||
startup_nodes=[{"host": "127.0.0.1", "port": authenticated_redis_cluster}],
|
||||
with (
|
||||
patch("azure.identity.DefaultAzureCredential", return_value=credential),
|
||||
patch("redis.RedisCluster", autospec=True) as cluster,
|
||||
):
|
||||
get_redis_client(
|
||||
startup_nodes=[{"host": "cluster-node", "port": 6379}],
|
||||
azure_redis_ad_token=True,
|
||||
password="stale-password",
|
||||
socket_timeout=1,
|
||||
socket_connect_timeout=1,
|
||||
) as client:
|
||||
assert client.ping() is True
|
||||
assert client.set("iam-regression", "success") is True
|
||||
assert client.get("iam-regression") == b"success"
|
||||
)
|
||||
|
||||
kwargs = cluster.call_args.kwargs
|
||||
provider = kwargs.get("credential_provider")
|
||||
assert isinstance(provider, AzureADCredentialProvider)
|
||||
assert provider.get_credentials() == ("identity-object-id", "azure-access-token")
|
||||
assert "username" not in kwargs
|
||||
assert "password" not in kwargs
|
||||
assert "redis_connect_func" not in kwargs
|
||||
credential.get_token.assert_called_once_with("https://redis.azure.com/.default")
|
||||
|
||||
|
||||
def test_sync_cluster_authenticates_with_gcp_credentials(
|
||||
clean_redis_environment: None, authenticated_redis_cluster: int
|
||||
) -> None:
|
||||
iam_client: Final = MagicMock()
|
||||
iam_client.generate_access_token.return_value = SimpleNamespace(access_token="local-fixture-token")
|
||||
|
||||
with patch("google.cloud.iam_credentials_v1.IAMCredentialsClient", return_value=iam_client):
|
||||
with get_redis_client(
|
||||
startup_nodes=[{"host": "127.0.0.1", "port": authenticated_redis_cluster}],
|
||||
gcp_service_account="projects/-/serviceAccounts/sa@project.iam.gserviceaccount.com",
|
||||
def test_sync_cluster_authenticates_with_gcp_credentials(clean_redis_environment):
|
||||
with patch("redis.RedisCluster", autospec=True) as cluster:
|
||||
get_redis_client(
|
||||
startup_nodes=[{"host": "cluster-node", "port": 6379}],
|
||||
redis_connect_func=_gcp_marker_callback(),
|
||||
username="stale-user",
|
||||
password="stale-password",
|
||||
socket_timeout=1,
|
||||
socket_connect_timeout=1,
|
||||
) as client:
|
||||
assert client.ping() is True
|
||||
assert client.set("iam-regression", "success") is True
|
||||
assert client.get("iam-regression") == b"success"
|
||||
)
|
||||
|
||||
kwargs = cluster.call_args.kwargs
|
||||
assert isinstance(kwargs.get("credential_provider"), GCPIAMCredentialProvider)
|
||||
assert "username" not in kwargs
|
||||
assert "password" not in kwargs
|
||||
assert "redis_connect_func" not in kwargs
|
||||
|
||||
|
||||
def test_async_cluster_preserves_credential_provider_identity(clean_redis_environment):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue