From 0676e6a460fdbf54492910ebca228cdd64dea0e9 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 26 Oct 2026 18:42:19 -0700 Subject: [PATCH] test(guardrails): key the test guardrail hook payloads by str --- .../policy_engine/test_pipeline_executor.py | 2 +- .../proxy_logging/test_guardrail_pipeline.py | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index c56b6364bb8..c5cb7b3444f 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1457,7 +1457,7 @@ class _LitellmMetadataOpeningGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 2426ee8b7f0..bd0480182b7 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1171,13 +1171,13 @@ class _PipelineStepGuardrail(CustomGuardrail): self.calls = 0 async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: self.calls += 1 return data async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes ) -> None: self.calls += 1 @@ -1252,8 +1252,8 @@ class _RaisingPipelineStepGuardrail(_PipelineStepGuardrail): self.error = error async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: self.calls += 1 raise self.error @@ -1312,8 +1312,8 @@ async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): async def async_pre_call_hook( - self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral - ) -> dict: + self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict[str, object], call_type: CallTypesLiteral + ) -> dict[str, object]: self.calls += 1 data["litellm_metadata"] = {} add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @@ -1366,7 +1366,7 @@ async def test_post_call_pipeline_steps_are_named_only_once_they_run( class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): async def async_post_call_success_hook( - self, data: dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes ) -> None: self.calls += 1 data["litellm_metadata"] = {"user_api_key_hash": "hashed"}