litellm/tests/unit/proxy/test_spend_log_cleanup.py

"""
Test cases for spend log cleanup functionality
"""

import asyncio
import logging
import math
import time
from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock

import pytest

from litellm.constants import (
    SPEND_LOG_CLEANUP_BATCH_SIZE,
    SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
    SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
    SPEND_LOG_CLEANUP_BOUND_SETTINGS,
    SizeCleanupState,
    SpendLogCleanup,
    SpendLogStorage,
    StorageProbe,
    TableCleanupResult,
    advance_size_state,
)
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
    SpendLogCleanupMetrics,
)
from tests.unit.proxy.db.fake_prisma_engine import engine_call


def _far_deadline() -> float:
    """A run deadline far enough out that only the other bounds can stop a batch loop."""
    return time.monotonic() + 3600


def _wire_tx(db, rejected_timeout=None):
    """
    Model the prisma seam the cleanup job actually uses.

    Every statement the job issues runs inside db.tx() so it can carry a SET
    LOCAL statement_timeout. Batch and probe statements are forwarded to
    db.execute_raw and db.query_raw, which is what tests configure and assert
    on, while the SET LOCAL statements are answered here so they neither consume
    a side_effect entry nor show up in the recorded call list. Lookup is
    deferred to call time so this can be wired before a test assigns its own
    execute_raw.
    """

    @asynccontextmanager
    async def _tx(**kwargs):
        if kwargs.get("timeout") == rejected_timeout:
            raise TypeError("transaction timeout rejected by fake")
        tx = MagicMock()

        async def _execute_raw(sql, *args):
            if sql.lstrip().upper().startswith("SET LOCAL"):
                return 0
            return await db.execute_raw(sql, *args)

        async def _query_raw(sql, *args):
            return await db.query_raw(sql, *args)

        tx.execute_raw = _execute_raw
        tx.query_raw = _query_raw
        yield tx

    db.tx = _tx
    db.query_raw = AsyncMock(return_value=[{"remaining": 0}])


def test_spend_log_cleanup_cron_scheduling():
    """Test that cron expressions are correctly parsed for spend log cleanup scheduling"""
    from apscheduler.triggers.cron import CronTrigger

    # Valid cron expressions
    cron_expr = "0 4 * * *"  # 4:00 AM daily
    trigger = CronTrigger.from_crontab(cron_expr)
    assert trigger is not None

    # Every minute (useful for testing)
    trigger_minute = CronTrigger.from_crontab("*/1 * * * *")
    assert trigger_minute is not None

    # Specific day and hour
    trigger_weekly = CronTrigger.from_crontab("0 3 * * 0")  # 3 AM every Sunday
    assert trigger_weekly is not None

    # Invalid cron expression should raise ValueError
    with pytest.raises(ValueError, match="Wrong number of fields; got"):
        CronTrigger.from_crontab("invalid cron")

    with pytest.raises(ValueError, match="is higher than the maximum value"):
        CronTrigger.from_crontab("60 25 * * *")  # Invalid minute and hour


def test_spend_log_cleanup_cron_scheduler_integration():
    """
    Integration test: Verify the proxy_server scheduler logic correctly adds
    cron-based cleanup job when maximum_spend_logs_cleanup_cron is configured.

    This tests the logic in proxy_server.py lines 4671-4717 without requiring
    a real database connection.
    """
    from unittest.mock import MagicMock

    from apscheduler.triggers.cron import CronTrigger

    # Mock scheduler
    mock_scheduler = MagicMock()
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_cleanup_instance = MagicMock()

    # Test Case 1: Cron-based scheduling
    general_settings_cron = {
        "maximum_spend_logs_retention_period": "7d",
        "maximum_spend_logs_cleanup_cron": "0 4 * * *",  # 4 AM daily
    }

    cleanup_cron = general_settings_cron.get("maximum_spend_logs_cleanup_cron")
    assert cleanup_cron is not None

    # Simulate the scheduler logic from proxy_server.py
    cron_trigger = CronTrigger.from_crontab(cleanup_cron)
    mock_scheduler.add_job(
        mock_cleanup_instance.cleanup_old_spend_logs,
        cron_trigger,
        args=[mock_prisma_client],
        id="spend_log_cleanup_job",
        replace_existing=True,
        misfire_grace_time=3600,
    )

    # Verify scheduler was called correctly
    mock_scheduler.add_job.assert_called_once()
    call_args = mock_scheduler.add_job.call_args

    # Verify the trigger is a CronTrigger
    assert isinstance(call_args[0][1], CronTrigger)

    # Verify job ID
    assert call_args[1]["id"] == "spend_log_cleanup_job"
    assert call_args[1]["replace_existing"] is True

    # Test Case 2: Interval-based scheduling (fallback)
    mock_scheduler.reset_mock()
    general_settings_interval = {
        "maximum_spend_logs_retention_period": "7d",
        # No cron, so it should fall back to interval
    }

    cleanup_cron_fallback = general_settings_interval.get("maximum_spend_logs_cleanup_cron")
    assert cleanup_cron_fallback is None  # No cron configured

    # Simulate interval-based scheduling fallback
    retention_interval = general_settings_interval.get("maximum_spend_logs_retention_interval", "1d")
    from litellm.litellm_core_utils.duration_parser import duration_in_seconds

    interval_seconds = duration_in_seconds(retention_interval)

    mock_scheduler.add_job(
        mock_cleanup_instance.cleanup_old_spend_logs,
        "interval",
        seconds=interval_seconds,
        args=[mock_prisma_client],
        id="spend_log_cleanup_job",
        replace_existing=True,
    )

    # Verify interval scheduling was called
    mock_scheduler.add_job.assert_called_once()
    interval_call_args = mock_scheduler.add_job.call_args
    assert interval_call_args[0][1] == "interval"
    assert interval_call_args[1]["seconds"] == 86400  # 1 day in seconds


@pytest.mark.asyncio
async def test_should_delete_spend_logs():
    # Test case 1: No retention set
    cleaner = SpendLogCleanup(general_settings={})
    assert cleaner._should_delete_spend_logs() is False

    # Test case 2: Valid seconds string
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "3600s"})
    assert cleaner._should_delete_spend_logs() is True

    # Test case 3: Valid days string
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
    assert cleaner._should_delete_spend_logs() is True

    # Test case 4: Valid hours string
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "24h"})
    assert cleaner._should_delete_spend_logs() is True

    # Test case 5: Invalid format
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "invalid"})
    assert cleaner._should_delete_spend_logs() is False


@pytest.mark.asyncio
async def test_delete_batch_transaction_timeout_covers_statement_timeout() -> None:
    transaction: Final = MagicMock()
    transaction.execute_raw = AsyncMock(side_effect=[0, 0, 3])

    @asynccontextmanager
    async def tx(timeout: timedelta):
        yield transaction

    prisma_client: Final = MagicMock()
    prisma_client.db.tx = MagicMock(side_effect=tx)
    cleaner: Final = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
    cleaner.batch_timeout_seconds = 12

    await cleaner._execute_delete_batch(
        prisma_client, "DELETE FROM table", datetime.now(), "LiteLLM_SpendLogs", _far_deadline()
    )

    assert prisma_client.db.tx.call_args.kwargs["timeout"] >= timedelta(seconds=12)


@pytest.mark.asyncio
async def test_count_remaining_transaction_timeout_covers_statement_timeout() -> None:
    transaction: Final = MagicMock()
    transaction.execute_raw = AsyncMock(return_value=0)
    transaction.query_raw = AsyncMock(return_value=[{"remaining": 7}])

    @asynccontextmanager
    async def tx(timeout: timedelta):
        yield transaction

    prisma_client: Final = MagicMock()
    prisma_client.db.tx = MagicMock(side_effect=tx)
    cleaner: Final = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "30d"})
    cleaner.batch_timeout_seconds = 12

    result: Final = await cleaner._count_remaining(
        prisma_client, datetime.now(), "LiteLLM_SpendLogs", "startTime", _far_deadline()
    )

    assert result == 7
    assert prisma_client.db.tx.call_args.kwargs["timeout"] >= timedelta(seconds=12)


@pytest.mark.asyncio
async def test_cleanup_old_spend_logs_batch_deletion():
    from unittest.mock import AsyncMock, MagicMock

    # Setup Prisma client
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)

    # Mock execute_raw to return deleted counts (3 spend-log batches, then the
    # tool-index cleanup's first batch returning 0)
    mock_db.execute_raw = AsyncMock(side_effect=[1000, 500, 0, 0])

    # Wire up mocks
    mock_prisma_client.db = mock_db

    # Mock Redis cache and pod_lock_manager
    mock_redis_cache = MagicMock()
    mock_pod_lock_manager = MagicMock()
    mock_pod_lock_manager.redis_cache = mock_redis_cache
    mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
    mock_pod_lock_manager.release_lock = AsyncMock()

    # Run cleanup with mocked pod_lock_manager
    test_settings = {"maximum_spend_logs_retention_period": "7d"}
    cleaner = SpendLogCleanup(general_settings=test_settings)
    cleaner.pod_lock_manager = mock_pod_lock_manager
    assert cleaner._should_delete_spend_logs() is True
    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    # Validate batching and deletion via raw SQL
    assert mock_db.execute_raw.call_count == 4

    # Check the first call argument
    call_args_sql = mock_db.execute_raw.call_args_list[0][0][0]
    assert 'DELETE FROM "LiteLLM_SpendLogs"' in call_args_sql
    # must match on the full composite identity: on a partitioned table
    # request_id alone is not unique, and deleting by it would let a client
    # reusing x-litellm-call-id take out a fresh row alongside the expired one
    assert 'WHERE ("request_id", "startTime") IN' in call_args_sql

    # After spend logs, the derived tool index rows expire on the same cutoff
    tool_index_sql = mock_db.execute_raw.call_args_list[3][0][0]
    assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in tool_index_sql

    # The LiteLLM_DailyToolSpend rollup must outlive spend-log retention: it is
    # the only copy of tool spend history once its per-request sources expire,
    # so spend-log cleanup must never touch it.
    for call in mock_db.execute_raw.call_args_list:
        assert "LiteLLM_DailyToolSpend" not in call[0][0]


@pytest.mark.asyncio
async def test_cleanup_old_spend_logs_retention_period_cutoff():
    """
    Test that logs are filtered using correct cutoff based on retention
    """
    # Setup Prisma client
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=0)
    mock_prisma_client.db = mock_db

    # Mock Redis cache and pod_lock_manager
    mock_redis_cache = MagicMock()
    mock_pod_lock_manager = MagicMock()
    mock_pod_lock_manager.redis_cache = mock_redis_cache
    mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
    mock_pod_lock_manager.release_lock = AsyncMock()

    # Run cleanup with mocked pod_lock_manager
    test_settings = {"maximum_spend_logs_retention_period": "24h"}
    cleaner = SpendLogCleanup(general_settings=test_settings)
    cleaner.pod_lock_manager = mock_pod_lock_manager
    assert cleaner._should_delete_spend_logs() is True
    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    # Verify the cutoff date is correct
    cutoff_date = mock_db.execute_raw.call_args[0][1]
    expected_cutoff = datetime.now(timezone.utc) - timedelta(seconds=86400)
    assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1  # Allow 1 second difference for test execution time


@pytest.mark.asyncio
async def test_cleanup_drops_partitions_when_enabled_and_partitioned():
    """
    With use_spend_logs_partitioning enabled and a partitioned table, cleanup
    must reclaim disk by dropping partitions AND still delete expired rows the
    drops cannot reach (DEFAULT partition, cutoff-spanning partitions), so
    retention is never bypassed.
    """
    from unittest.mock import AsyncMock, MagicMock

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_prisma_client.db.execute_raw = AsyncMock(return_value=0)

    partition_manager = MagicMock()
    partition_manager.is_partitioned = AsyncMock(return_value=True)
    partition_manager.ensure_partitions = AsyncMock(return_value=["p1"])
    partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"])

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "use_spend_logs_partitioning": True,
        },
        partition_manager=partition_manager,
    )
    cleaner.pod_lock_manager = MagicMock()
    cleaner.pod_lock_manager.redis_cache = None

    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    partition_manager.ensure_partitions.assert_awaited_once()
    partition_manager.drop_partitions_older_than.assert_awaited_once()
    delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0]
    assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql
    # Partition drops only reclaim spend logs; the tool index must still be
    # cleaned row-wise on the same run
    all_sql = [c[0][0] for c in mock_prisma_client.db.execute_raw.call_args_list]
    assert any('DELETE FROM "LiteLLM_SpendLogToolIndex"' in s for s in all_sql)


@pytest.mark.asyncio
async def test_cleanup_uses_delete_when_partitioning_not_enabled():
    """
    Even against a partitioned table, the partition path must stay off until
    use_spend_logs_partitioning is explicitly enabled, so existing deployments
    see zero behavior change. The catalog must not even be queried.
    """
    from unittest.mock import AsyncMock, MagicMock

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])

    partition_manager = MagicMock()
    partition_manager.is_partitioned = AsyncMock(return_value=True)
    partition_manager.ensure_partitions = AsyncMock()
    partition_manager.drop_partitions_older_than = AsyncMock()

    cleaner = SpendLogCleanup(
        general_settings={"maximum_spend_logs_retention_period": "7d"},
        partition_manager=partition_manager,
    )
    cleaner.pod_lock_manager = MagicMock()
    cleaner.pod_lock_manager.redis_cache = None

    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    partition_manager.is_partitioned.assert_not_awaited()
    partition_manager.drop_partitions_older_than.assert_not_awaited()
    delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0]
    assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql


@pytest.mark.asyncio
async def test_cleanup_uses_delete_when_not_partitioned():
    """
    With the feature enabled but the table not actually partitioned (script not
    run yet), cleanup must keep using the batched DELETE path.
    """
    from unittest.mock import AsyncMock, MagicMock

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])

    partition_manager = MagicMock()
    partition_manager.is_partitioned = AsyncMock(return_value=False)
    partition_manager.drop_partitions_older_than = AsyncMock()

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "use_spend_logs_partitioning": True,
        },
        partition_manager=partition_manager,
    )
    cleaner.pod_lock_manager = MagicMock()
    cleaner.pod_lock_manager.redis_cache = None

    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    partition_manager.drop_partitions_older_than.assert_not_awaited()
    assert mock_prisma_client.db.execute_raw.await_count == 3
    delete_sql = mock_prisma_client.db.execute_raw.call_args_list[0][0][0]
    assert 'DELETE FROM "LiteLLM_SpendLogs"' in delete_sql


@pytest.mark.asyncio
async def test_cleanup_old_spend_logs_no_retention_period():
    """
    Test that no logs are deleted when no retention period is set
    """
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_prisma_client.db.execute_raw = AsyncMock()

    cleaner = SpendLogCleanup(general_settings={})  # no retention
    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    mock_prisma_client.db.execute_raw.assert_not_called()


@pytest.mark.asyncio
async def test_lock_not_released_when_not_acquired():
    """
    Lock release should be skipped when _should_delete_spend_logs returns False
    before the lock is ever acquired.
    """
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_prisma_client.db.execute_raw = AsyncMock()

    mock_redis_cache = MagicMock()
    mock_pod_lock_manager = MagicMock()
    mock_pod_lock_manager.redis_cache = mock_redis_cache
    mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
    mock_pod_lock_manager.release_lock = AsyncMock()

    # No retention setting → _should_delete_spend_logs() returns False before lock is acquired
    cleaner = SpendLogCleanup(general_settings={})
    cleaner.pod_lock_manager = mock_pod_lock_manager

    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    mock_pod_lock_manager.acquire_lock.assert_not_called()
    mock_pod_lock_manager.release_lock.assert_not_called()


@pytest.mark.asyncio
async def test_integer_retention_treated_as_days():
    """
    An integer value for maximum_spend_logs_retention_period should be treated
    as days (e.g., 3 → '3d' → 259200 seconds).
    """
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": 3})
    result = cleaner._should_delete_spend_logs()
    assert result is True
    assert cleaner.retention_seconds == 3 * 86400  # 3 days in seconds


def test_string_retention_still_works():
    """
    String values like '3d', '24h', '3600s' should continue to parse correctly.
    """
    cases = [
        ("3d", 3 * 86400),
        ("24h", 24 * 3600),
        ("3600s", 3600),
        ("2w", 2 * 604800),
    ]
    for setting, expected_seconds in cases:
        cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": setting})
        assert cleaner._should_delete_spend_logs() is True, f"Failed for {setting}"
        assert cleaner.retention_seconds == expected_seconds, (
            f"Expected {expected_seconds} for {setting}, got {cleaner.retention_seconds}"
        )


@pytest.mark.asyncio
async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
    """should abort deletion loop immediately when execute_raw returns a non-int
    (e.g. None or dict), preventing an infinite loop."""
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=None)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())

    assert mock_db.execute_raw.call_count == 1
    assert result.rows_deleted == 0


@pytest.mark.asyncio
async def test_delete_old_logs_continues_on_valid_int_return():
    """should continue deletion loop across batches when execute_raw returns valid int counts."""
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0])
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())

    assert mock_db.execute_raw.call_count == 3
    assert result.rows_deleted == 800


@pytest.mark.asyncio
async def test_delete_old_rows_stops_at_max_batches():
    """The batch cap must halt a cleanup that keeps finding rows, so a huge
    backlog is spread across scheduled runs instead of one unbounded loop, and
    the operator-facing knob must mean exactly the number of statements it names."""
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=1000)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_max_batches": 2,
        }
    )

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())

    assert mock_db.execute_raw.call_count == 2
    assert result.rows_deleted == 2000
    assert result.stop_reason == "batch_cap_reached"


@pytest.mark.asyncio
async def test_delete_old_tool_index_rows_deletes_on_composite_key():
    """Tool index rows are derived from spend logs and expire on the same cutoff;
    the delete must match on the table's composite primary key."""
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(side_effect=[5, 0])
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline())

    assert result.rows_deleted == 5
    delete_sql = mock_db.execute_raw.call_args_list[0][0][0]
    assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in delete_sql
    assert 'WHERE ("request_id", "tool_name") IN' in delete_sql
    assert '"start_time" <' in delete_sql
    assert mock_db.execute_raw.call_args_list[0][0][1] == cutoff_date


@pytest.mark.asyncio
async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch):
    """A single batch failure (e.g. DB timeout) must not abort the whole run —
    subsequent batches should still execute and their counts accumulate."""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    # Zero out the failure backoff so the test doesn't take ~0.5s of real sleep.
    monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed,
    # batch 5 returns 0 → loop exits naturally.
    mock_db.execute_raw = AsyncMock(side_effect=[100, TimeoutError("simulated DB timeout"), 200, 50, 0])
    mock_prisma_client.db = mock_db

    cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())

    # All 5 batches should have been attempted; 100 + 200 + 50 = 350 deleted.
    assert mock_db.execute_raw.call_count == 5
    assert result.rows_deleted == 350


@pytest.mark.asyncio
async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch):
    """If batch failures persist for SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES
    in a row (e.g. DB is down), the loop must abort instead of hot-looping."""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    # Lower the threshold so the test is fast and deterministic.
    monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
    monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    # Every batch raises — must abort after exactly 3 attempts, not loop forever.
    mock_db.execute_raw = AsyncMock(side_effect=ConnectionError("simulated persistent DB outage"))
    mock_prisma_client.db = mock_db

    cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())

    assert mock_db.execute_raw.call_count == 3
    assert result.rows_deleted == 0


@pytest.mark.asyncio
async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatch):
    """A success between failures must reset the consecutive-failure counter so
    intermittent timeouts don't trip the abort threshold."""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)
    monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    # Pattern: fail, fail, success (resets counter), fail, fail, success, done.
    # Without reset, three of these would trip abort; with reset, they don't.
    mock_db.execute_raw = AsyncMock(
        side_effect=[
            TimeoutError("t1"),
            TimeoutError("t2"),
            100,
            TimeoutError("t3"),
            TimeoutError("t4"),
            50,
            0,
        ]
    )
    mock_prisma_client.db = mock_db

    cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())

    assert mock_db.execute_raw.call_count == 7
    assert result.rows_deleted == 150


@pytest.mark.asyncio
async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch):
    """The outer error handler must call logger.exception() (not .error(str(e)))
    so Prisma/DB timeouts surface a full traceback and exception type."""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    mock_logger = MagicMock()
    monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    # Force the outer try/except to fire by making _should_delete_spend_logs raise.
    cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = None

    def boom():
        raise RuntimeError("simulated prisma timeout")

    cleaner._should_delete_spend_logs = boom  # type: ignore[assignment]

    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    assert mock_logger.exception.called, "expected logger.exception() to be called"
    # The exception type name must appear in the formatted args so operators can
    # tell *what* failed, not just "Error during cleanup:".
    call_args = mock_logger.exception.call_args
    formatted = call_args[0][0] % call_args[0][1:]
    assert "RuntimeError" in formatted
    assert "simulated prisma timeout" in formatted


@pytest.mark.asyncio
async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch):
    """Even when batch deletion aborts due to consecutive failures, the pod lock
    must still be released so the next scheduled run isn't permanently blocked."""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 2)
    monkeypatch.setattr(cleanup_module, "SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.0)

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
    mock_prisma_client.db = mock_db

    mock_pod_lock_manager = MagicMock()
    mock_pod_lock_manager.redis_cache = MagicMock()
    mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
    mock_pod_lock_manager.release_lock = AsyncMock()

    cleaner = cleanup_module.SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = mock_pod_lock_manager

    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    # Cleanup didn't crash; the abort-after-failures path returned cleanly.
    mock_pod_lock_manager.release_lock.assert_awaited_once()


def test_cleanup_batch_size_env_var(monkeypatch):
    """Ensure batch size is configurable via environment variable"""
    import importlib

    import litellm.constants as constants_module
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    # Set env var and reload modules to pick up new value
    monkeypatch.setenv("SPEND_LOG_CLEANUP_BATCH_SIZE", "25")
    importlib.reload(constants_module)
    importlib.reload(cleanup_module)

    cleaner = cleanup_module.SpendLogCleanup(general_settings={})
    assert cleaner.batch_size == 25

    # Remove env var and reload to restore default for other tests
    monkeypatch.delenv("SPEND_LOG_CLEANUP_BATCH_SIZE", raising=False)
    importlib.reload(constants_module)
    importlib.reload(cleanup_module)


def _mock_prisma_for_retention(side_effect: list) -> "MagicMock":
    from unittest.mock import AsyncMock, MagicMock

    client = MagicMock()
    _wire_tx(client.db, rejected_timeout=timedelta(seconds=10))
    client.db.execute_raw = AsyncMock(side_effect=side_effect)
    return client


@pytest.mark.asyncio
async def test_spend_logs_retention_alone_does_not_touch_the_session_rollup():
    client = _mock_prisma_for_retention([0, 0])
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(client)
    tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
    assert any('"LiteLLM_SpendLogs"' in sql for sql in tables)
    assert not any('"LiteLLM_AutoRouterSession"' in sql for sql in tables)
    assert not any('"LiteLLM_AutoRouterUserSession"' in sql for sql in tables)
    assert not any('"LiteLLM_AutoRouterDailySpend"' in sql for sql in tables)
    assert not any('"LiteLLM_HealthCheckTable"' in sql for sql in tables)


@pytest.mark.asyncio
async def test_session_retention_alone_cleans_both_session_rollups_and_the_daily_rollup():
    client = _mock_prisma_for_retention([0, 0, 0])
    cleaner = SpendLogCleanup(general_settings={"maximum_autorouter_session_retention_period": "365d"})
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(client)
    calls = client.db.execute_raw.call_args_list
    tables = [call[0][0] for call in calls]
    assert len(tables) == 3
    assert '"LiteLLM_AutoRouterSession"' in tables[0]
    assert '"LiteLLM_AutoRouterUserSession"' in tables[1]
    assert '"LiteLLM_AutoRouterDailySpend"' in tables[2]
    assert calls[2][0][1] == calls[0][0][1].date().isoformat()


@pytest.mark.asyncio
async def test_health_check_retention_alone_cleans_only_the_health_check_table():
    client = _mock_prisma_for_retention([0])
    cleaner = SpendLogCleanup(general_settings={"maximum_health_check_retention_period": "30d"})
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(client)
    tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
    assert len(tables) == 1
    assert '"LiteLLM_HealthCheckTable"' in tables[0]
    assert '"health_check_id"' in tables[0]
    assert '"checked_at"' in tables[0]
    cutoff_date = client.db.execute_raw.call_args[0][1]
    expected_cutoff = datetime.now(timezone.utc) - timedelta(days=30)
    assert abs((cutoff_date - expected_cutoff).total_seconds()) < 1


@pytest.mark.asyncio
async def test_daily_tag_spend_retention_alone_prunes_only_that_table_by_calendar_day():
    client = _mock_prisma_for_retention([0])
    cleaner = SpendLogCleanup(general_settings={"maximum_daily_tag_spend_retention_period": "90d"})
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(client)
    tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
    assert len(tables) == 1
    assert '"LiteLLM_DailyTagSpend"' in tables[0]
    cutoff_day = client.db.execute_raw.call_args[0][1]
    assert cutoff_day == (datetime.now(timezone.utc) - timedelta(days=90)).date().isoformat()


@pytest.mark.asyncio
async def test_spend_logs_retention_alone_keeps_daily_tag_spend_forever():
    client = _mock_prisma_for_retention([0, 0])
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(client)
    tables = [call[0][0] for call in client.db.execute_raw.call_args_list]
    assert not any('"LiteLLM_DailyTagSpend"' in sql for sql in tables)


@pytest.mark.asyncio
async def test_each_retention_key_cuts_off_at_its_own_horizon():
    client = _mock_prisma_for_retention([0, 0, 0, 0, 0, 0])
    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_autorouter_session_retention_period": "365d",
            "maximum_health_check_retention_period": "30d",
        }
    )
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(client)
    cutoffs = {
        (
            "LiteLLM_AutoRouterSession"
            if '"LiteLLM_AutoRouterSession"' in call[0][0]
            else "LiteLLM_AutoRouterUserSession"
            if '"LiteLLM_AutoRouterUserSession"' in call[0][0]
            else "LiteLLM_AutoRouterDailySpend"
            if '"LiteLLM_AutoRouterDailySpend"' in call[0][0]
            else "LiteLLM_HealthCheckTable"
            if '"LiteLLM_HealthCheckTable"' in call[0][0]
            else "logs"
        ): call[0][1]
        for call in client.db.execute_raw.call_args_list
    }
    now = datetime.now(timezone.utc)
    assert (now - cutoffs["logs"]).days == 7
    assert (now - cutoffs["LiteLLM_AutoRouterSession"]).days == 365
    assert cutoffs["LiteLLM_AutoRouterUserSession"] == cutoffs["LiteLLM_AutoRouterSession"]
    assert cutoffs["LiteLLM_AutoRouterDailySpend"] == cutoffs["LiteLLM_AutoRouterSession"].date().isoformat()
    assert (now - cutoffs["LiteLLM_HealthCheckTable"]).days == 30


@pytest.mark.asyncio
async def test_no_retention_keys_means_no_cleanup_at_all():
    client = _mock_prisma_for_retention([])
    cleaner = SpendLogCleanup(general_settings={})
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(client)
    assert client.db.execute_raw.await_count == 0


@pytest.mark.asyncio
async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run():
    """
    The wall-clock budget is the bound that keeps a large backlog from turning
    into one multi-hour run. With rows always available, the loop must stop on
    the deadline rather than on the batch cap, and must report that reason so
    operators can tell a budgeted stop from a drained table.
    """
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=1000)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            # Comfortably more batches than a sub-second budget can reach (each
            # batch sleeps 0.1s), but small enough that a broken deadline fails
            # this test in seconds instead of hanging it
            "maximum_spend_logs_cleanup_max_batches": 50,
        }
    )

    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
    started_at = time.monotonic()
    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 1.0)
    elapsed = time.monotonic() - started_at

    assert result.stop_reason == "budget_exhausted"
    assert elapsed < 3, f"budgeted run overran its deadline: {elapsed}s"
    assert mock_db.execute_raw.call_count < 50
    assert result.rows_deleted > 0


@pytest.mark.asyncio
async def test_run_budget_is_shared_across_tables_not_granted_per_table():
    """
    A per-table budget would let a run take N times the configured bound. The
    deadline is computed once per run, so once it is spent on the first table
    the later tables must stop immediately rather than each getting a fresh one.
    """
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=1000)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_autorouter_session_retention_period": "365d",
            # Comfortably more batches than a sub-second budget can reach (each
            # batch sleeps 0.1s), but small enough that a broken deadline fails
            # this test in seconds instead of hanging it
            "maximum_spend_logs_cleanup_max_batches": 50,
            "maximum_spend_logs_cleanup_run_budget": "1s",
        }
    )
    cleaner.pod_lock_manager = None

    started_at = time.monotonic()
    await cleaner.cleanup_old_spend_logs(mock_prisma_client)
    elapsed = time.monotonic() - started_at

    # three tables are eligible; a per-table budget would push this past 3s
    assert elapsed < 2.5, f"budget was granted per table, not per run: {elapsed}s"
    tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list}
    assert "LiteLLM_SpendLogs" in tables_touched


@pytest.mark.asyncio
async def test_cleanup_groups_share_budget_so_health_checks_still_get_a_delete():
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=1000)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_health_check_retention_period": "30d",
            "maximum_spend_logs_cleanup_max_batches": 50,
            "maximum_spend_logs_cleanup_run_budget": "1s",
        }
    )
    cleaner.pod_lock_manager = None

    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list}
    assert "LiteLLM_SpendLogs" in tables_touched
    assert "LiteLLM_HealthCheckTable" in tables_touched


@pytest.mark.asyncio
async def test_each_batch_carries_a_statement_and_lock_timeout():
    """
    A Prisma transaction timeout cannot interrupt a statement already running,
    so the Postgres statement_timeout and lock_timeout are the only things
    stopping one batch from holding row locks and a pooled connection
    indefinitely. Both must be set, inside the batch's own transaction, and
    scoped with SET LOCAL so the pooled connection is left unchanged.
    """
    recorded: list[str] = []

    mock_prisma_client = MagicMock()
    mock_db = MagicMock()

    @asynccontextmanager
    async def _tx(**kwargs):
        tx = MagicMock()

        async def _execute_raw(sql, *args):
            recorded.append(sql.strip())
            return 0

        tx.execute_raw = _execute_raw
        yield tx

    mock_db.tx = _tx
    mock_db.query_raw = AsyncMock(return_value=[{"remaining": 0}])
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_batch_timeout": "12s",
        }
    )

    await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline())

    assert "SET LOCAL statement_timeout = 12000" in recorded
    assert "SET LOCAL lock_timeout = 12000" in recorded
    # the timeouts must precede the delete they are meant to bound
    assert recorded.index("SET LOCAL statement_timeout = 12000") < next(
        i for i, sql in enumerate(recorded) if sql.startswith("DELETE")
    )


@pytest.mark.parametrize(
    "setting_value",
    ["inf", "-inf", "nan", "1e400", "0s", "-5m", "not-a-duration"],
)
def test_a_non_finite_or_non_positive_budget_falls_back_to_the_default(setting_value):
    """
    The knob must not be able to remove the bound it exists to enforce.

    'inf', 'nan' and '1e400' are the spellings that would turn the deadline
    into no deadline at all, and '0s' and '-5m' would make every run stop before
    deleting anything. All of them must land on the default rather than being
    honoured, and the resulting budget must be usable arithmetic.
    """
    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_run_budget": setting_value,
        }
    )

    assert cleaner.run_budget_seconds == SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS
    assert math.isfinite(cleaner.run_budget_seconds)
    assert cleaner.run_budget_seconds > 0


@pytest.mark.parametrize("setting_value", [0, -1, "abc", "", 2.9])
def test_a_bad_batch_size_falls_back_to_the_default(setting_value):
    """A zero or negative batch size would make every DELETE a no-op and the
    loop spin, so unusable values must fall back rather than be honoured."""
    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_batch_size": setting_value,
        }
    )

    assert cleaner.batch_size >= 1


def test_operator_knobs_override_the_env_defaults():
    """The knobs are meant to be reachable from general_settings (and therefore
    from the admin UI), not only from environment variables."""
    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_batch_size": 250,
            "maximum_spend_logs_cleanup_max_batches": 7,
            "maximum_spend_logs_cleanup_run_budget": "90s",
            "maximum_spend_logs_cleanup_batch_timeout": "2m",
        }
    )

    assert cleaner.batch_size == 250
    assert cleaner.max_batches == 7
    assert cleaner.run_budget_seconds == 90
    assert cleaner.batch_timeout_seconds == 120


_BOUND_SETTING_CASES = (
    ("maximum_spend_logs_cleanup_batch_size", 137, "batch_size", 137),
    ("maximum_spend_logs_cleanup_max_batches", 9, "max_batches", 9),
    ("maximum_spend_logs_cleanup_run_budget", "45s", "run_budget_seconds", 45.0),
    ("maximum_spend_logs_cleanup_batch_timeout", "8s", "batch_timeout_seconds", 8.0),
)


@pytest.mark.parametrize("setting_name, setting_value, attribute, expected", _BOUND_SETTING_CASES)
@pytest.mark.asyncio
async def test_a_bound_changed_after_construction_reaches_the_next_run(
    setting_name, setting_value, attribute, expected
):
    """The scheduler holds one long-lived instance and the config reload mutates
    general_settings in place, so a bound captured at construction would leave
    every dashboard change inert until the process restarts."""
    settings = {"maximum_spend_logs_retention_period": "7d"}
    cleaner = SpendLogCleanup(general_settings=settings)
    cleaner.pod_lock_manager = None
    assert getattr(cleaner, attribute) != expected

    settings[setting_name] = setting_value

    await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))

    assert getattr(cleaner, attribute) == expected


@pytest.mark.parametrize("cleared_to_none", [True, False])
@pytest.mark.asyncio
async def test_a_bound_cleared_after_construction_falls_back_to_its_default(cleared_to_none):
    """Blanking the field in the dashboard has to restore the shipped default
    rather than leave the operator's old bound in force, whether the reload
    spells the clear as an explicit None or as an absent key."""
    settings = {"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_batch_size": 137}
    cleaner = SpendLogCleanup(general_settings=settings)
    cleaner.pod_lock_manager = None
    assert cleaner.batch_size == 137

    if cleared_to_none:
        settings["maximum_spend_logs_cleanup_batch_size"] = None
    else:
        del settings["maximum_spend_logs_cleanup_batch_size"]

    await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))

    assert cleaner.batch_size == SPEND_LOG_CLEANUP_BATCH_SIZE


def test_every_declared_bound_setting_is_covered_by_a_live_reread_case():
    """A bound added to the declared set without a live-reread case would be
    propagated by the proxy and then ignored by the running job."""
    assert {case[0] for case in _BOUND_SETTING_CASES} == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS)


@pytest.mark.asyncio
async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table():
    """The remaining-eligible-rows metric must never itself become the long
    scan this job exists to avoid, so its probe carries a LIMIT."""
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=0)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    await cleaner._delete_old_logs(mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline())

    count_sql = mock_db.query_raw.call_args[0][0]
    assert "count(*)" in count_sql
    assert "LIMIT $2" in count_sql
    assert mock_db.query_raw.call_args[0][2] == SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP


@pytest.mark.asyncio
async def test_a_run_skipped_because_another_pod_holds_the_lock_is_reported():
    """Operators need to tell "nothing to do" apart from "someone else is doing
    it", so a lock-skipped run is recorded under its own outcome."""
    recorded: list[str] = []
    original_record_run = SpendLogCleanupMetrics.record_run

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = MagicMock()
    cleaner.pod_lock_manager.redis_cache = MagicMock()
    cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
    cleaner.pod_lock_manager.release_lock = AsyncMock()

    SpendLogCleanupMetrics.record_run = classmethod(lambda cls, outcome: recorded.append(outcome))
    try:
        await cleaner.cleanup_old_spend_logs(mock_prisma_client)
    finally:
        SpendLogCleanupMetrics.record_run = original_record_run

    assert recorded == ["skipped_locked"]
    cleaner.pod_lock_manager.release_lock.assert_not_awaited()


@pytest.mark.asyncio
async def test_the_outstanding_rows_probe_carries_a_statement_timeout():
    """
    The probe is a statement like any other, so if it were issued bare a slow one
    would hold a connection past the budget the job advertises, which is exactly
    what the bounds exist to prevent. With budget to spare it carries the same
    per-statement timeout the delete batches do.
    """
    recorded: list[str] = []

    mock_prisma_client = MagicMock()
    mock_db = MagicMock()

    @asynccontextmanager
    async def _tx(**kwargs):
        tx = MagicMock()

        async def _execute_raw(sql, *args):
            recorded.append(sql.strip())
            return 0

        async def _query_raw(sql, *args):
            recorded.append(sql.strip())
            return [{"remaining": 7}]

        tx.execute_raw = _execute_raw
        tx.query_raw = _query_raw
        yield tx

    mock_db.tx = _tx
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_batch_timeout": "8s",
        }
    )

    remaining = await cleaner._count_remaining(
        mock_prisma_client,
        datetime.now(timezone.utc) - timedelta(days=7),
        "LiteLLM_SpendLogs",
        "startTime",
        _far_deadline(),
    )

    assert remaining == 7
    count_index = next(i for i, sql in enumerate(recorded) if sql.startswith("SELECT count(*)"))
    assert "SET LOCAL statement_timeout = 8000" in recorded[:count_index], (
        f"the probe ran without a statement timeout: {recorded}"
    )


@pytest.mark.asyncio
async def test_a_statement_timeout_is_clamped_to_the_budget_that_is_left():
    """
    Postgres has no 'stop at time T', only a per-statement duration, so a batch
    issued just under the deadline would run a whole batch timeout past it and
    the run budget would be advisory. Clamping the timeout to the remaining
    budget is what makes the budget a real wall clock.
    """
    recorded: list[str] = []
    client = MagicMock()

    @asynccontextmanager
    async def _tx(**kwargs):
        tx = MagicMock()

        async def _execute_raw(sql, *args):
            recorded.append(sql.strip())
            return 0

        tx.execute_raw = _execute_raw
        tx.query_raw = AsyncMock(return_value=[{"remaining": 0}])
        yield tx

    client.db.tx = _tx

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_batch_timeout": "30s",
        }
    )

    # Only 2s of budget left against a 30s batch timeout.
    await cleaner._execute_delete_batch(
        client, "DELETE FROM x", datetime.now(timezone.utc), "LiteLLM_SpendLogs", time.monotonic() + 2
    )

    timeouts = [sql for sql in recorded if "statement_timeout" in sql]
    assert timeouts, f"no statement timeout was issued: {recorded}"
    issued_ms = int(timeouts[0].split("=")[1].strip())
    assert issued_ms <= 2000, f"the batch was given {issued_ms}ms with only 2000ms of budget left"


@pytest.mark.asyncio
async def test_no_statement_is_issued_once_the_budget_is_spent():
    """
    Every table exits through _finish_table, including the ones a spent run never
    started, so an unconditional probe there would put one more statement per
    table past the bound.
    """
    client = _mock_prisma_for_retention([0, 0])
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})

    result = await cleaner._finish_table(
        client,
        datetime.now(timezone.utc) - timedelta(days=7),
        "LiteLLM_SpendLogs",
        "startTime",
        123,
        "budget_exhausted",
        time.monotonic() - 1,
    )

    assert result.rows_deleted == 123
    assert result.stop_reason == "budget_exhausted"
    client.db.query_raw.assert_not_called()


@pytest.mark.asyncio
async def test_a_batch_cancelled_by_the_deadline_is_budget_exhaustion_not_a_failure(monkeypatch):
    """
    Clamping the timeout means the last batch of a budget-exhausted run is
    cancelled by the deadline itself. Counting that as a batch failure would
    inflate the failure metric on every such run and walk it toward the abort
    threshold, so it has to be classified as the bound working.
    """
    failures: list[str] = []
    client = MagicMock()
    _wire_tx(client.db)

    # The deadline has to pass DURING the batch, not before it: a deadline
    # already spent is caught by the loop's own check and no batch is ever
    # issued, which would exercise none of the classification under test.
    async def _cancelled_after_the_deadline(sql, *args):
        await asyncio.sleep(0.05)
        raise Exception("canceling statement due to statement timeout")

    client.db.execute_raw = _cancelled_after_the_deadline

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    monkeypatch.setattr(SpendLogCleanupMetrics, "record_batch_failure", lambda table: failures.append(table))

    result = await cleaner._delete_old_logs(
        client, datetime.now(timezone.utc) - timedelta(days=7), time.monotonic() + 0.02
    )

    assert result.stop_reason == "budget_exhausted"
    assert failures == [], f"a deadline cancellation was recorded as a batch failure: {failures}"


@pytest.mark.asyncio
async def test_partition_maintenance_is_skipped_once_the_run_budget_is_spent():
    """
    Dropping a partition is DDL holding an ACCESS EXCLUSIVE lock, and unlike a
    delete batch it cannot be cut short once it has started. A run whose budget is
    already gone must therefore not start it at all; the next tick picks it up.
    """
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=0)
    mock_prisma_client.db = mock_db

    partition_manager = MagicMock()
    partition_manager.is_partitioned = AsyncMock(return_value=True)
    partition_manager.ensure_partitions = AsyncMock()
    partition_manager.drop_partitions_older_than = AsyncMock(return_value=[])

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "use_spend_logs_partitioning": True,
        },
        partition_manager=partition_manager,
    )
    cleaner._should_delete_spend_logs()

    # a deadline already in the past is what a run that spent its budget on an
    # earlier table looks like
    await cleaner._clean_spend_log_tables(mock_prisma_client, time.monotonic() - 1)

    partition_manager.ensure_partitions.assert_not_awaited()
    partition_manager.drop_partitions_older_than.assert_not_awaited()


@pytest.mark.asyncio
async def test_partition_maintenance_still_runs_while_the_run_has_budget():
    """The skip above must be caused by the spent budget, not by breaking the
    partition path outright."""
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=0)
    mock_prisma_client.db = mock_db

    partition_manager = MagicMock()
    partition_manager.is_partitioned = AsyncMock(return_value=True)
    partition_manager.ensure_partitions = AsyncMock()
    partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"])

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "use_spend_logs_partitioning": True,
        },
        partition_manager=partition_manager,
    )
    cleaner._should_delete_spend_logs()

    await cleaner._clean_spend_log_tables(mock_prisma_client, _far_deadline())

    partition_manager.ensure_partitions.assert_awaited_once()
    partition_manager.drop_partitions_older_than.assert_awaited_once()


@pytest.mark.parametrize(
    "stop_reasons, expected",
    [
        (("exhausted",), "completed"),
        (("exhausted", "exhausted"), "completed"),
        (("exhausted", "batch_cap_reached"), "batch_cap_reached"),
        (("batch_cap_reached", "exhausted"), "batch_cap_reached"),
        (("exhausted", "budget_exhausted"), "budget_exhausted"),
        (("budget_exhausted", "exhausted"), "budget_exhausted"),
        (("batch_cap_reached", "budget_exhausted"), "budget_exhausted"),
        (("budget_exhausted", "batch_cap_reached"), "budget_exhausted"),
        (("exhausted", "aborted"), "aborted"),
        (("aborted", "exhausted"), "aborted"),
        (("budget_exhausted", "aborted"), "aborted"),
        (("aborted", "budget_exhausted"), "aborted"),
        (("aborted", "budget_exhausted", "batch_cap_reached"), "aborted"),
    ],
)
def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(stop_reasons, expected):
    """
    The run outcome answers "why did this run stop", so a table that merely ran
    dry must never mask one that hit a bound, and an abort must outrank both.

    Both orders of every pair are covered because this folds several per-table
    results into one answer: a first-match-wins implementation would pass on
    whichever order happened to be written and fail on its mirror.
    """
    results = tuple(
        TableCleanupResult(table_name=f"t{i}", rows_deleted=0, stop_reason=reason)
        for i, reason in enumerate(stop_reasons)
    )
    assert SpendLogCleanup._run_outcome(results) == expected


_OTHER_OUTCOMES: Final = ("completed", "budget_exhausted", "batch_cap_reached", "skipped_locked", "skipped_disabled")


def _runs_recorded(outcome: str) -> float:
    """The real ``litellm_spend_log_cleanup_runs_total`` sample for one outcome, 0 when unset"""
    from prometheus_client import REGISTRY

    return REGISTRY.get_sample_value("litellm_spend_log_cleanup_runs_total", {"outcome": outcome}) or 0.0


@pytest.mark.asyncio
async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_raising(monkeypatch):
    """A run cut short by shutdown must leave its outcome and how far it got behind"""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    mock_logger = MagicMock()
    monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
    aborted_runs_before = _runs_recorded("aborted")
    other_runs_before = {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES}

    third_batch_reached = asyncio.Event()

    async def _execute_raw(sql, *args):
        if third_batch_reached.is_set():
            raise AssertionError("no batch may be issued after the cancelled one")
        if _execute_raw.calls < 2:
            _execute_raw.calls += 1
            return 150
        third_batch_reached.set()
        await asyncio.Event().wait()

    _execute_raw.calls = 0
    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_prisma_client.db.execute_raw = _execute_raw

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = MagicMock()
    cleaner.pod_lock_manager.redis_cache = MagicMock()
    cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
    cleaner.pod_lock_manager.release_lock = AsyncMock()

    run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(mock_prisma_client))
    await asyncio.wait_for(third_batch_reached.wait(), timeout=5)
    run.cancel()
    with pytest.raises(asyncio.CancelledError):
        await run

    assert _runs_recorded("aborted") == aborted_runs_before + 1
    assert {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} == other_runs_before
    cleaner.pod_lock_manager.release_lock.assert_awaited_once()
    mock_logger.exception.assert_not_called()
    (error_call,) = mock_logger.error.call_args_list
    rendered = error_call[0][0] % error_call[0][1:]
    assert rendered.startswith("Spend log cleanup cancelled after ")
    assert "s (rows_deleted=300, batches=2)" in rendered


@pytest.mark.asyncio
async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatch):
    """The scheduler holds one cleaner for the life of the process, so progress must not carry over"""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    mock_logger = MagicMock()
    monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)

    mock_prisma_client = MagicMock()
    _wire_tx(mock_prisma_client.db)
    mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0])
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = None
    await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, asyncio.CancelledError()])
    with pytest.raises(asyncio.CancelledError):
        await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    (error_call,) = mock_logger.error.call_args_list
    rendered = error_call[0][0] % error_call[0][1:]
    assert "(rows_deleted=150, batches=1)" in rendered


@pytest.mark.asyncio
async def test_a_size_pass_that_fires_while_an_age_pass_is_running_in_the_same_pod_waits_for_it(monkeypatch):
    """Without Redis there is no pod lock, and the two jobs are scheduled on two SpendLogCleanup instances"""
    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module

    mock_logger = MagicMock()
    monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)

    first_batch_done = asyncio.Event()
    second_run_done = asyncio.Event()

    async def _slow_execute_raw(sql, *args):
        first_batch_done.set()
        await second_run_done.wait()
        return 100

    slow_client = MagicMock()
    _wire_tx(slow_client.db)
    slow_client.db.execute_raw = _slow_execute_raw
    size_db = MagicMock()
    _wire_tx(size_db)
    size_db.execute_raw = AsyncMock(return_value=1000)
    _wire_size_probes(size_db, [_storage(allocated=10 * 1024**3, live=1000, dead=0)], cutoff=datetime.now(timezone.utc))
    size_client = MagicMock()
    size_client.db = size_db

    age_cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    age_cleaner.pod_lock_manager = None
    size_cleaner = _size_cleaner({})

    slow_run = asyncio.ensure_future(age_cleaner.cleanup_old_spend_logs(slow_client))
    await asyncio.wait_for(first_batch_done.wait(), timeout=5)
    size_run = asyncio.ensure_future(size_cleaner.cleanup_spend_logs_by_size(size_client))
    await asyncio.sleep(0.05)
    assert size_db.deletes.await_count == 0 and not size_run.done(), "the size pass ran while the age pass held the pod"

    second_run_done.set()
    await asyncio.sleep(0)
    slow_run.cancel()
    with pytest.raises(asyncio.CancelledError):
        await slow_run
    await asyncio.wait_for(size_run, timeout=5)

    assert size_db.deletes.await_count > 0, "the size pass never ran once the age pass let go"
    (error_call,) = mock_logger.error.call_args_list
    rendered = error_call[0][0] % error_call[0][1:]
    assert "(rows_deleted=100, batches=1)" in rendered


@pytest.mark.asyncio
async def test_spend_logs_backlog_cannot_starve_tool_index_cleanup():
    """
    Both spend-log tables share one run budget. Before the fix the spend-log
    loop ran against the whole deadline, so a backlog that outlasted the budget
    meant LiteLLM_SpendLogToolIndex never received a single delete batch, run
    after run. The index table must still get its own share of the budget.
    """
    mock_prisma_client = MagicMock()
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=1000)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_max_batches": 500,
            "maximum_spend_logs_cleanup_run_budget": "1s",
        }
    )
    cleaner.pod_lock_manager = None

    started_at = time.monotonic()
    await cleaner.cleanup_old_spend_logs(mock_prisma_client)
    elapsed = time.monotonic() - started_at

    tables = [call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list]
    assert tables.count("LiteLLM_SpendLogs") > 0
    assert tables.count("LiteLLM_SpendLogToolIndex") > 0, "tool index cleanup was starved by the spend-log backlog"
    assert elapsed < 2.5, f"splitting the budget must not extend the run: {elapsed}s"


@pytest.mark.asyncio
async def test_run_that_leaves_backlog_logs_a_warning_summary_naming_each_table(caplog):
    """
    Operators running at warning or error level saw nothing when a run stopped
    with expired rows still present. A run that ends on a bound must emit one
    WARNING line that names every table, its rows deleted and its stop reason.
    """
    mock_prisma_client = MagicMock()
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=1000)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(
        general_settings={
            "maximum_spend_logs_retention_period": "7d",
            "maximum_spend_logs_cleanup_max_batches": 2,
        }
    )
    cleaner.pod_lock_manager = None

    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    summaries = [record for record in caplog.records if "Spend log cleanup run finished" in record.getMessage()]
    assert len(summaries) == 1
    summary = summaries[0]
    assert summary.levelno == logging.WARNING
    message = summary.getMessage()
    assert "outcome=batch_cap_reached" in message
    assert "LiteLLM_SpendLogs: deleted=2000 stop_reason=batch_cap_reached" in message
    assert "LiteLLM_SpendLogToolIndex: deleted=2000 stop_reason=batch_cap_reached" in message


@pytest.mark.asyncio
async def test_run_that_drains_every_table_logs_the_summary_at_info_not_warning(caplog):
    """A healthy run must not page anyone: the summary stays at INFO."""
    mock_prisma_client = MagicMock()
    mock_db = MagicMock()
    _wire_tx(mock_db)
    mock_db.execute_raw = AsyncMock(return_value=0)
    mock_prisma_client.db = mock_db

    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
    cleaner.pod_lock_manager = None

    with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
        await cleaner.cleanup_old_spend_logs(mock_prisma_client)

    summaries = [record for record in caplog.records if "Spend log cleanup run finished" in record.getMessage()]
    assert len(summaries) == 1
    assert summaries[0].levelno == logging.INFO
    assert "outcome=completed" in summaries[0].getMessage()


@pytest.mark.asyncio
async def test_a_cleanup_delete_batch_renders_a_postgres_delete_span_for_its_table(
    postgres_span_names: Callable[[], Awaitable[tuple[str, ...]]],
) -> None:
    client = MagicMock()
    _wire_tx(client.db)
    client.db.execute_raw = engine_call(5)

    await SpendLogCleanup(general_settings={})._execute_delete_batch(
        client, "DELETE FROM x", datetime(2026, 1, 1, tzinfo=timezone.utc), "LiteLLM_SpendLogs", time.monotonic() + 2
    )

    assert await postgres_span_names() == ("postgres.delete LiteLLM_SpendLogs",)


# --- size rule (maximum_spend_logs_size) ---------------------------------------------------------------


def _storage(allocated: int, live: int, dead: int, inserted: int | None = None, deleted: int | None = None) -> dict:
    return {
        "allocated_bytes": allocated,
        "live_rows": live,
        "dead_rows": dead,
        "inserted_total": live + dead if inserted is None else inserted,
        "deleted_total": dead if deleted is None else deleted,
    }


def _wire_size_probes(
    db, storage_rows: list[dict], cutoff: datetime, saved_state: SizeCleanupState | None = None
) -> list[SizeCleanupState]:
    """
    Answer the size pass's catalog reads: the storage probe walks `storage_rows` one measurement per call
    (repeating the last), the step-cutoff query returns `cutoff`, the remaining-rows probe returns 0, and the
    LiteLLM_Config row holds `saved_state`. Writes of that row land in the returned list instead of
    db.execute_raw, so DELETE counts stay what tests configure.
    """
    measurements = iter(storage_rows)
    last = [storage_rows[-1]]
    saves: list[SizeCleanupState] = []
    deletes = db.execute_raw

    async def _query_raw(sql, *args):
        if "pg_total_relation_size" in sql:
            row = next(measurements, None)
            if row is None:
                row = last[0]
            return [row]
        if "OFFSET" in sql:
            return [{"startTime": cutoff}]
        if '"LiteLLM_Config"' in sql:
            return [{"state_json": saved_state.model_dump_json()}] if saved_state is not None else []
        return [{"remaining": 0}]

    async def _execute_raw(sql, *args):
        if '"LiteLLM_Config"' in sql:
            saves.append(SizeCleanupState.model_validate_json(args[1]))
            return 1
        return await deletes(sql, *args)

    db.query_raw = AsyncMock(side_effect=_query_raw)
    db.execute_raw = _execute_raw
    db.deletes = deletes
    return saves


def _size_cleaner(settings: dict, **kwargs) -> SpendLogCleanup:
    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_size": "1GB", **settings}, **kwargs)
    cleaner.pod_lock_manager = None
    return cleaner


@pytest.mark.parametrize(
    ("value", "expected"),
    [
        ("180GB", 180 * 1024**3),
        ("20 GiB", 20 * 1024**3),
        ("512MB", 512 * 1024**2),
        ("1.5TB", int(1.5 * 1024**4)),
        ("4096", 4096),
        (4096, 4096),
    ],
)
def test_size_in_bytes_reads_the_units_pg_size_pretty_prints(value, expected):
    from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import size_in_bytes

    assert size_in_bytes(value) == expected


@pytest.mark.parametrize("value", ["", "big", "-1GB", "0", 0, "1PB"])
def test_size_in_bytes_rejects_what_cannot_be_a_size(value):
    from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import size_in_bytes

    with pytest.raises(ValueError, match="size"):
        size_in_bytes(value)


def test_an_unset_size_means_no_size_rule_and_a_bad_target_falls_back_to_the_size(caplog):
    assert SpendLogCleanup(general_settings={})._size_rule() is None
    rule = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "200MB"})._size_rule()
    assert rule is not None
    assert (rule.limit_bytes, rule.target_bytes) == (1024**3, 200 * 1024**2)
    with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
        rule_without_room = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "2GB"})._size_rule()
    assert rule_without_room is not None
    assert rule_without_room.target_bytes == rule_without_room.limit_bytes == 1024**3
    assert any("must be below maximum_spend_logs_size" in record.getMessage() for record in caplog.records)


@pytest.mark.asyncio
async def test_size_pass_under_the_limit_issues_no_delete():
    db = MagicMock()
    _wire_tx(db)
    db.execute_raw = AsyncMock(return_value=0)
    _wire_size_probes(db, [_storage(allocated=500 * 1024**2, live=1000, dead=0)], cutoff=datetime.now(timezone.utc))
    client = MagicMock()
    client.db = db

    await _size_cleaner({}).cleanup_spend_logs_by_size(client)

    assert db.deletes.await_count == 0


@pytest.mark.asyncio
async def test_size_pass_deletes_oldest_first_in_steps_until_the_live_estimate_is_under_the_target(caplog):
    """
    Over the limit, the pass turns the next step of oldest rows into a cutoff and runs the retention DELETE
    for the spend logs and their tool-index rows, re-measuring after each step; it stops on the first
    measurement under the target, not when the file size (which deletes never shrink) moves.
    """
    db = MagicMock()
    _wire_tx(db)
    # 3 GiB over 3000 rows is 1 MiB a row, so 2560 rows take it from 3 GiB to the 512 MiB target: one step
    # of 3 batches; the tool index is cleaned on the same cutoff
    db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 300, 0])
    step_cutoff = datetime(2026, 9, 1, tzinfo=timezone.utc)
    allocated = 3 * 1024**3
    _wire_size_probes(db, [_storage(allocated, live=3000, dead=0)], cutoff=step_cutoff)
    client = MagicMock()
    client.db = db
    cleaner = _size_cleaner({"maximum_spend_logs_cleanup_target_size": "512MB"})

    with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
        await cleaner.cleanup_spend_logs_by_size(client)

    deletes = [(call[0][0], call[0][1:]) for call in db.deletes.call_args_list]
    assert len(deletes) == 5
    assert 'DELETE FROM "LiteLLM_SpendLogs"' in deletes[0][0]
    assert 'WHERE "startTime" < $1' in deletes[0][0]
    assert deletes[0][1][0] == step_cutoff
    assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in deletes[3][0]
    assert deletes[3][1][0] == step_cutoff
    offset_query = next(call for call in db.query_raw.call_args_list if "OFFSET" in call[0][0])
    assert offset_query[0][1] == 3000
    summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
    assert "outcome=completed" in summary
    assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=target_reached" in summary
    assert "LiteLLM_SpendLogToolIndex (size rule): deleted=300 stop_reason=target_reached" in summary
    assert f"allocated_before={allocated}" in summary
    assert "estimated_live_before=3221225472" in summary
    assert "estimated_live_after=0" in summary
    assert "target_reached=true above_target=0" in summary
    assert "VACUUM FULL or pg_repack" in summary


@pytest.mark.asyncio
async def test_size_pass_stops_at_the_run_batch_cap_and_reports_how_far_above_the_target_it_is(caplog):
    db = MagicMock()
    _wire_tx(db)
    # 4 GiB over 4000 rows: 3584 rows must go to reach the 512 MiB target, the cap allows 3 batches
    db.execute_raw = AsyncMock(side_effect=[1000, 1000, 1000, 0])
    allocated = 4 * 1024**3
    saves = _wire_size_probes(db, [_storage(allocated, live=4000, dead=0)], cutoff=datetime.now(timezone.utc))
    client = MagicMock()
    client.db = db
    cleaner = _size_cleaner(
        {"maximum_spend_logs_cleanup_target_size": "512MB", "maximum_spend_logs_cleanup_max_batches": 3}
    )

    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        await cleaner.cleanup_spend_logs_by_size(client)

    assert db.deletes.await_count == 4
    summary = next(r for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
    assert summary.levelno == logging.WARNING
    assert "outcome=batch_cap_reached" in summary.getMessage()
    assert "LiteLLM_SpendLogs (size rule): deleted=3000 stop_reason=batch_cap_reached" in summary.getMessage()
    # 1000 of 4000 rows left is 1 GiB live, 512 MiB above the target
    assert f"target_reached=false above_target={512 * 1024**2}" in summary.getMessage()
    # the next check starts from what this run freed, with the deletes already counted out of n_tup_del
    assert saves == [
        SizeCleanupState(
            allocated_bytes=allocated,
            reclaimable_bytes=round(3000 * allocated / 4000),
            bytes_per_row=allocated / 4000,
            inserted_total=4000,
            deleted_total=3000,
        )
    ]


@pytest.mark.asyncio
async def test_size_pass_analyzes_and_measures_again_when_the_row_statistics_were_reset(caplog):
    """A crash zeroes pg_stat_user_tables under a full file; the pass rebuilds the counters instead of skipping.
    The first read after ANALYZE still shows zeros, as the stats collector publishes them a moment later"""
    db = MagicMock()
    _wire_tx(db)
    statements = AsyncMock(return_value=0)
    db.execute_raw = statements
    saves = _wire_size_probes(
        db,
        [
            _storage(allocated=10 * 1024**3, live=0, dead=0, inserted=0, deleted=0),
            _storage(allocated=10 * 1024**3, live=0, dead=0, inserted=0, deleted=0),
            _storage(10 * 1024**3, 1000, 0, 0, 0),
        ],
        cutoff=datetime.now(timezone.utc),
        saved_state=SizeCleanupState(
            allocated_bytes=10 * 1024**3,
            reclaimable_bytes=9 * 1024**3 + 512 * 1024**2,
            bytes_per_row=1024.0,
            inserted_total=50_000,
            deleted_total=49_000,
        ),
    )
    client = MagicMock()
    client.db = db

    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        await _size_cleaner({}).cleanup_spend_logs_by_size(client)

    issued = [call.args[0] for call in statements.await_args_list]
    assert issued == ['ANALYZE "LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"'], (
        "the only statement is the ANALYZE: no DELETE, the carried count puts the table under 1GB"
    )
    assert any("row statistics are stale" in record.getMessage() for record in caplog.records)
    assert saves[0].reclaimable_bytes == 9 * 1024**3 + 512 * 1024**2


@pytest.mark.asyncio
async def test_size_pass_does_not_refire_after_autovacuum_forgets_the_rows_it_deleted():
    """
    The run before this one deleted 15,000 of 20,000 rows and saved what it freed. Autovacuum has since
    turned those dead tuples into free space, so pg_stat_user_tables reports 5,000 live and 0 dead rows while
    the files still measure the full 32 MB. Trusting the live share would call all 32 MB live and empty the
    table; the saved count says 8 MB are live, under the limit, so nothing is deleted.
    """
    db = MagicMock()
    _wire_tx(db)
    db.execute_raw = AsyncMock(return_value=1000)
    after_last_run = SizeCleanupState(
        allocated_bytes=32_030_720,
        reclaimable_bytes=24_015_000,
        bytes_per_row=1601.0,
        inserted_total=20_000,
        deleted_total=15_000,
    )
    saves = _wire_size_probes(
        db,
        [_storage(32_030_720, live=5000, dead=0, inserted=20_000, deleted=15_000)],
        cutoff=datetime.now(timezone.utc),
        saved_state=after_last_run,
    )
    client = MagicMock()
    client.db = db
    cleaner = _size_cleaner({"maximum_spend_logs_size": "16MB", "maximum_spend_logs_cleanup_target_size": "10MB"})

    await cleaner.cleanup_spend_logs_by_size(client)

    assert db.deletes.await_count == 0
    assert saves[0].reclaimable_bytes == 24_015_000


_RUN_STATE = SizeCleanupState(
    allocated_bytes=32_030_720,
    reclaimable_bytes=24_015_000,
    bytes_per_row=1601.0,
    inserted_total=20_000,
    deleted_total=15_000,
)


@pytest.mark.parametrize(
    ("name", "probe", "reclaimable_expected"),
    [
        (
            "vacuum zeroed the dead rows, nothing else changed",
            _storage(32_030_720, 5000, 0, 20_000, 15_000),
            24_015_000,
        ),
        ("7000 new rows written into the free space", _storage(32_030_720, 12_000, 0, 27_000, 15_000), 12_808_000),
        (
            "7000 new rows that all grew the files",
            _storage(32_030_720 + 7000 * 1601, 12_000, 0, 27_000, 15_000),
            24_015_000,
        ),
        ("VACUUM FULL handed the free space back", _storage(8_015_720, 5000, 0, 20_000, 15_000), 0),
        (
            "statistics wiped while the files still hold data: the count is carried forward, not reseeded",
            _storage(32_030_720, 0, 0, 0, 0),
            24_015_000,
        ),
        ("the retention job deleted 2000 more rows", _storage(32_030_720, 3000, 2000, 20_000, 17_000), 27_217_000),
        (
            "statistics reset (a crash, then ANALYZE): the count is carried forward, not restarted from the dead rows",
            _storage(32_030_720, 5000, 1000, 5000, 1000),
            24_015_000,
        ),
        (
            "statistics reset on a file VACUUM FULL shrank meanwhile: the carried count loses the shrink",
            _storage(8_015_720, 5000, 0, 5000, 0),
            0,
        ),
        (
            "statistics reset with more dead rows than the carried count explains: the dead rows win",
            _storage(32_030_720, 1000, 19_000, 1000, 19_000),
            30_429_184,
        ),
    ],
)
def test_advance_size_state_tracks_free_space_through_vacuum_inserts_and_other_deletes(
    name, probe, reclaimable_expected
):
    advanced = advance_size_state(_RUN_STATE, StorageProbe(**probe))

    assert advanced.reclaimable_bytes == reclaimable_expected, name
    assert 0 <= probe["allocated_bytes"] - advanced.reclaimable_bytes <= probe["allocated_bytes"], name


@pytest.mark.asyncio
async def test_size_only_run_leaves_the_age_rules_to_the_retention_job():
    db = MagicMock()
    _wire_tx(db)
    db.execute_raw = AsyncMock(return_value=0)
    _wire_size_probes(db, [_storage(allocated=10 * 1024**3, live=100, dead=0)], cutoff=datetime.now(timezone.utc))
    client = MagicMock()
    client.db = db
    cleaner = _size_cleaner(
        {"maximum_spend_logs_retention_period": "7d", "maximum_health_check_retention_period": "30d"}
    )

    await cleaner.cleanup_spend_logs_by_size(client)

    tables = {call[0][0].split('"')[1] for call in db.deletes.call_args_list}
    assert tables == {"LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"}


@pytest.mark.asyncio
@pytest.mark.parametrize(("cron", "size_pass_expected"), [(None, False), ("0 3 * * *", True)])
async def test_the_retention_run_carries_the_size_pass_only_when_one_cron_drives_both(cron, size_pass_expected):
    db = MagicMock()
    _wire_tx(db)
    db.execute_raw = AsyncMock(return_value=0)
    _wire_size_probes(db, [_storage(allocated=10 * 1024**3, live=100, dead=0)], cutoff=datetime.now(timezone.utc))
    client = MagicMock()
    client.db = db
    cleaner = _size_cleaner({"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_cron": cron})

    await cleaner.cleanup_old_spend_logs(client)

    measured = any("pg_total_relation_size" in call[0][0] for call in db.query_raw.call_args_list)
    assert measured is size_pass_expected


@pytest.mark.asyncio
async def test_partitioned_size_pass_drops_the_oldest_partitions_first_and_skips_row_deletes_once_under_target(
    caplog,
):
    db = MagicMock()
    _wire_tx(db)
    db.execute_raw = AsyncMock(return_value=1000)
    saves = _wire_size_probes(
        db,
        [_storage(10 * 1024**3, live=1000, dead=0), _storage(256 * 1024**2, live=25, dead=0)],
        cutoff=datetime.now(timezone.utc),
    )
    client = MagicMock()
    client.db = db
    verdicts = []

    async def _drop(prisma_client, now, remaining_timeout_ms, still_over_target):
        verdicts.append(await still_over_target())
        return ("LiteLLM_SpendLogs_p20260601",)

    partition_manager = MagicMock()
    partition_manager.is_partitioned = AsyncMock(return_value=True)
    partition_manager.ensure_partitions = AsyncMock()
    partition_manager.drop_oldest_partitions_while = AsyncMock(side_effect=_drop)
    cleaner = _size_cleaner(
        {"use_spend_logs_partitioning": True, "maximum_spend_logs_cleanup_target_size": "512MB"},
        partition_manager=partition_manager,
    )

    with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
        await cleaner.cleanup_spend_logs_by_size(client)

    partition_manager.ensure_partitions.assert_awaited_once(), "a size-only setup still needs tomorrow's partition"
    partition_manager.drop_oldest_partitions_while.assert_awaited_once()
    assert verdicts == [False]
    assert db.deletes.await_count == 0
    summary = next(r.getMessage() for r in caplog.records if "Spend log cleanup run finished" in r.getMessage())
    assert "partitions_dropped=1" in summary
    assert "target_reached=true" in summary
    assert (saves[-1].inserted_total, saves[-1].deleted_total, saves[-1].allocated_bytes) == (25, 0, 256 * 1024**2), (
        "the saved counters must be the post-drop ones, or the next check reads the drop as a statistics reset"
    )


@pytest.mark.asyncio
async def test_a_failed_remeasure_after_a_partition_drop_stops_the_pass_instead_of_deleting_on_old_numbers(caplog):
    db = MagicMock()
    _wire_tx(db)
    db.execute_raw = AsyncMock(return_value=1000)
    saves = _wire_size_probes(db, [_storage(10 * 1024**3, live=1000, dead=0)], cutoff=datetime.now(timezone.utc))
    client = MagicMock()
    client.db = db

    async def _drop(prisma_client, now, remaining_timeout_ms, still_over_target):
        db.query_raw = AsyncMock(side_effect=RuntimeError("connection reset"))
        assert await still_over_target() is False
        return ("LiteLLM_SpendLogs_p20260601",)

    partition_manager = MagicMock()
    partition_manager.is_partitioned = AsyncMock(return_value=True)
    partition_manager.ensure_partitions = AsyncMock()
    partition_manager.drop_oldest_partitions_while = AsyncMock(side_effect=_drop)
    cleaner = _size_cleaner(
        {"use_spend_logs_partitioning": True, "maximum_spend_logs_cleanup_target_size": "512MB"},
        partition_manager=partition_manager,
    )

    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        await cleaner.cleanup_spend_logs_by_size(client)

    assert db.deletes.await_count == 0, "the row step ran on the pre-drop measurement"
    assert saves == [], "a count saved from pre-drop counters reads the drop as a statistics reset next check"
    assert any("could not be re-measured after dropping partitions" in r.getMessage() for r in caplog.records)


def test_drop_oldest_partitions_while_only_drops_fully_elapsed_partitions_oldest_first():
    from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import SpendLogsPartitionManager

    manager = SpendLogsPartitionManager()
    now = datetime(2026, 10, 17, 12, tzinfo=timezone.utc)
    manager._list_partitions = AsyncMock(
        return_value=[
            ("LiteLLM_SpendLogs_p20261017", datetime(2026, 10, 18)),  # today's, still receiving writes
            ("LiteLLM_SpendLogs_p20261015", datetime(2026, 10, 16)),
            ("LiteLLM_SpendLogs_default", None),
            ("LiteLLM_SpendLogs_p20261014", datetime(2026, 10, 15)),
        ]
    )
    dropped_sql: list[str] = []

    async def _ddl(prisma_client, statement, timeout_ms):
        dropped_sql.append(statement)

    manager._execute_bounded_ddl = _ddl
    verdicts = iter([True, False])

    async def still_over_target():
        return next(verdicts)

    dropped = asyncio.run(manager.drop_oldest_partitions_while(MagicMock(), now, lambda: 1000, still_over_target))

    assert dropped == ("LiteLLM_SpendLogs_p20261014", "LiteLLM_SpendLogs_p20261015")
    assert dropped_sql == [
        'DROP TABLE IF EXISTS "LiteLLM_SpendLogs_p20261014"',
        'DROP TABLE IF EXISTS "LiteLLM_SpendLogs_p20261015"',
    ]


@pytest.mark.parametrize(
    ("settings", "retention_job", "size_job"),
    [
        ({}, False, False),
        ({"maximum_spend_logs_retention_period": "7d"}, True, False),
        ({"maximum_spend_logs_size": "180GB"}, False, True),
        ({"maximum_spend_logs_size": "180GB", "maximum_spend_logs_retention_period": "7d"}, True, True),
        ({"maximum_spend_logs_size": "180GB", "maximum_spend_logs_cleanup_cron": "0 3 * * *"}, True, False),
    ],
)
def test_which_scheduler_jobs_the_cleanup_settings_ask_for(settings, retention_job, size_job):
    from litellm.proxy.proxy_server import _wants_retention_job, _wants_size_check_job

    assert _wants_retention_job(settings) is retention_job
    assert _wants_size_check_job(settings) is size_job


def test_after_deleting_counts_the_rows_and_any_file_growth_as_reclaimable_and_to_state_keeps_the_counters():
    storage = SpendLogStorage(allocated_bytes=32_030_720, reclaimable_bytes=0, bytes_per_row=1601.0, live_rows=20_000)

    after = storage.after_deleting(15_000, 32_030_720 + 8192)

    assert (after.reclaimable_bytes, after.live_rows, after.estimated_live_bytes) == (24_023_192, 5000, 8_015_720)
    assert after.to_state(inserted_total=20_000, deleted_total=15_000) == SizeCleanupState(
        allocated_bytes=32_038_912,
        reclaimable_bytes=24_023_192,
        bytes_per_row=1601.0,
        inserted_total=20_000,
        deleted_total=15_000,
    )
    assert StorageProbe(
        allocated_bytes=8192, live_rows=0, dead_rows=0, inserted_total=0, deleted_total=0
    ).stats_are_stale


@pytest.mark.asyncio
async def test_probe_storage_reads_the_catalog_row_for_both_tables_and_a_failed_probe_is_none(caplog):
    prisma_client = MagicMock()
    _wire_tx(prisma_client.db)
    prisma_client.db.query_raw = AsyncMock(return_value=[_storage(32_030_720, 20_000, 0)])
    cleaner = SpendLogCleanup(general_settings={})

    probe = await cleaner._probe_storage(prisma_client, _far_deadline())

    assert probe is not None
    assert (probe.allocated_bytes, probe.live_rows, probe.dead_rows, probe.inserted_total, probe.deleted_total) == (
        32_030_720,
        20_000,
        0,
        20_000,
        0,
    )
    assert prisma_client.db.query_raw.call_args.args[1] == ["LiteLLM_SpendLogs", "LiteLLM_SpendLogToolIndex"]

    prisma_client.db.query_raw = AsyncMock(side_effect=RuntimeError("pg_stat_user_tables unavailable"))
    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        assert await cleaner._probe_storage(prisma_client, _far_deadline()) is None
    assert any("Could not measure spend-log storage" in record.getMessage() for record in caplog.records)


@pytest.mark.asyncio
async def test_size_state_round_trips_through_its_litellm_config_row_and_a_lost_save_is_only_a_warning(caplog):
    saved = {}
    prisma_client = MagicMock()
    _wire_tx(prisma_client.db)

    async def execute_raw(sql, *args):
        saved[args[0]] = args[1]
        return 1

    prisma_client.db.execute_raw = execute_raw
    prisma_client.db.query_raw = AsyncMock(
        side_effect=lambda sql, param: [{"state_json": saved[param]}] if param in saved else []
    )
    cleaner = SpendLogCleanup(general_settings={})

    assert await cleaner._load_size_state(prisma_client, _far_deadline()) is None
    await cleaner._save_size_state(prisma_client, _far_deadline(), _RUN_STATE)
    assert list(saved) == ["spend_logs_size_cleanup_state"]
    loaded = await cleaner._load_size_state(prisma_client, _far_deadline())
    assert loaded is not None
    assert loaded.model_dump() == _RUN_STATE.model_dump()

    prisma_client.db.execute_raw = AsyncMock(side_effect=RuntimeError("LiteLLM_Config is read only"))
    with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
        await cleaner._save_size_state(prisma_client, _far_deadline(), _RUN_STATE)
    assert any("Could not save the spend-log size cleanup state" in record.getMessage() for record in caplog.records)