From f2b0b18f43581a3358037e2603ebe8f341934f89 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 22 Jan 2026 19:00:02 -0800
Subject: [PATCH] test fixes

---
 litellm/__init__.py                           |  6 ++--
 .../proxy/test_litellm_pre_call_utils.py      | 33 +++++++++++++------
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/litellm/__init__.py b/litellm/__init__.py
index 76377d9caa0..e5c09702b9b 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -376,8 +376,10 @@ public_model_groups_links: Dict[str, Union[str, Dict[str, Any]]] = {}
 priority_reservation: Optional[
     Dict[str, Union[float, "PriorityReservationDict"]]
 ] = None
-# priority_reservation_settings is lazy-loaded via __getattr__, but declared here for type checking
-priority_reservation_settings: Optional["PriorityReservationSettings"] = None
+# priority_reservation_settings is lazy-loaded via __getattr__
+# Only declare for type checking - at runtime __getattr__ handles it
+if TYPE_CHECKING:
+    priority_reservation_settings: Optional["PriorityReservationSettings"] = None
 
 
 ######## Networking Settings ########
diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
index c0a4fd2e5a2..b9485a2e4cb 100644
--- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
+++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py
@@ -1485,8 +1485,13 @@ def test_add_guardrails_from_policy_engine():
     Test that add_guardrails_from_policy_engine adds guardrails from matching policies
     and tracks applied policies in metadata.
     """
+    from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
     from litellm.proxy.policy_engine.policy_registry import get_policy_registry
-    from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails, PolicyScope
+    from litellm.types.proxy.policy_engine import (
+        Policy,
+        PolicyAttachment,
+        PolicyGuardrails,
+    )
 
     # Setup test data
     data = {
@@ -1501,19 +1506,25 @@ def test_add_guardrails_from_policy_engine():
         key_alias="my-key",
     )
 
-    # Setup mock policies in the registry (directly set parsed Policy objects)
-    registry = get_policy_registry()
-    registry._policies = {
+    # Setup mock policies in the registry (policies define WHAT guardrails to apply)
+    policy_registry = get_policy_registry()
+    policy_registry._policies = {
         "global-baseline": Policy(
             guardrails=PolicyGuardrails(add=["pii_blocker"]),
-            scope=PolicyScope(teams=["*"]),
         ),
         "healthcare": Policy(
             guardrails=PolicyGuardrails(add=["hipaa_audit"]),
-            scope=PolicyScope(teams=["healthcare-team"]),
         ),
     }
-    registry._initialized = True
+    policy_registry._initialized = True
+
+    # Setup attachments in the attachment registry (attachments define WHERE policies apply)
+    attachment_registry = get_attachment_registry()
+    attachment_registry._attachments = [
+        PolicyAttachment(policy="global-baseline", scope="*"),  # applies to all
+        PolicyAttachment(policy="healthcare", teams=["healthcare-team"]),  # applies to healthcare team
+    ]
+    attachment_registry._initialized = True
 
     # Call the function
     add_guardrails_from_policy_engine(
@@ -1532,6 +1543,8 @@ def test_add_guardrails_from_policy_engine():
     assert "global-baseline" in data["metadata"]["applied_policies"]
     assert "healthcare" in data["metadata"]["applied_policies"]
 
-    # Clean up registry
-    registry._policies = {}
-    registry._initialized = False
+    # Clean up registries
+    policy_registry._policies = {}
+    policy_registry._initialized = False
+    attachment_registry._attachments = []
+    attachment_registry._initialized = False