From d02bad23c1ef7b1140ca2a0473855e474f8fba7b Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 5 Feb 2026 13:57:29 -0800 Subject: [PATCH] test fixes --- .../openai/test_moderations.py | 50 +++++++++---------- 1 file changed, 25 insertions(+), 25 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index b707805cfcb..cebba2ff5e1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -83,7 +83,9 @@ async def test_openai_moderation_guardrail_adds_to_litellm_callbacks(): @pytest.mark.asyncio async def test_openai_moderation_guardrail_safe_content(): - """Test OpenAI moderation guardrail with safe content""" + """Test OpenAI moderation guardrail with safe content via apply_guardrail""" + from litellm.types.utils import GenericGuardrailAPIInputs + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", @@ -122,23 +124,21 @@ async def test_openai_moderation_guardrail_safe_content(): ) with patch.object(guardrail, 'async_make_request', return_value=mock_response): - # Test pre-call hook with safe content - user_api_key_dict = UserAPIKeyAuth(api_key="test") - data = { - "messages": [ + # Test apply_guardrail with safe content using structured_messages + inputs = GenericGuardrailAPIInputs( + structured_messages=[ {"role": "user", "content": "Hello, how are you today?"} ] - } - - result = await guardrail.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=None, - data=data, - call_type="completion" ) - # Should return the original data unchanged - assert result == data + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"messages": [{"role": "user", "content": "Hello, how are you today?"}]}, + input_type="request" + ) + + # Should return the original inputs unchanged + assert result == inputs @pytest.mark.asyncio @@ -201,7 +201,9 @@ async def test_openai_moderation_guardrail_apply_guardrail(): @pytest.mark.asyncio async def test_openai_moderation_guardrail_harmful_content(): - """Test OpenAI moderation guardrail with harmful content""" + """Test OpenAI moderation guardrail with harmful content via apply_guardrail""" + from litellm.types.utils import GenericGuardrailAPIInputs + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", @@ -240,22 +242,20 @@ async def test_openai_moderation_guardrail_harmful_content(): ) with patch.object(guardrail, 'async_make_request', return_value=mock_response): - # Test pre-call hook with harmful content - user_api_key_dict = UserAPIKeyAuth(api_key="test") - data = { - "messages": [ + # Test apply_guardrail with harmful content using structured_messages + inputs = GenericGuardrailAPIInputs( + structured_messages=[ {"role": "user", "content": "This is hateful content"} ] - } + ) # Should raise HTTPException from fastapi import HTTPException with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=None, - data=data, - call_type="completion" + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"messages": [{"role": "user", "content": "This is hateful content"}]}, + input_type="request" ) assert exc_info.value.status_code == 400