diff --git a/tests/litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index cde43158379..697d985dc93 100644 --- a/tests/litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -2,7 +2,7 @@ import json import os import sys from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient @@ -15,306 +15,251 @@ from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager -# Mock Prisma client class -class MockPrismaClient: +class MockRedisCache: def __init__(self): - self.db = MagicMock() - self.db.litellm_cronjob = AsyncMock() + self.async_set_cache = AsyncMock() + self.async_get_cache = AsyncMock() + self.async_delete_cache = AsyncMock() @pytest.fixture -def mock_prisma(monkeypatch): - mock_client = MockPrismaClient() - - # Mock the prisma_client import in proxy_server - def mock_get_prisma(): - return mock_client - - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_client) - return mock_client +def mock_redis(): + return MockRedisCache() @pytest.fixture -def pod_lock_manager(): - return PodLockManager(cronjob_id="test_job") +def pod_lock_manager(mock_redis): + return PodLockManager(cronjob_id="test_job", redis_cache=mock_redis) @pytest.mark.asyncio -async def test_acquire_lock_success(pod_lock_manager, mock_prisma): +async def test_acquire_lock_success(pod_lock_manager, mock_redis): """ Test that the lock is acquired successfully when no existing lock exists """ - # Mock find_unique to return None (no existing lock) - mock_prisma.db.litellm_cronjob.find_unique.return_value = None - - # Mock successful creation of new lock - mock_response = AsyncMock() - mock_response.status = "ACTIVE" - mock_response.pod_id = pod_lock_manager.pod_id - mock_prisma.db.litellm_cronjob.create.return_value = mock_response + # Mock successful acquisition (SET NX returns True) + mock_redis.async_set_cache.return_value = True result = await pod_lock_manager.acquire_lock() assert result == True - # Verify find_unique was called - mock_prisma.db.litellm_cronjob.find_unique.assert_called_once() - # Verify create was called with correct parameters - mock_prisma.db.litellm_cronjob.create.assert_called_once() - call_args = mock_prisma.db.litellm_cronjob.create.call_args[1] - assert call_args["data"]["cronjob_id"] == "test_job" - assert call_args["data"]["pod_id"] == pod_lock_manager.pod_id - assert call_args["data"]["status"] == "ACTIVE" + # Verify set_cache was called with correct parameters + mock_redis.async_set_cache.assert_called_once_with( + pod_lock_manager.lock_key, + pod_lock_manager.pod_id, + nx=True, + ttl=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS, + ) @pytest.mark.asyncio -async def test_acquire_lock_existing_active(pod_lock_manager, mock_prisma): +async def test_acquire_lock_existing_active(pod_lock_manager, mock_redis): """ Test that the lock is not acquired if there's an active lock by different pod """ - # Mock existing active lock - mock_existing = AsyncMock() - mock_existing.status = "ACTIVE" - mock_existing.pod_id = "different_pod_id" - mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) # Future TTL - mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + # Mock failed acquisition (SET NX returns False) + mock_redis.async_set_cache.return_value = False + # Mock get_cache to return a different pod's ID + mock_redis.async_get_cache.return_value = "different_pod_id" result = await pod_lock_manager.acquire_lock() assert result == False - # Verify find_unique was called but update/create were not - mock_prisma.db.litellm_cronjob.find_unique.assert_called_once() - mock_prisma.db.litellm_cronjob.update.assert_not_called() - mock_prisma.db.litellm_cronjob.create.assert_not_called() + # Verify set_cache was called + mock_redis.async_set_cache.assert_called_once() + # Verify get_cache was called to check existing lock + mock_redis.async_get_cache.assert_called_once_with(pod_lock_manager.lock_key) @pytest.mark.asyncio -async def test_acquire_lock_expired(pod_lock_manager, mock_prisma): +async def test_acquire_lock_expired(pod_lock_manager, mock_redis): """ Test that the lock can be acquired if existing lock is expired """ - # Mock existing expired lock - mock_existing = AsyncMock() - mock_existing.status = "ACTIVE" - mock_existing.pod_id = "different_pod_id" - mock_existing.ttl = datetime.now(timezone.utc) - timedelta(seconds=30) # Past TTL - mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + # Mock failed acquisition first (SET NX returns False) + mock_redis.async_set_cache.return_value = False - # Mock successful update - mock_updated = AsyncMock() - mock_updated.pod_id = pod_lock_manager.pod_id - mock_prisma.db.litellm_cronjob.update.return_value = mock_updated + # Simulate an expired lock by having the TTL return a value + # Since Redis auto-expires keys, an expired lock would be absent + # So we'll simulate a retry after the first check fails + # First check returns a value (lock exists) + mock_redis.async_get_cache.return_value = "different_pod_id" + + # Then set succeeds on retry (simulating key expiring between checks) + mock_redis.async_set_cache.side_effect = [False, True] + + result = await pod_lock_manager.acquire_lock() + assert result == False # First attempt fails + + # Reset mock for a second attempt + mock_redis.async_set_cache.reset_mock() + mock_redis.async_set_cache.return_value = True + + # Try again (simulating the lock expired) result = await pod_lock_manager.acquire_lock() assert result == True - # Verify both find_unique and update were called - mock_prisma.db.litellm_cronjob.find_unique.assert_called_once() - mock_prisma.db.litellm_cronjob.update.assert_called_once() + # Verify set_cache was called again + mock_redis.async_set_cache.assert_called_once() @pytest.mark.asyncio -async def test_renew_lock(pod_lock_manager, mock_prisma): +async def test_release_lock_success(pod_lock_manager, mock_redis): """ - Test that the renew lock calls the DB update method with the correct parameters + Test that the release lock works when the current pod holds the lock """ - mock_prisma.db.litellm_cronjob.update.return_value = AsyncMock() - - await pod_lock_manager.renew_lock() - - # Verify update was called with correct parameters - mock_prisma.db.litellm_cronjob.update.assert_called_once() - call_args = mock_prisma.db.litellm_cronjob.update.call_args[1] - assert call_args["where"]["cronjob_id"] == "test_job" - assert call_args["where"]["pod_id"] == pod_lock_manager.pod_id - assert "ttl" in call_args["data"] - assert "last_updated" in call_args["data"] - - -@pytest.mark.asyncio -async def test_release_lock(pod_lock_manager, mock_prisma): - """ - Test that the release lock calls the DB update method with the correct parameters - - specifically, the status should be set to INACTIVE - """ - mock_prisma.db.litellm_cronjob.update.return_value = AsyncMock() + # Mock get_cache to return this pod's ID + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id + # Mock successful deletion + mock_redis.async_delete_cache.return_value = 1 await pod_lock_manager.release_lock() - # Verify update was called with correct parameters - mock_prisma.db.litellm_cronjob.update.assert_called_once() - call_args = mock_prisma.db.litellm_cronjob.update.call_args[1] - assert call_args["where"]["cronjob_id"] == "test_job" - assert call_args["where"]["pod_id"] == pod_lock_manager.pod_id - assert call_args["data"]["status"] == "INACTIVE" + # Verify get_cache was called + mock_redis.async_get_cache.assert_called_once_with(pod_lock_manager.lock_key) + # Verify delete_cache was called + mock_redis.async_delete_cache.assert_called_once_with(pod_lock_manager.lock_key) @pytest.mark.asyncio -async def test_prisma_client_none(pod_lock_manager, monkeypatch): - # Mock prisma_client as None - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) +async def test_release_lock_different_pod(pod_lock_manager, mock_redis): + """ + Test that the release lock doesn't delete when a different pod holds the lock + """ + # Mock get_cache to return a different pod's ID + mock_redis.async_get_cache.return_value = "different_pod_id" - # Test all methods with None client - assert await pod_lock_manager.acquire_lock() == False - assert await pod_lock_manager.renew_lock() == False - assert await pod_lock_manager.release_lock() == False + await pod_lock_manager.release_lock() + + # Verify get_cache was called + mock_redis.async_get_cache.assert_called_once_with(pod_lock_manager.lock_key) + # Verify delete_cache was NOT called + mock_redis.async_delete_cache.assert_not_called() @pytest.mark.asyncio -async def test_database_error_handling(pod_lock_manager, mock_prisma): - # Mock database errors - mock_prisma.db.litellm_cronjob.upsert.side_effect = Exception("Database error") - mock_prisma.db.litellm_cronjob.update.side_effect = Exception("Database error") +async def test_release_lock_no_lock(pod_lock_manager, mock_redis): + """ + Test release lock behavior when no lock exists + """ + # Mock get_cache to return None (no lock) + mock_redis.async_get_cache.return_value = None - # Test error handling in all methods - assert await pod_lock_manager.acquire_lock() == False - await pod_lock_manager.renew_lock() # Should not raise exception - await pod_lock_manager.release_lock() # Should not raise exception + await pod_lock_manager.release_lock() + + # Verify get_cache was called + mock_redis.async_get_cache.assert_called_once_with(pod_lock_manager.lock_key) + # Verify delete_cache was NOT called + mock_redis.async_delete_cache.assert_not_called() @pytest.mark.asyncio -async def test_acquire_lock_inactive_status(pod_lock_manager, mock_prisma): +async def test_redis_none(monkeypatch): """ - Test that the lock can be acquired if existing lock is INACTIVE + Test behavior when redis_cache is None """ - # Mock existing inactive lock - mock_existing = AsyncMock() - mock_existing.status = "INACTIVE" - mock_existing.pod_id = "different_pod_id" - mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) - mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + pod_lock_manager = PodLockManager(cronjob_id="test_job", redis_cache=None) - # Mock successful update - mock_updated = AsyncMock() - mock_updated.pod_id = pod_lock_manager.pod_id - mock_prisma.db.litellm_cronjob.update.return_value = mock_updated + # Test acquire_lock with None redis_cache + assert await pod_lock_manager.acquire_lock() is None - result = await pod_lock_manager.acquire_lock() - assert result == True - - mock_prisma.db.litellm_cronjob.update.assert_called_once() + # Test release_lock with None redis_cache (should not raise exception) + await pod_lock_manager.release_lock() @pytest.mark.asyncio -async def test_acquire_lock_same_pod(pod_lock_manager, mock_prisma): +async def test_redis_error_handling(pod_lock_manager, mock_redis): """ - Test that the lock returns True if the same pod already holds the lock + Test error handling in Redis operations """ - # Mock existing active lock held by same pod - mock_existing = AsyncMock() - mock_existing.status = "ACTIVE" - mock_existing.pod_id = pod_lock_manager.pod_id - mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) - mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing - - result = await pod_lock_manager.acquire_lock() - assert result == True - - # Verify no update was needed - mock_prisma.db.litellm_cronjob.update.assert_not_called() - mock_prisma.db.litellm_cronjob.create.assert_not_called() - - -@pytest.mark.asyncio -async def test_acquire_lock_race_condition(pod_lock_manager, mock_prisma): - """ - Test handling of potential race conditions during lock acquisition - """ - # First find_unique returns None - mock_prisma.db.litellm_cronjob.find_unique.return_value = None - - # But create raises unique constraint violation - mock_prisma.db.litellm_cronjob.create.side_effect = Exception( - "Unique constraint violation" - ) + # Mock exceptions for Redis operations + mock_redis.async_set_cache.side_effect = Exception("Redis error") + mock_redis.async_get_cache.side_effect = Exception("Redis error") + mock_redis.async_delete_cache.side_effect = Exception("Redis error") + # Test acquire_lock error handling result = await pod_lock_manager.acquire_lock() assert result == False + # Reset side effect for get_cache for the release test + mock_redis.async_get_cache.side_effect = None + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id -@pytest.mark.asyncio -async def test_ttl_calculation(pod_lock_manager, mock_prisma): - """ - Test that TTL is calculated correctly when acquiring lock - """ - mock_prisma.db.litellm_cronjob.find_unique.return_value = None - mock_prisma.db.litellm_cronjob.create.return_value = AsyncMock() - - await pod_lock_manager.acquire_lock() - - call_args = mock_prisma.db.litellm_cronjob.create.call_args[1] - ttl = call_args["data"]["ttl"] - - # Verify TTL is in the future by DEFAULT_CRON_JOB_LOCK_TTL_SECONDS - expected_ttl = datetime.now(timezone.utc) + timedelta( - seconds=DEFAULT_CRON_JOB_LOCK_TTL_SECONDS - ) - assert abs((ttl - expected_ttl).total_seconds()) < 1 # Allow 1 second difference + # Test release_lock error handling (should not raise exception) + await pod_lock_manager.release_lock() @pytest.mark.asyncio -async def test_concurrent_lock_acquisition_simulation(mock_prisma): +async def test_bytes_handling(pod_lock_manager, mock_redis): + """ + Test handling of bytes values from Redis + """ + # Mock failed acquisition + mock_redis.async_set_cache.return_value = False + # Mock get_cache to return bytes + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id.encode("utf-8") + + result = await pod_lock_manager.acquire_lock() + assert result == True + + # Reset for release test + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id.encode("utf-8") + mock_redis.async_delete_cache.return_value = 1 + + await pod_lock_manager.release_lock() + mock_redis.async_delete_cache.assert_called_once() + + +@pytest.mark.asyncio +async def test_concurrent_lock_acquisition_simulation(): """ Simulate multiple pods trying to acquire the lock simultaneously """ - pod1 = PodLockManager(cronjob_id="test_job") - pod2 = PodLockManager(cronjob_id="test_job") - pod3 = PodLockManager(cronjob_id="test_job") + mock_redis = MockRedisCache() + pod1 = PodLockManager(cronjob_id="test_job", redis_cache=mock_redis) + pod2 = PodLockManager(cronjob_id="test_job", redis_cache=mock_redis) + pod3 = PodLockManager(cronjob_id="test_job", redis_cache=mock_redis) # Simulate first pod getting the lock - mock_prisma.db.litellm_cronjob.find_unique.return_value = None - mock_response = AsyncMock() - mock_response.pod_id = pod1.pod_id - mock_response.status = "ACTIVE" - mock_prisma.db.litellm_cronjob.create.return_value = mock_response + mock_redis.async_set_cache.return_value = True # First pod should get the lock result1 = await pod1.acquire_lock() assert result1 == True - # Simulate other pods trying to acquire same lock immediately after - mock_existing = AsyncMock() - mock_existing.status = "ACTIVE" - mock_existing.pod_id = pod1.pod_id - mock_existing.ttl = datetime.now(timezone.utc) + timedelta(seconds=30) - mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + # Simulate other pods failing to get the lock + mock_redis.async_set_cache.return_value = False + mock_redis.async_get_cache.return_value = pod1.pod_id # Other pods should fail to acquire result2 = await pod2.acquire_lock() result3 = await pod3.acquire_lock() + + # Since other pods don't have the lock, they should get False assert result2 == False assert result3 == False @pytest.mark.asyncio -async def test_lock_takeover_race_condition(mock_prisma): +async def test_lock_takeover_race_condition(mock_redis): """ - Test scenario where multiple pods try to take over an expired lock + Test scenario where multiple pods try to take over an expired lock using Redis """ - pod1 = PodLockManager(cronjob_id="test_job") - pod2 = PodLockManager(cronjob_id="test_job") + pod1 = PodLockManager(cronjob_id="test_job", redis_cache=mock_redis) + pod2 = PodLockManager(cronjob_id="test_job", redis_cache=mock_redis) - # Simulate expired lock - mock_existing = AsyncMock() - mock_existing.status = "ACTIVE" - mock_existing.pod_id = "old_pod" - mock_existing.ttl = datetime.now(timezone.utc) - timedelta(seconds=30) - mock_prisma.db.litellm_cronjob.find_unique.return_value = mock_existing + # Simulate first pod's acquisition succeeding + mock_redis.async_set_cache.return_value = True - # Simulate pod1's update succeeding - mock_update1 = AsyncMock() - mock_update1.pod_id = pod1.pod_id - mock_prisma.db.litellm_cronjob.update.return_value = mock_update1 - - # First pod should successfully take over + # First pod should successfully acquire result1 = await pod1.acquire_lock() assert result1 == True - # Simulate pod2's update failing due to race condition - mock_prisma.db.litellm_cronjob.update.side_effect = Exception( - "Row was updated by another transaction" - ) + # Simulate race condition: second pod tries but fails + mock_redis.async_set_cache.return_value = False + mock_redis.async_get_cache.return_value = pod1.pod_id - # Second pod should fail to take over + # Second pod should fail to acquire result2 = await pod2.acquire_lock() assert result2 == False