scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix: req changes on tests

Browse source
This commit is contained in:
Harshit28j 2026-03-16 16:38:53 +05:30
parent 7481473cd4
commit 8d4b5ed571
2 changed files with 38 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/proxy/guardrails/guardrail_endpoints.py
View file

@ -462,6 +462,7 @@ async def update_guardrail(
f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}"
)
# Rollback: restore previous guardrail data in DB
rollback_ok = False
try:
await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
@ -473,14 +474,20 @@ async def update_guardrail(
guardrail_id=guardrail_id,
guardrail=cast(Guardrail, existing_guardrail),
)
rollback_ok = True
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail DB entry after update failure"
)
status = (
"rolled back"
if rollback_ok
else "rollback also failed, DB/memory may be inconsistent"
)
raise HTTPException(
status_code=422,
detail=f"Guardrail update failed, rolled back: {update_error}",
detail=f"Guardrail update failed, {status}: {update_error}",
)
return result
@ -557,19 +564,26 @@ async def delete_guardrail(
f"Immediate sync: Failed to remove guardrail '{guardrail_name}' (ID: {guardrail_id}) from memory: {delete_error}"
)
# Rollback: re-create the DB entry with the ORIGINAL guardrail_id
rollback_ok = False
try:
await GUARDRAIL_REGISTRY.add_guardrail_to_db(
guardrail=cast(Guardrail, existing_guardrail),
prisma_client=prisma_client,
guardrail_id=guardrail_id,
)
rollback_ok = True
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail DB deletion after memory removal failure"
)
status = (
"rolled back"
if rollback_ok
else "rollback also failed, DB/memory may be inconsistent"
)
raise HTTPException(
status_code=422,
detail=f"Guardrail delete failed, rolled back: {delete_error}",
detail=f"Guardrail delete failed, {status}: {delete_error}",
)
return result
@ -1155,6 +1169,7 @@ async def patch_guardrail(
f"Immediate sync: Failed to update '{guardrail_name}' (ID: {guardrail_id}) in memory: {update_error}"
)
# Rollback: restore previous guardrail data in DB
rollback_ok = False
try:
await GUARDRAIL_REGISTRY.update_guardrail_in_db(
guardrail_id=guardrail_id,
@ -1165,14 +1180,20 @@ async def patch_guardrail(
IN_MEMORY_GUARDRAIL_HANDLER.sync_guardrail_from_db(
guardrail=cast(Guardrail, existing_guardrail),
)
rollback_ok = True
except Exception:
verbose_proxy_logger.error(
"Failed to rollback guardrail DB entry after patch failure"
)
status = (
"rolled back"
if rollback_ok
else "rollback also failed, DB/memory may be inconsistent"
)
raise HTTPException(
status_code=422,
detail=f"Guardrail patch failed, rolled back: {update_error}",
detail=f"Guardrail patch failed, {status}: {update_error}",
)
return result

16
tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py
View file

@ -837,7 +837,11 @@ async def test_update_guardrail_endpoint(
assert "Prisma client not initialized" in str(exc_info.value.detail)
elif scenario == "success_sync_fails":
assert exc_info.value.status_code == 422
assert "rolled back" in str(exc_info.value.detail)
assert "Guardrail update failed" in str(exc_info.value.detail)
# Verify rollback was attempted: update_in_db called twice (initial + rollback)
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
# Verify rollback attempted in-memory re-init with old config
assert mock_in_memory_handler.update_in_memory_guardrail.call_count == 2
else:
result = await update_guardrail(
"test-guardrail-id", MOCK_UPDATE_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
@ -942,7 +946,11 @@ async def test_patch_guardrail_endpoint(
if scenario == "success_sync_fails":
assert exc_info.value.status_code == 422
assert "rolled back" in str(exc_info.value.detail)
assert "Guardrail patch failed" in str(exc_info.value.detail)
# Verify rollback was attempted: update_in_db called twice (initial + rollback)
assert mock_guardrail_registry.update_guardrail_in_db.call_count == 2
# Verify rollback attempted in-memory re-sync with old config
assert mock_in_memory_handler.sync_guardrail_from_db.call_count == 2
else:
result = await patch_guardrail(
"test-guardrail-id", MOCK_PATCH_REQUEST, user_api_key_dict=MOCK_ADMIN_USER
@ -1017,6 +1025,10 @@ async def test_delete_guardrail_endpoint(
if scenario == "success_sync_fails":
assert exc_info.value.status_code == 422
assert "rolled back" in str(exc_info.value.detail)
# Verify rollback was attempted: add_guardrail_to_db called with original ID
mock_guardrail_registry.add_guardrail_to_db.assert_called_once()
call_kwargs = mock_guardrail_registry.add_guardrail_to_db.call_args
assert call_kwargs.kwargs.get("guardrail_id") == expected_result
else:
result = await delete_guardrail(
guardrail_id=expected_result, user_api_key_dict=MOCK_ADMIN_USER