mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
test pod lock manager
This commit is contained in:
Ishaan Jaff
parent
8b12a2e5dc
commit
8405fcb748
1 changed files with 146 additions and 201 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue