mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
39e31958f8
* test(proxy): move auth, hooks, policy_engine and client tests into tests/unit/proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): stub HIBP through respx by disabling the aiohttp transport Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): share the httpx transport fixture across proxy unit tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): restore proxy globals without a missing-value sentinel Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): package moved dirs and stub the login breach check at the HTTP boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): isolate the mcp server manager per test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
633 lines
24 KiB
Python
633 lines
24 KiB
Python
"""
|
|
Unit tests for policy versioning: registry behavior, status transitions, and version CRUD.
|
|
"""
|
|
|
|
from datetime import datetime, timezone
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from litellm.proxy.policy_engine.policy_registry import (
|
|
PolicyRegistry,
|
|
_row_to_policy_db_response,
|
|
get_policy_registry,
|
|
)
|
|
from litellm.types.proxy.policy_engine import (
|
|
Policy,
|
|
PolicyCreateRequest,
|
|
PolicyDBResponse,
|
|
PolicyGuardrails,
|
|
PolicyUpdateRequest,
|
|
)
|
|
|
|
|
|
def _make_row(
|
|
policy_id="pid-1",
|
|
policy_name="test-policy",
|
|
version_number=1,
|
|
version_status="production",
|
|
parent_version_id=None,
|
|
is_latest=True,
|
|
published_at=None,
|
|
production_at=None,
|
|
inherit=None,
|
|
description="desc",
|
|
guardrails_add=None,
|
|
guardrails_remove=None,
|
|
condition=None,
|
|
pipeline=None,
|
|
created_at=None,
|
|
updated_at=None,
|
|
created_by=None,
|
|
updated_by=None,
|
|
):
|
|
row = MagicMock()
|
|
row.policy_id = policy_id
|
|
row.policy_name = policy_name
|
|
row.version_number = version_number
|
|
row.version_status = version_status
|
|
row.parent_version_id = parent_version_id
|
|
row.is_latest = is_latest
|
|
row.published_at = published_at
|
|
row.production_at = production_at
|
|
row.inherit = inherit
|
|
row.description = description
|
|
row.guardrails_add = guardrails_add or []
|
|
row.guardrails_remove = guardrails_remove or []
|
|
row.condition = condition
|
|
row.pipeline = pipeline
|
|
row.created_at = created_at or datetime.now(timezone.utc)
|
|
row.updated_at = updated_at or datetime.now(timezone.utc)
|
|
row.created_by = created_by
|
|
row.updated_by = updated_by
|
|
return row
|
|
|
|
|
|
class TestRowToPolicyDBResponse:
|
|
"""Test _row_to_policy_db_response includes all version fields."""
|
|
|
|
def test_includes_version_fields(self):
|
|
row = _make_row(
|
|
version_number=2,
|
|
version_status="draft",
|
|
parent_version_id="pid-0",
|
|
is_latest=True,
|
|
published_at=None,
|
|
production_at=None,
|
|
)
|
|
resp = _row_to_policy_db_response(row)
|
|
assert isinstance(resp, PolicyDBResponse)
|
|
assert resp.policy_id == "pid-1"
|
|
assert resp.policy_name == "test-policy"
|
|
assert resp.version_number == 2
|
|
assert resp.version_status == "draft"
|
|
assert resp.parent_version_id == "pid-0"
|
|
assert resp.is_latest is True
|
|
assert resp.published_at is None
|
|
assert resp.production_at is None
|
|
|
|
def test_backward_compat_missing_version_attrs(self):
|
|
row = _make_row()
|
|
del row.version_number
|
|
del row.version_status
|
|
del row.parent_version_id
|
|
del row.is_latest
|
|
del row.published_at
|
|
del row.production_at
|
|
resp = _row_to_policy_db_response(row)
|
|
assert resp.version_number == 1
|
|
assert resp.version_status == "production"
|
|
assert resp.parent_version_id is None
|
|
assert resp.is_latest is True
|
|
|
|
|
|
class TestSyncPoliciesFromDbProductionOnly:
|
|
"""Test that sync_policies_from_db only loads production versions."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_all_policies_with_version_status_calls_find_many_with_where(
|
|
self,
|
|
):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
prod_row = _make_row(policy_id="prod-1", version_status="production")
|
|
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row])
|
|
|
|
result = await registry.get_all_policies_from_db(
|
|
prisma, version_status="production"
|
|
)
|
|
|
|
assert len(result) == 1
|
|
assert result[0].version_status == "production"
|
|
prisma.db.litellm_policytable.find_many.assert_called_once()
|
|
call_kw = prisma.db.litellm_policytable.find_many.call_args[1]
|
|
assert call_kw.get("where") == {"version_status": "production"}
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_policies_from_db_only_loads_production(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
prod_row = _make_row(
|
|
policy_id="prod-1",
|
|
policy_name="foo",
|
|
version_status="production",
|
|
guardrails_add=["g1"],
|
|
)
|
|
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row])
|
|
|
|
await registry.sync_policies_from_db(prisma)
|
|
|
|
assert registry.has_policy("foo")
|
|
policy = registry.get_policy("foo")
|
|
assert policy is not None
|
|
assert policy.guardrails.add == ["g1"]
|
|
# find_many was called with version_status=production (via get_all_policies_from_db)
|
|
find_many_calls = prisma.db.litellm_policytable.find_many.call_args_list
|
|
assert len(find_many_calls) >= 1
|
|
assert find_many_calls[0][1].get("where") == {"version_status": "production"}
|
|
|
|
|
|
class TestUpdatePolicyDraftOnly:
|
|
"""Test that update_policy_in_db only allows draft versions."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_production_raises(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
prod_row = _make_row(policy_id="pid-1", version_status="production")
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row)
|
|
|
|
with pytest.raises(Exception, match='Error updating policy in DB: Only draft versions can be') as exc_info:
|
|
await registry.update_policy_in_db(
|
|
policy_id="pid-1",
|
|
policy_request=PolicyUpdateRequest(description="new"),
|
|
prisma_client=prisma,
|
|
)
|
|
assert (
|
|
"Only draft" in str(exc_info.value)
|
|
or "draft" in str(exc_info.value).lower()
|
|
)
|
|
prisma.db.litellm_policytable.update.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_draft_succeeds_and_does_not_update_registry(self):
|
|
registry = PolicyRegistry()
|
|
registry.add_policy("test-policy", MagicMock()) # in-memory state
|
|
prisma = MagicMock()
|
|
draft_row = _make_row(
|
|
policy_id="draft-1",
|
|
policy_name="test-policy",
|
|
version_status="draft",
|
|
description="old",
|
|
)
|
|
updated_row = _make_row(
|
|
policy_id="draft-1",
|
|
policy_name="test-policy",
|
|
version_status="draft",
|
|
description="new",
|
|
)
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft_row)
|
|
prisma.db.litellm_policytable.update = AsyncMock(return_value=updated_row)
|
|
|
|
result = await registry.update_policy_in_db(
|
|
policy_id="draft-1",
|
|
policy_request=PolicyUpdateRequest(description="new"),
|
|
prisma_client=prisma,
|
|
)
|
|
|
|
assert result.description == "new"
|
|
prisma.db.litellm_policytable.update.assert_called_once()
|
|
# Registry still has old in-memory policy (drafts are not in registry; we don't add)
|
|
assert registry.has_policy("test-policy")
|
|
|
|
|
|
class TestDeletePolicyFromDb:
|
|
"""Test delete_policy_from_db removes production from registry and returns warning."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_production_removes_from_registry_and_returns_warning(self):
|
|
registry = PolicyRegistry()
|
|
registry.add_policy("deleted-policy", MagicMock())
|
|
prisma = MagicMock()
|
|
prod_row = _make_row(
|
|
policy_id="prod-1",
|
|
policy_name="deleted-policy",
|
|
version_status="production",
|
|
)
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row)
|
|
prisma.db.litellm_policytable.delete = AsyncMock()
|
|
|
|
result = await registry.delete_policy_from_db(
|
|
policy_id="prod-1",
|
|
prisma_client=prisma,
|
|
)
|
|
|
|
assert result["message"]
|
|
assert "warning" in result
|
|
assert "Production" in result["warning"] or "production" in result["warning"]
|
|
assert not registry.has_policy("deleted-policy")
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_draft_does_not_remove_from_registry_no_warning(self):
|
|
registry = PolicyRegistry()
|
|
registry.add_policy("my-policy", MagicMock())
|
|
prisma = MagicMock()
|
|
draft_row = _make_row(
|
|
policy_id="draft-1",
|
|
policy_name="my-policy",
|
|
version_status="draft",
|
|
)
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft_row)
|
|
prisma.db.litellm_policytable.delete = AsyncMock()
|
|
|
|
result = await registry.delete_policy_from_db(
|
|
policy_id="draft-1",
|
|
prisma_client=prisma,
|
|
)
|
|
|
|
assert "warning" not in result
|
|
assert registry.has_policy("my-policy")
|
|
|
|
|
|
class TestCreateNewVersion:
|
|
"""Test create_new_version copies all fields and sets draft."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_new_version_from_production_increments_version(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
prod = _make_row(
|
|
policy_id="prod-1",
|
|
policy_name="foo",
|
|
version_number=1,
|
|
version_status="production",
|
|
guardrails_add=["g1"],
|
|
description="base",
|
|
inherit=None,
|
|
pipeline={"mode": "pre_call", "steps": []},
|
|
)
|
|
# find_first for production
|
|
prisma.db.litellm_policytable.find_first = AsyncMock(return_value=prod)
|
|
# find_first for latest version number
|
|
prisma.db.litellm_policytable.find_first.side_effect = [
|
|
prod, # production lookup
|
|
prod, # latest version_number lookup
|
|
]
|
|
# update_many for is_latest=False
|
|
prisma.db.litellm_policytable.update_many = AsyncMock()
|
|
new_row = _make_row(
|
|
policy_id="new-id",
|
|
policy_name="foo",
|
|
version_number=2,
|
|
version_status="draft",
|
|
parent_version_id="prod-1",
|
|
is_latest=True,
|
|
guardrails_add=["g1"],
|
|
description="base",
|
|
pipeline={"mode": "pre_call", "steps": []},
|
|
)
|
|
prisma.db.litellm_policytable.create = AsyncMock(return_value=new_row)
|
|
|
|
result = await registry.create_new_version(
|
|
policy_name="foo",
|
|
prisma_client=prisma,
|
|
source_policy_id=None,
|
|
created_by="user",
|
|
)
|
|
|
|
assert result.version_number == 2
|
|
assert result.version_status == "draft"
|
|
assert result.parent_version_id == "prod-1"
|
|
assert result.guardrails_add == ["g1"]
|
|
assert result.description == "base"
|
|
create_call = prisma.db.litellm_policytable.create.call_args[1]["data"]
|
|
assert create_call["version_number"] == 2
|
|
assert create_call["version_status"] == "draft"
|
|
assert create_call["parent_version_id"] == "prod-1"
|
|
assert create_call["guardrails_add"] == ["g1"]
|
|
|
|
|
|
class TestUpdateVersionStatus:
|
|
"""Test status transitions: valid succeed, invalid return error."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_draft_to_published_sets_published_at(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
draft = _make_row(policy_id="d-1", version_status="draft")
|
|
updated = _make_row(
|
|
policy_id="d-1",
|
|
version_status="published",
|
|
published_at=datetime.now(timezone.utc),
|
|
)
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft)
|
|
prisma.db.litellm_policytable.update = AsyncMock(return_value=updated)
|
|
|
|
result = await registry.update_version_status(
|
|
policy_id="d-1",
|
|
new_status="published",
|
|
prisma_client=prisma,
|
|
)
|
|
|
|
assert result.version_status == "published"
|
|
update_data = prisma.db.litellm_policytable.update.call_args[1]["data"]
|
|
assert update_data["version_status"] == "published"
|
|
assert "published_at" in update_data
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_draft_to_production_raises(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
draft = _make_row(policy_id="d-1", version_status="draft")
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=draft)
|
|
|
|
with pytest.raises(Exception, match='Error updating version status: Cannot promote draft') as exc_info:
|
|
await registry.update_version_status(
|
|
policy_id="d-1",
|
|
new_status="production",
|
|
prisma_client=prisma,
|
|
)
|
|
assert (
|
|
"publish" in str(exc_info.value).lower()
|
|
or "draft" in str(exc_info.value).lower()
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_published_to_production_demotes_old_and_updates_registry(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
published_row = _make_row(
|
|
policy_id="pub-1",
|
|
policy_name="foo",
|
|
version_status="published",
|
|
)
|
|
updated_row = _make_row(
|
|
policy_id="pub-1",
|
|
policy_name="foo",
|
|
version_status="production",
|
|
production_at=datetime.now(timezone.utc),
|
|
)
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(
|
|
return_value=published_row
|
|
)
|
|
prisma.db.litellm_policytable.update_many = AsyncMock()
|
|
prisma.db.litellm_policytable.update = AsyncMock(return_value=updated_row)
|
|
|
|
result = await registry.update_version_status(
|
|
policy_id="pub-1",
|
|
new_status="production",
|
|
prisma_client=prisma,
|
|
)
|
|
|
|
assert result.version_status == "production"
|
|
# update_many should have been called to demote current production
|
|
assert prisma.db.litellm_policytable.update_many.called
|
|
# Registry should have been updated with new production
|
|
assert registry.has_policy("foo")
|
|
|
|
|
|
class TestCompareVersions:
|
|
"""Test compare_versions returns correct field diffs."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compare_versions_returns_diffs(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
a = _make_row(
|
|
policy_id="a",
|
|
policy_name="p",
|
|
description="desc A",
|
|
guardrails_add=["g1"],
|
|
)
|
|
b = _make_row(
|
|
policy_id="b",
|
|
policy_name="p",
|
|
description="desc B",
|
|
guardrails_add=["g1", "g2"],
|
|
)
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(side_effect=[a, b])
|
|
|
|
result = await registry.compare_versions(
|
|
policy_id_a="a",
|
|
policy_id_b="b",
|
|
prisma_client=prisma,
|
|
)
|
|
|
|
assert result.version_a.policy_id == "a"
|
|
assert result.version_b.policy_id == "b"
|
|
assert "description" in result.field_diffs
|
|
assert result.field_diffs["description"]["version_a"] == "desc A"
|
|
assert result.field_diffs["description"]["version_b"] == "desc B"
|
|
assert "guardrails_add" in result.field_diffs
|
|
|
|
|
|
class TestResolveGuardrailsProductionOnly:
|
|
"""Test that resolve_guardrails_from_db uses only production versions."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resolve_guardrails_calls_get_all_with_production_filter(self):
|
|
registry = PolicyRegistry()
|
|
prisma = MagicMock()
|
|
prod_row = _make_row(
|
|
policy_name="base",
|
|
version_status="production",
|
|
guardrails_add=["g1"],
|
|
)
|
|
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[prod_row])
|
|
|
|
result = await registry.resolve_guardrails_from_db(
|
|
policy_name="base",
|
|
prisma_client=prisma,
|
|
)
|
|
|
|
assert "g1" in result
|
|
call_kw = prisma.db.litellm_policytable.find_many.call_args[1]
|
|
assert call_kw.get("where") == {"version_status": "production"}
|
|
|
|
|
|
class TestGetPolicyRegistrySingleton:
|
|
"""Test get_policy_registry returns same instance."""
|
|
|
|
def test_returns_singleton(self):
|
|
a = get_policy_registry()
|
|
b = get_policy_registry()
|
|
assert a is b
|
|
|
|
|
|
def _prisma_with_policy_rows(production_rows, non_production_rows=None):
|
|
prisma = MagicMock()
|
|
prisma.db.litellm_policytable.find_many = AsyncMock(side_effect=[production_rows, non_production_rows or []])
|
|
return prisma
|
|
|
|
|
|
class TestConfigPoliciesPreservedAcrossDbSync:
|
|
"""Config-defined policies must survive sync_policies_from_db (regression for issue #35255)."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_with_empty_db_preserves_config_policies(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"config-policy": {"description": "from config", "guardrails": {"add": ["tooling"]}}})
|
|
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
|
|
|
|
assert registry.has_policy("config-policy")
|
|
policy = registry.get_policy("config-policy")
|
|
assert policy is not None
|
|
assert policy.guardrails.add == ["tooling"]
|
|
assert registry.get_source("config-policy") == "config"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sync_merges_db_policies_with_config_policies(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
|
|
db_row = _make_row(policy_id="db-1", policy_name="db-policy", guardrails_add=["db-guard"])
|
|
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))
|
|
|
|
assert registry.has_policy("config-policy")
|
|
assert registry.has_policy("db-policy")
|
|
assert registry.get_source("config-policy") == "config"
|
|
assert registry.get_source("db-policy") == "db"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_db_wins_on_policy_name_conflict(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
|
|
db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"])
|
|
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))
|
|
|
|
policy = registry.get_policy("shared-name")
|
|
assert policy is not None
|
|
assert policy.guardrails.add == ["db-guard"]
|
|
assert registry.get_source("shared-name") == "db"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_config_policy_restored_after_conflicting_db_row_deleted(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
|
|
db_row = _make_row(policy_id="db-1", policy_name="shared-name", guardrails_add=["db-guard"])
|
|
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([db_row]))
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
|
|
|
|
policy = registry.get_policy("shared-name")
|
|
assert policy is not None
|
|
assert policy.guardrails.add == ["config-guard"]
|
|
assert registry.get_source("shared-name") == "config"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_config_policy_resolves_guardrails_after_sync(self):
|
|
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
|
|
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
|
|
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
|
|
|
|
resolved = PolicyResolver.resolve_policy_guardrails(
|
|
policy_name="config-policy",
|
|
policies=registry.get_all_policies(),
|
|
context=None,
|
|
)
|
|
assert resolved.guardrails == ["tooling"]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_add_policy_with_config_source_survives_sync(self):
|
|
registry = PolicyRegistry()
|
|
registry.add_policy(
|
|
"late-config-policy",
|
|
Policy(guardrails=PolicyGuardrails(add=["tooling"])),
|
|
source="config",
|
|
)
|
|
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
|
|
|
|
assert registry.has_policy("late-config-policy")
|
|
assert registry.get_source("late-config-policy") == "config"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_clear_removes_config_snapshot_so_sync_does_not_resurrect(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"config-policy": {"guardrails": {"add": ["tooling"]}}})
|
|
|
|
registry.clear()
|
|
await registry.sync_policies_from_db(_prisma_with_policy_rows([]))
|
|
|
|
assert not registry.has_policy("config-policy")
|
|
assert registry.get_source("config-policy") is None
|
|
|
|
|
|
class TestRemovePolicyRestoresConfigFallback:
|
|
"""Deleting a same-named DB override must re-activate the config policy immediately, not at the next sync."""
|
|
|
|
def test_remove_policy_restores_config_version_immediately(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
|
|
registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
|
|
|
|
assert registry.remove_policy("shared-name") is True
|
|
|
|
policy = registry.get_policy("shared-name")
|
|
assert policy is not None
|
|
assert policy.guardrails.add == ["config-guard"]
|
|
assert registry.get_source("shared-name") == "config"
|
|
|
|
def test_remove_policy_without_config_fallback_removes_entirely(self):
|
|
registry = PolicyRegistry()
|
|
registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"])))
|
|
|
|
assert registry.remove_policy("db-only") is True
|
|
|
|
assert not registry.has_policy("db-only")
|
|
assert registry.get_source("db-only") is None
|
|
|
|
def test_remove_missing_policy_returns_false(self):
|
|
registry = PolicyRegistry()
|
|
|
|
assert registry.remove_policy("missing") is False
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_production_override_reactivates_config_policy_and_says_so(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
|
|
registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
|
|
prisma = MagicMock()
|
|
prod_row = _make_row(policy_id="prod-1", policy_name="shared-name", version_status="production")
|
|
prisma.db.litellm_policytable.find_unique = AsyncMock(return_value=prod_row)
|
|
prisma.db.litellm_policytable.delete = AsyncMock()
|
|
|
|
result = await registry.delete_policy_from_db(policy_id="prod-1", prisma_client=prisma)
|
|
|
|
assert "config" in result["warning"]
|
|
policy = registry.get_policy("shared-name")
|
|
assert policy is not None
|
|
assert policy.guardrails.add == ["config-guard"]
|
|
assert registry.get_source("shared-name") == "config"
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_all_versions_reactivates_config_policy(self):
|
|
registry = PolicyRegistry()
|
|
registry.load_policies({"shared-name": {"guardrails": {"add": ["config-guard"]}}})
|
|
registry.add_policy("shared-name", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
|
|
prisma = MagicMock()
|
|
prisma.db.litellm_policytable.delete_many = AsyncMock()
|
|
|
|
result = await registry.delete_all_versions(policy_name="shared-name", prisma_client=prisma)
|
|
|
|
assert registry.get_source("shared-name") == "config"
|
|
policy = registry.get_policy("shared-name")
|
|
assert policy is not None
|
|
assert policy.guardrails.add == ["config-guard"]
|
|
assert "config" in result["warning"]
|
|
|
|
async def test_delete_all_versions_without_config_twin_has_no_warning(self):
|
|
registry = PolicyRegistry()
|
|
registry.add_policy("db-only", Policy(guardrails=PolicyGuardrails(add=["db-guard"])), source="db")
|
|
prisma = MagicMock()
|
|
prisma.db.litellm_policytable.delete_many = AsyncMock()
|
|
|
|
result = await registry.delete_all_versions(policy_name="db-only", prisma_client=prisma)
|
|
|
|
assert registry.get_policy("db-only") is None
|
|
assert "warning" not in result
|