mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-29 01:42:19 +00:00
test(redis): exercise IAM cluster authentication over TCP
Run Azure and GCP regressions against a real local cluster with only cloud token issuance stubbed. Build a checksum-verified Redis server in the compatibility workflow and report its coverage. Signed-off-by: Silu Panda <31051721+SiluPanda@users.noreply.github.com>
This commit is contained in:
Silu Panda
parent
28499f4a0a
commit
cda44a021f
2 changed files with 108 additions and 34 deletions
26
.github/workflows/test-redis-compat.yml
vendored
26
.github/workflows/test-redis-compat.yml
vendored
|
|
@ -28,6 +28,9 @@ jobs:
|
|||
name: "redis-py ${{ matrix.redis-version }}"
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
permissions:
|
||||
contents: read
|
||||
id-token: write
|
||||
|
||||
strategy:
|
||||
fail-fast: false
|
||||
|
|
@ -57,7 +60,7 @@ jobs:
|
|||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra extra_proxy --extra semantic-router
|
||||
|
||||
- name: Pin redis-py to the matrix version
|
||||
env:
|
||||
|
|
@ -66,12 +69,31 @@ jobs:
|
|||
uv pip install "redis==${REDIS_VERSION:?}"
|
||||
uv run --no-sync python -c "import redis; assert redis.__version__ == '${REDIS_VERSION:?}', redis.__version__; print('redis-py', redis.__version__)"
|
||||
|
||||
- name: Build Redis for cluster authentication tests
|
||||
run: |
|
||||
curl --fail --location --retry 3 https://download.redis.io/releases/redis-7.2.16.tar.gz -o "$RUNNER_TEMP/redis-7.2.16.tar.gz"
|
||||
echo "960a8ec15e34ff40e57ff16837b26b33bd81f2da6d24497bb63de532a323a18e $RUNNER_TEMP/redis-7.2.16.tar.gz" | sha256sum --check
|
||||
tar -xzf "$RUNNER_TEMP/redis-7.2.16.tar.gz" -C "$RUNNER_TEMP"
|
||||
make -C "$RUNNER_TEMP/redis-7.2.16" -j2 MALLOC=libc OPTIMIZATION=-O1 redis-server
|
||||
echo "$RUNNER_TEMP/redis-7.2.16/src" >> "$GITHUB_PATH"
|
||||
|
||||
- name: Run redis unit tests
|
||||
run: |
|
||||
redis-server --version
|
||||
uv run --no-sync pytest \
|
||||
tests/test_litellm/test_redis.py \
|
||||
tests/test_litellm/caching/test_redis_connection_pool.py \
|
||||
--tb=short -vv \
|
||||
--reruns 2 \
|
||||
--reruns-delay 1 \
|
||||
--durations=20
|
||||
--durations=20 \
|
||||
--cov=./litellm --cov-report=xml:coverage-redis.xml
|
||||
|
||||
- name: Upload Redis coverage
|
||||
if: matrix.redis-version == '5.3.1'
|
||||
uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4
|
||||
with:
|
||||
use_oidc: true
|
||||
files: coverage-redis.xml
|
||||
flags: redis-compat
|
||||
fail_ci_if_error: false
|
||||
|
|
|
|||
|
|
@ -1,6 +1,12 @@
|
|||
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
|
||||
|
|
@ -243,45 +249,91 @@ 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)]
|
||||
|
||||
|
||||
def test_sync_cluster_authenticates_with_azure_credentials(clean_redis_environment, monkeypatch):
|
||||
monkeypatch.setenv("REDIS_USERNAME", "identity-object-id")
|
||||
credential = MagicMock()
|
||||
credential.get_token.return_value = SimpleNamespace(token="azure-access-token")
|
||||
@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()
|
||||
|
||||
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}],
|
||||
|
||||
def test_sync_cluster_authenticates_with_azure_credentials(
|
||||
clean_redis_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",
|
||||
)
|
||||
|
||||
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")
|
||||
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_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(),
|
||||
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",
|
||||
username="stale-user",
|
||||
password="stale-password",
|
||||
)
|
||||
|
||||
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
|
||||
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_async_cluster_preserves_credential_provider_identity(clean_redis_environment):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue