test pod lock manager

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