From 193052ed7054292e7d01794cf3b3fb76bbb83503 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 28 Mar 2025 15:05:17 -0700 Subject: [PATCH] test pod lock manager --- .../litellm/proxy/db/test_pod_lock_manager.py | 215 ++++++++++++++++-- 1 file changed, 200 insertions(+), 15 deletions(-) diff --git a/tests/litellm/proxy/db/test_pod_lock_manager.py b/tests/litellm/proxy/db/test_pod_lock_manager.py index e685172c1ed..bce7b66409e 100644 --- a/tests/litellm/proxy/db/test_pod_lock_manager.py +++ b/tests/litellm/proxy/db/test_pod_lock_manager.py @@ -42,38 +42,75 @@ def pod_lock_manager(): @pytest.mark.asyncio async def test_acquire_lock_success(pod_lock_manager, mock_prisma): """ - Test that the lock is acquired successfully if the DB response is successful + 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.upsert.return_value = mock_response + mock_prisma.db.litellm_cronjob.create.return_value = mock_response result = await pod_lock_manager.acquire_lock() assert result == True - # Verify upsert was called with correct parameters - mock_prisma.db.litellm_cronjob.upsert.assert_called_once() - call_args = mock_prisma.db.litellm_cronjob.upsert.call_args[1] - assert call_args["where"]["cronjob_id"] == "test_job" - assert "create" in call_args["data"] - assert "update" in call_args["data"] + # 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" @pytest.mark.asyncio -async def test_acquire_lock_failure(pod_lock_manager, mock_prisma): +async def test_acquire_lock_existing_active(pod_lock_manager, mock_prisma): """ - Test that the lock is not acquired if the lock is held by a different pod + Test that the lock is not acquired if there's an active lock by different pod """ - # Mock failed lock acquisition (different pod holds the lock) - mock_response = AsyncMock() - mock_response.status = "ACTIVE" - mock_response.pod_id = "different_pod_id" - mock_prisma.db.litellm_cronjob.upsert.return_value = mock_response + # 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 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() + + +@pytest.mark.asyncio +async def test_acquire_lock_expired(pod_lock_manager, mock_prisma): + """ + 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 successful update + mock_updated = AsyncMock() + mock_updated.pod_id = pod_lock_manager.pod_id + mock_prisma.db.litellm_cronjob.update.return_value = mock_updated + + 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() + @pytest.mark.asyncio async def test_renew_lock(pod_lock_manager, mock_prisma): @@ -133,3 +170,151 @@ async def test_database_error_handling(pod_lock_manager, mock_prisma): 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 + + +@pytest.mark.asyncio +async def test_acquire_lock_inactive_status(pod_lock_manager, mock_prisma): + """ + Test that the lock can be acquired if existing lock is INACTIVE + """ + # 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 + + # 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 + + result = await pod_lock_manager.acquire_lock() + assert result == True + + mock_prisma.db.litellm_cronjob.update.assert_called_once() + + +@pytest.mark.asyncio +async def test_acquire_lock_same_pod(pod_lock_manager, mock_prisma): + """ + Test that the lock returns True if the same pod already holds the lock + """ + # 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" + ) + + result = await pod_lock_manager.acquire_lock() + assert result == False + + +@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 + + +@pytest.mark.asyncio +async def test_concurrent_lock_acquisition_simulation(mock_prisma): + """ + 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") + + # 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 + + # 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 + + # Other pods should fail to acquire + result2 = await pod2.acquire_lock() + result3 = await pod3.acquire_lock() + assert result2 == False + assert result3 == False + + +@pytest.mark.asyncio +async def test_lock_takeover_race_condition(mock_prisma): + """ + Test scenario where multiple pods try to take over an expired lock + """ + pod1 = PodLockManager(cronjob_id="test_job") + pod2 = PodLockManager(cronjob_id="test_job") + + # 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 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 + 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" + ) + + # Second pod should fail to take over + result2 = await pod2.acquire_lock() + assert result2 == False