mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
43 lines
1.8 KiB
Python
43 lines
1.8 KiB
Python
from typing import Final
|
|
|
|
import pytest
|
|
|
|
from .containers import Containers, ready
|
|
from .database import Database
|
|
from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision
|
|
|
|
SPEND_ROWS: Final = 20_000
|
|
|
|
pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup]
|
|
|
|
|
|
def seed_spend_logs(database: Database, rows: int) -> None:
|
|
database.execute(
|
|
'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime") '
|
|
"SELECT 'upgrade-shape-' || g, 'acompletion', now() - (g || ' seconds')::interval, "
|
|
"now() - (g || ' seconds')::interval FROM generate_series(1, %s) AS g",
|
|
(rows,),
|
|
)
|
|
assert database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((rows,),)
|
|
|
|
|
|
class TestPopulatedDatabaseUpgrade:
|
|
def test_upgrade_completes_and_preserves_a_populated_spend_log(
|
|
self, containers: Containers, baseline_image: str, baseline_database: Database
|
|
) -> None:
|
|
with containers.using(baseline_image).start(baseline_database) as old:
|
|
ready((old,), baseline_database)
|
|
key, alias = provision(old)
|
|
seed_spend_logs(baseline_database, SPEND_ROWS)
|
|
before: Final = migration_names(baseline_database)
|
|
with containers.start(baseline_database) as new:
|
|
ready((new,), baseline_database)
|
|
assert_upgraded(before, migration_names(baseline_database))
|
|
confirm(new, key, alias)
|
|
assert_history_clean(baseline_database)
|
|
assert baseline_database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((SPEND_ROWS,),), (
|
|
"The upgrade lost spend rows"
|
|
)
|
|
assert baseline_database.query(
|
|
'SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE "startTime" IS NULL OR "endTime" IS NULL'
|
|
) == ((0,),), "The upgrade nulled timestamps on existing spend rows"
|