From 4bf583010663a9b3de3b0473cdabc3dd93240813 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 26 Nov 2025 12:07:58 -0800 Subject: [PATCH] test fix --- .../guardrail_hooks/test_tool_permission.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 0926fe10c9e..5468dcf9491 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -6,12 +6,14 @@ import os import sys from unittest.mock import patch -from litellm.caching.dual_cache import DualCache import pytest +from litellm.caching.dual_cache import DualCache + sys.path.insert(0, os.path.abspath("../../../../../..")) from fastapi import HTTPException + from litellm.exceptions import GuardrailRaisedException from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ( @@ -22,9 +24,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import ( PermissionError, ) from litellm.types.utils import ( - ModelResponse, - Choices, ChatCompletionMessageToolCall, + Choices, + ModelResponse, ) @@ -179,6 +181,7 @@ class TestToolPermissionGuardrail: @pytest.mark.asyncio async def test_async_post_call_success_hook_no_tools(self): + """Test that async_post_call_success_hook returns response when no tool calls are present.""" response = ModelResponse(choices=[Choices(message={})]) user_api_key_dict = UserAPIKeyAuth() data = {"guardrails": ["test-tool-permission"]} @@ -187,10 +190,11 @@ class TestToolPermissionGuardrail: result = await self.guardrail.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response ) - assert result is None + assert result is response @pytest.mark.asyncio async def test_async_post_call_success_hook_with_allowed_tools(self): + """Test that async_post_call_success_hook returns response when tool calls are allowed.""" tool_call = { "function": {"name": "Bash", "arguments": "{}"}, "type": "function", @@ -203,7 +207,7 @@ class TestToolPermissionGuardrail: result = await self.guardrail.async_post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response ) - assert result is None + assert result is response @pytest.mark.asyncio async def test_async_post_call_success_hook_with_denied_tools_raises(self):