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

test(policy_engine): annotate pipeline test helper parameters

Browse source
This commit is contained in:
mateo-berri 2026-08-31 16:19:29 -07:00
parent e18c19f7ab
commit 55e71a53d2
1 changed files with 16 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -5,6 +5,7 @@ import os
import time
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import TYPE_CHECKING
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -44,6 +45,12 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.types.utils import CredentialItem
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry
from litellm.proxy.policy_engine.policy_registry import PolicyRegistry
from litellm.types.proxy.policy_engine import Policy
def test_check_if_token_is_service_account():
@ -4226,7 +4233,11 @@ async def test_add_guardrails_from_policy_engine():
attachment_registry._initialized = False
def _policy_engine_pipeline_registries(policies, monkeypatch, callbacks):
def _policy_engine_pipeline_registries(
policies: "dict[str, Policy]",
monkeypatch: pytest.MonkeyPatch,
callbacks: "list[CustomGuardrail]",
) -> "tuple[PolicyRegistry, AttachmentRegistry]":
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 PolicyAttachment
@ -4244,14 +4255,16 @@ def _policy_engine_pipeline_registries(policies, monkeypatch, callbacks):
return policy_registry, attachment_registry
def _reset_policy_engine_registries(policy_registry, attachment_registry):
def _reset_policy_engine_registries(
policy_registry: "PolicyRegistry", attachment_registry: "AttachmentRegistry"
) -> None:
policy_registry._policies = {}
policy_registry._initialized = False
attachment_registry._attachments = []
attachment_registry._initialized = False
def _word_guard_pipeline_policy(mode: str):
def _word_guard_pipeline_policy(mode: str) -> "Policy":
from litellm.types.proxy.policy_engine import (
GuardrailPipeline,
PipelineStep,