test(proxy): verify cancellation progress and lock release via prisma seam (#39730)

tests/test_litellm/proxy/test_spend_log_cleanup.py

@@ -1424,45 +1424,20 @@ async def test_cancelled_cleanup_records_and_logs_aborted_run(monkeypatch):
     """
     When cleanup_old_spend_logs is cancelled (e.g. during pod shutdown),
     asyncio.CancelledError must record outcome='aborted', emit an error-level
-    diagnostic containing progress metrics, and re-raise CancelledError.
+    diagnostic containing progress metrics (including rows deleted and batches
+    completed prior to cancellation), and re-raise CancelledError.
     """
     import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
 
-    cleaner = SpendLogCleanup(
-        general_settings={"maximum_spend_logs_retention_period": "7d"}
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    # First batch deletes 150 rows. Second batch simulates cancellation mid-run.
+    mock_db.execute_raw = AsyncMock(
+        side_effect=[150, asyncio.CancelledError("simulated shutdown mid-run")]
     )
-    cleaner.pod_lock_manager = None
-    cleaner._clean_spend_log_tables = AsyncMock(
-        side_effect=asyncio.CancelledError("simulated shutdown")
-    )
-    prisma_client = MagicMock()
-    logger = MagicMock()
-    recorded_runs: list[str] = []
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db = mock_db
 
-    monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", logger)
-    monkeypatch.setattr(
-        cleanup_module.SpendLogCleanupMetrics, "record_run", lambda outcome: recorded_runs.append(outcome)
-    )
-
-    with pytest.raises(asyncio.CancelledError):
-        await cleaner.cleanup_old_spend_logs(prisma_client)
-
-    assert recorded_runs == ["aborted"]
-    logger.exception.assert_called_once()
-    log_call = logger.exception.call_args[0]
-    rendered_log = log_call[0] % log_call[1:]
-    assert "cancelled" in rendered_log.lower()
-    assert "elapsed=" in rendered_log
-    assert "rows_deleted=" in rendered_log
-    assert "batches=" in rendered_log
-
-
-@pytest.mark.asyncio
-async def test_cancelled_cleanup_releases_lock():
-    """
-    When a cleanup run holds a distributed pod lock and is cancelled,
-    the finally block must release the lock before propagating CancelledError.
-    """
     mock_lock_manager = MagicMock()
     mock_lock_manager.redis_cache = MagicMock()
     mock_lock_manager.acquire_lock = AsyncMock(return_value=True)
@@ -1472,13 +1447,58 @@ async def test_cancelled_cleanup_releases_lock():
         general_settings={"maximum_spend_logs_retention_period": "7d"}
     )
     cleaner.pod_lock_manager = mock_lock_manager
-    cleaner._clean_spend_log_tables = AsyncMock(
-        side_effect=asyncio.CancelledError("shutdown")
+
+    logger = MagicMock()
+    recorded_runs: list[str] = []
+
+    monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", logger)
+    monkeypatch.setattr(
+        cleanup_module.SpendLogCleanupMetrics, "record_run", lambda outcome: recorded_runs.append(outcome)
     )
-    prisma_client = MagicMock()
 
     with pytest.raises(asyncio.CancelledError):
-        await cleaner.cleanup_old_spend_logs(prisma_client)
+        await cleaner.cleanup_old_spend_logs(mock_prisma_client)
+
+    assert recorded_runs == ["aborted"]
+    assert cleaner._run_rows_deleted == 150
+    assert cleaner._run_batches == 1
+    mock_lock_manager.release_lock.assert_awaited_once()
+
+    logger.exception.assert_called_once()
+    log_call = logger.exception.call_args[0]
+    rendered_log = log_call[0] % log_call[1:]
+    assert "cancelled" in rendered_log.lower()
+    assert "elapsed=" in rendered_log
+    assert "rows_deleted=150" in rendered_log
+    assert "batches=1" in rendered_log
+
+
+@pytest.mark.asyncio
+async def test_cancelled_cleanup_releases_lock():
+    """
+    When a cleanup run holds a distributed pod lock and is cancelled,
+    the finally block must release the lock before propagating CancelledError.
+    """
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(
+        side_effect=asyncio.CancelledError("simulated immediate shutdown")
+    )
+    mock_prisma_client = MagicMock()
+    mock_prisma_client.db = mock_db
+
+    mock_lock_manager = MagicMock()
+    mock_lock_manager.redis_cache = MagicMock()
+    mock_lock_manager.acquire_lock = AsyncMock(return_value=True)
+    mock_lock_manager.release_lock = AsyncMock()
+
+    cleaner = SpendLogCleanup(
+        general_settings={"maximum_spend_logs_retention_period": "7d"}
+    )
+    cleaner.pod_lock_manager = mock_lock_manager
+
+    with pytest.raises(asyncio.CancelledError):
+        await cleaner.cleanup_old_spend_logs(mock_prisma_client)
 
     mock_lock_manager.release_lock.assert_awaited_once()