test(accounting): verify delivery continues while maintenance is locked (#50843)

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-11-10 06:52:48 +00:00 • committed by GitHub
parent d1a0cc6b24
commit b3f75884a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -978,14 +978,16 @@ async def test_fallback_charges_only_groups_declared_by_the_served_deployment(
@pytest.mark.asyncio
async def test_worker_skips_another_leader_and_releases_its_redis_lock_after_retry(
async def test_worker_delivers_under_another_leaders_lock_then_cleans_reports_and_releases_its_lock(
parity_database: str, monkeypatch: pytest.MonkeyPatch
):
from prometheus_client import REGISTRY
from litellm.accounting.worker import flush_accounting_deliveries
from litellm.caching.redis_cache import RedisCache
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
monkeypatch.setenv("DATABASE_URL", parity_database)
monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")
cache = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), namespace=uuid.uuid4().hex)
other_leader = PodLockManager(cache)
writer = DBSpendUpdateWriter(redis_cache=cache)
@ -993,16 +995,39 @@ async def test_worker_skips_another_leader_and_releases_its_redis_lock_after_ret
client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
await client.db.connect()
try:
reports_before: Final = REGISTRY.get_sample_value("litellm_accounting_batch_seconds_count")
assert reports_before is not None
with psycopg.connect(parity_database, autocommit=True) as connection:
connection.execute('INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch) VALUES ($$invalid-row$$, $$ {"metadata": {"accounting_windows": [{}]}} $$::jsonb, false)')
connection.execute(
'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, attempts) '
"VALUES ('fresh', %s::jsonb, false, 0), ('retry', %s::jsonb, false, 1)",
(
json.dumps({"spend": 0.25, "metadata": {"accounting_entities": {"user": "user"}}}),
json.dumps({"spend": 0.5, "metadata": {"accounting_entities": {"user": "user"}}}),
),
)
connection.execute(
'INSERT INTO "LiteLLM_SpendLogDelivery" '
'(request_id, accounting_payload, is_batch, state, applied_at) '
"VALUES ('expired', '{}'::jsonb, false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days')"
)
assert await other_leader.acquire_lock(cronjob_id="accounting_delivery", allow_reentrant=False)
await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj)
with psycopg.connect(parity_database) as connection:
assert connection.execute('SELECT attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(0,)]
assert connection.execute(
'SELECT request_id, state::text, attempts FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id'
).fetchall() == [("expired", "applied", 0), ("fresh", "applied", 0), ("retry", "applied", 1)]
assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0.75,)
assert REGISTRY.get_sample_value("litellm_accounting_batch_seconds_count") == reports_before
assert await writer.pod_lock_manager.acquire_lock(cronjob_id="accounting_delivery", allow_reentrant=False) is False
await other_leader.release_lock(cronjob_id="accounting_delivery")
await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj)
with psycopg.connect(parity_database) as connection:
assert connection.execute('SELECT attempts FROM "LiteLLM_SpendLogDelivery"').fetchall() == [(1,)]
assert connection.execute(
'SELECT request_id, state::text, attempts FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id'
).fetchall() == [("fresh", "applied", 0), ("retry", "applied", 1)]
assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (0.75,)
assert REGISTRY.get_sample_value("litellm_accounting_batch_seconds_count") == reports_before + 1
assert await other_leader.acquire_lock(cronjob_id="accounting_delivery", allow_reentrant=False)
finally:
await other_leader.release_lock(cronjob_id="accounting_delivery")