mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix: harden CORS, create_views exception handling, and spend log cleanup loop
- proxy_server.py: disable allow_credentials when allow_origins=['*'] (wildcard + credentials is a browser security misconfiguration). Add LITELLM_CORS_ORIGINS env var to configure explicit allowed origins. - create_views.py: narrow broad 'except Exception' to only catch genuine 'view does not exist' errors; re-raise all other DB errors (auth, connection, etc.) that were previously silently swallowed. - spend_log_cleanup.py: validate execute_raw() return type is int before using it as a deletion count; break loop safely on unexpected types to prevent infinite deletion loops.
This commit is contained in:
parent
4e12d3c562
commit
e079ee779f
6 changed files with 324 additions and 48 deletions
|
|
@ -18,14 +18,17 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
|
||||
If the view doesn't exist, one will be created.
|
||||
"""
|
||||
|
||||
try:
|
||||
# Try to select one row from the view
|
||||
await db.query_raw("""SELECT 1 FROM "LiteLLM_VerificationTokenView" LIMIT 1""")
|
||||
print("LiteLLM_VerificationTokenView Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("LiteLLM_VerificationTokenView Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
# If an error occurs, the view does not exist, so create it
|
||||
await db.execute_raw(
|
||||
"""
|
||||
await db.execute_raw("""
|
||||
CREATE VIEW "LiteLLM_VerificationTokenView" AS
|
||||
SELECT
|
||||
v.*,
|
||||
|
|
@ -37,15 +40,17 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
FROM "LiteLLM_VerificationToken" v
|
||||
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id
|
||||
LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id;
|
||||
"""
|
||||
)
|
||||
""")
|
||||
|
||||
print("LiteLLM_VerificationTokenView Created!") # noqa
|
||||
verbose_logger.debug("LiteLLM_VerificationTokenView Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""")
|
||||
print("MonthlyGlobalSpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("MonthlyGlobalSpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "MonthlyGlobalSpend" AS
|
||||
SELECT
|
||||
|
|
@ -60,12 +65,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("MonthlyGlobalSpend Created!") # noqa
|
||||
verbose_logger.debug("MonthlyGlobalSpend Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""")
|
||||
print("Last30dKeysBySpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("Last30dKeysBySpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "Last30dKeysBySpend" AS
|
||||
SELECT
|
||||
|
|
@ -88,12 +96,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("Last30dKeysBySpend Created!") # noqa
|
||||
verbose_logger.debug("Last30dKeysBySpend Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""")
|
||||
print("Last30dModelsBySpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("Last30dModelsBySpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "Last30dModelsBySpend" AS
|
||||
SELECT
|
||||
|
|
@ -111,11 +122,14 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("Last30dModelsBySpend Created!") # noqa
|
||||
verbose_logger.debug("Last30dModelsBySpend Created!")
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1""")
|
||||
print("MonthlyGlobalSpendPerKey Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerKey Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerKey" AS
|
||||
SELECT
|
||||
|
|
@ -132,13 +146,16 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("MonthlyGlobalSpendPerKey Created!") # noqa
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerKey Created!")
|
||||
try:
|
||||
await db.query_raw(
|
||||
"""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1"""
|
||||
)
|
||||
print("MonthlyGlobalSpendPerUserPerKey Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "MonthlyGlobalSpendPerUserPerKey" AS
|
||||
SELECT
|
||||
|
|
@ -157,12 +174,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("MonthlyGlobalSpendPerUserPerKey Created!") # noqa
|
||||
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""")
|
||||
print("DailyTagSpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("DailyTagSpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE OR REPLACE VIEW "DailyTagSpend" AS
|
||||
SELECT
|
||||
|
|
@ -175,12 +195,15 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("DailyTagSpend Created!") # noqa
|
||||
verbose_logger.debug("DailyTagSpend Created!")
|
||||
|
||||
try:
|
||||
await db.query_raw("""SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1""")
|
||||
print("Last30dTopEndUsersSpend Exists!") # noqa
|
||||
except Exception:
|
||||
verbose_logger.debug("Last30dTopEndUsersSpend Exists!")
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "does not exist" not in error_msg and "undefined" not in error_msg:
|
||||
raise
|
||||
sql_query = """
|
||||
CREATE VIEW "Last30dTopEndUsersSpend" AS
|
||||
SELECT end_user, COUNT(*) AS total_events, SUM(spend) AS total_spend
|
||||
|
|
@ -193,7 +216,7 @@ async def create_missing_views(db: _db): # noqa: PLR0915
|
|||
"""
|
||||
await db.execute_raw(query=sql_query)
|
||||
|
||||
print("Last30dTopEndUsersSpend Created!") # noqa
|
||||
verbose_logger.debug("Last30dTopEndUsersSpend Created!")
|
||||
|
||||
return
|
||||
|
||||
|
|
|
|||
|
|
@ -82,7 +82,7 @@ class SpendLogCleanup:
|
|||
break
|
||||
# Step 1: Find logs and delete them in one go without fetching to application
|
||||
# Delete in batches, limited by self.batch_size
|
||||
deleted_count = await prisma_client.db.execute_raw(
|
||||
deleted_result = await prisma_client.db.execute_raw(
|
||||
"""
|
||||
DELETE FROM "LiteLLM_SpendLogs"
|
||||
WHERE "request_id" IN (
|
||||
|
|
@ -94,6 +94,17 @@ class SpendLogCleanup:
|
|||
cutoff_date,
|
||||
self.batch_size,
|
||||
)
|
||||
|
||||
deleted_count = 0
|
||||
if isinstance(deleted_result, int):
|
||||
deleted_count = deleted_result
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
f"Unexpected execute_raw return type for spend log cleanup: {type(deleted_result)}; "
|
||||
"aborting cleanup to avoid infinite loop"
|
||||
)
|
||||
break
|
||||
|
||||
verbose_proxy_logger.info(f"Deleted {deleted_count} logs in this batch")
|
||||
|
||||
if deleted_count == 0:
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ from litellm.constants import (
|
|||
LITELLM_SETTINGS_SAFE_DB_OVERRIDES,
|
||||
LITELLM_UI_ALLOW_HEADERS,
|
||||
LITELLM_UI_SESSION_DURATION,
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_init_custom_logger_compatible_class,
|
||||
|
|
@ -1140,7 +1140,13 @@ async def openai_exception_handler(request: Request, exc: ProxyException):
|
|||
|
||||
|
||||
router = APIRouter()
|
||||
origins = ["*"]
|
||||
_cors_origins_env = os.getenv("LITELLM_CORS_ORIGINS")
|
||||
if _cors_origins_env is None or _cors_origins_env.strip() == "":
|
||||
origins = ["*"]
|
||||
else:
|
||||
origins = [o.strip() for o in _cors_origins_env.split(",") if o.strip()]
|
||||
|
||||
allow_cors_credentials = "*" not in origins
|
||||
|
||||
|
||||
# get current directory
|
||||
|
|
@ -1467,7 +1473,7 @@ current_dir = os.path.dirname(os.path.abspath(__file__))
|
|||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=origins,
|
||||
allow_credentials=True,
|
||||
allow_credentials=allow_cors_credentials,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
expose_headers=LITELLM_UI_ALLOW_HEADERS,
|
||||
|
|
@ -2316,9 +2322,13 @@ def _write_health_state_to_router_cache(
|
|||
|
||||
exception_status = getattr(original_exception, "status_code", 500)
|
||||
|
||||
if llm_router.health_check_ignore_transient_errors and exception_status in (
|
||||
429,
|
||||
408,
|
||||
if (
|
||||
llm_router.health_check_ignore_transient_errors
|
||||
and exception_status
|
||||
in (
|
||||
429,
|
||||
408,
|
||||
)
|
||||
):
|
||||
continue
|
||||
|
||||
|
|
@ -6287,7 +6297,9 @@ class ProxyStartupEvent:
|
|||
|
||||
### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ###
|
||||
## Reduces QPS as there are more tags for a single request
|
||||
tag_spend_update_interval = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER)
|
||||
tag_spend_update_interval = int(
|
||||
batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER
|
||||
)
|
||||
from litellm.proxy.utils import update_daily_tag_spend
|
||||
|
||||
scheduler.add_job(
|
||||
|
|
@ -7132,9 +7144,9 @@ async def chat_completion( # noqa: PLR0915
|
|||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
data["metadata"][
|
||||
"user_api_key_org_alias"
|
||||
] = user_api_key_dict.organization_alias
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
|
|
@ -7313,9 +7325,9 @@ async def completion( # noqa: PLR0915
|
|||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
data["metadata"][
|
||||
"user_api_key_org_alias"
|
||||
] = user_api_key_dict.organization_alias
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
|
|
@ -7562,9 +7574,9 @@ async def embeddings( # noqa: PLR0915
|
|||
hasattr(user_api_key_dict, "organization_alias")
|
||||
and user_api_key_dict.organization_alias is not None
|
||||
):
|
||||
data["metadata"]["user_api_key_org_alias"] = (
|
||||
user_api_key_dict.organization_alias
|
||||
)
|
||||
data["metadata"][
|
||||
"user_api_key_org_alias"
|
||||
] = user_api_key_dict.organization_alias
|
||||
if (
|
||||
hasattr(user_api_key_dict, "agent_id")
|
||||
and user_api_key_dict.agent_id is not None
|
||||
|
|
|
|||
112
tests/test_litellm/proxy/db/test_create_views.py
Normal file
112
tests/test_litellm/proxy/db/test_create_views.py
Normal file
|
|
@ -0,0 +1,112 @@
|
|||
"""
|
||||
Tests for create_missing_views exception handling fix.
|
||||
|
||||
Verifies that real DB errors (auth failures, connection errors, etc.)
|
||||
are re-raised instead of being silently swallowed, while genuine
|
||||
"view not found" errors still trigger view creation.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, call
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_reraises_connection_error():
|
||||
"""should re-raise exceptions that are NOT 'does not exist' errors (e.g. connection errors)."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=Exception("connection refused: unable to connect to database")
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock()
|
||||
|
||||
with pytest.raises(Exception, match="connection refused"):
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_reraises_permission_error():
|
||||
"""should re-raise permission denied errors, not treat them as missing views."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=Exception(
|
||||
"permission denied for table LiteLLM_VerificationTokenView"
|
||||
)
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock()
|
||||
|
||||
with pytest.raises(Exception, match="permission denied"):
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_creates_view_on_does_not_exist():
|
||||
"""should call execute_raw to create view when error contains 'does not exist'."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=[
|
||||
Exception('relation "LiteLLM_VerificationTokenView" does not exist'),
|
||||
None, # MonthlyGlobalSpend exists
|
||||
None, # Last30dKeysBySpend exists
|
||||
None, # Last30dModelsBySpend exists
|
||||
None, # MonthlyGlobalSpendPerKey exists
|
||||
None, # MonthlyGlobalSpendPerUserPerKey exists
|
||||
None, # DailyTagSpend exists
|
||||
None, # Last30dTopEndUsersSpend exists
|
||||
]
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock(return_value=None)
|
||||
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_called_once()
|
||||
created_sql = mock_db.execute_raw.call_args[0][0]
|
||||
assert 'CREATE VIEW "LiteLLM_VerificationTokenView"' in created_sql
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_creates_view_on_undefined_error():
|
||||
"""should treat 'undefined' errors as 'view not found' and attempt creation."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(
|
||||
side_effect=[
|
||||
Exception("undefined table LiteLLM_VerificationTokenView"),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
]
|
||||
)
|
||||
mock_db.execute_raw = AsyncMock(return_value=None)
|
||||
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_views_skips_creation_when_view_exists():
|
||||
"""should not call execute_raw when all views already exist."""
|
||||
from litellm.proxy.db.create_views import create_missing_views
|
||||
|
||||
mock_db = MagicMock()
|
||||
mock_db.query_raw = AsyncMock(return_value=[{"?column?": 1}])
|
||||
mock_db.execute_raw = AsyncMock()
|
||||
|
||||
await create_missing_views(mock_db)
|
||||
|
||||
mock_db.execute_raw.assert_not_called()
|
||||
79
tests/test_litellm/proxy/test_cors_config.py
Normal file
79
tests/test_litellm/proxy/test_cors_config.py
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
"""
|
||||
Tests for CORS configuration security fix.
|
||||
|
||||
Verifies that allow_credentials is automatically disabled when
|
||||
allow_origins=["*"] (wildcard) to prevent credentialed cross-origin
|
||||
requests from arbitrary origins.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
def _compute_cors_config(cors_origins_env):
|
||||
"""
|
||||
Mirror of the CORS config logic in proxy_server.py.
|
||||
Kept here so tests remain isolated from module-level side-effects.
|
||||
"""
|
||||
if cors_origins_env is None or cors_origins_env.strip() == "":
|
||||
origins = ["*"]
|
||||
else:
|
||||
origins = [o.strip() for o in cors_origins_env.split(",") if o.strip()]
|
||||
allow_cors_credentials = "*" not in origins
|
||||
return origins, allow_cors_credentials
|
||||
|
||||
|
||||
def test_cors_wildcard_disables_credentials():
|
||||
"""should disable credentials when LITELLM_CORS_ORIGINS is not set (defaults to wildcard)."""
|
||||
origins, allow_credentials = _compute_cors_config(None)
|
||||
assert origins == ["*"]
|
||||
assert allow_credentials is False
|
||||
|
||||
|
||||
def test_cors_empty_string_disables_credentials():
|
||||
"""should disable credentials when LITELLM_CORS_ORIGINS is an empty or whitespace string."""
|
||||
for empty in ("", " ", "\t"):
|
||||
origins, allow_credentials = _compute_cors_config(empty)
|
||||
assert origins == ["*"], f"Expected wildcard for input {repr(empty)}"
|
||||
assert (
|
||||
allow_credentials is False
|
||||
), f"Expected no credentials for input {repr(empty)}"
|
||||
|
||||
|
||||
def test_cors_single_specific_origin_enables_credentials():
|
||||
"""should enable credentials when a single explicit origin is configured."""
|
||||
origins, allow_credentials = _compute_cors_config("https://admin.example.com")
|
||||
assert origins == ["https://admin.example.com"]
|
||||
assert allow_credentials is True
|
||||
|
||||
|
||||
def test_cors_multiple_specific_origins_enables_credentials():
|
||||
"""should enable credentials and correctly parse comma-separated origins."""
|
||||
origins, allow_credentials = _compute_cors_config(
|
||||
"https://app.example.com, https://admin.example.com, https://api.example.com"
|
||||
)
|
||||
assert origins == [
|
||||
"https://app.example.com",
|
||||
"https://admin.example.com",
|
||||
"https://api.example.com",
|
||||
]
|
||||
assert allow_credentials is True
|
||||
|
||||
|
||||
def test_cors_wildcard_string_in_env_disables_credentials():
|
||||
"""should disable credentials when LITELLM_CORS_ORIGINS is explicitly set to '*'."""
|
||||
origins, allow_credentials = _compute_cors_config("*")
|
||||
assert "*" in origins
|
||||
assert allow_credentials is False
|
||||
|
||||
|
||||
def test_cors_origins_strips_whitespace():
|
||||
"""should strip surrounding whitespace from each origin entry."""
|
||||
origins, _ = _compute_cors_config(" https://a.com , https://b.com ")
|
||||
assert origins == ["https://a.com", "https://b.com"]
|
||||
|
||||
|
||||
def test_cors_origins_skips_blank_entries():
|
||||
"""should skip blank entries caused by trailing/double commas."""
|
||||
origins, allow_credentials = _compute_cors_config("https://a.com,,https://b.com,")
|
||||
assert origins == ["https://a.com", "https://b.com"]
|
||||
assert allow_credentials is True
|
||||
|
|
@ -287,9 +287,48 @@ def test_string_retention_still_works():
|
|||
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}"
|
||||
)
|
||||
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()
|
||||
mock_db = MagicMock()
|
||||
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)
|
||||
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
|
||||
|
||||
assert mock_db.execute_raw.call_count == 1
|
||||
assert total_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()
|
||||
mock_db = MagicMock()
|
||||
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)
|
||||
total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
|
||||
|
||||
assert mock_db.execute_raw.call_count == 3
|
||||
assert total_deleted == 800
|
||||
|
||||
|
||||
def test_cleanup_batch_size_env_var(monkeypatch):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue