From f6827a5312523a60e5d322cfa16ec2dbdb2b32a1 Mon Sep 17 00:00:00 2001 From: hxlir Date: Sun, 1 Mar 2026 10:57:25 +0200 Subject: [PATCH] fixed comments --- tests/guardrails_tests/test_aim_guardrail_base_url.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/guardrails_tests/test_aim_guardrail_base_url.py b/tests/guardrails_tests/test_aim_guardrail_base_url.py index 3e230bb32cc..105ddc89b17 100644 --- a/tests/guardrails_tests/test_aim_guardrail_base_url.py +++ b/tests/guardrails_tests/test_aim_guardrail_base_url.py @@ -1,21 +1,24 @@ import pytest +from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail import os -from litellm.proxy.guardrails.guardrail_hooks.aim.aim import CustomGuardrail + +# Set a dummy API key for testing +os.environ["AIM_API_KEY"] = "test-key" def test_aim_base_url_trailing_slash(): # Test with trailing slash - guardrail = CustomGuardrail(api_base="https://api.aim.security/") + guardrail = AimGuardrail(api_base="https://api.aim.security/") assert guardrail.api_base == "https://api.aim.security" assert guardrail.ws_api_base.startswith("wss://api.aim.security") # Test without trailing slash - guardrail2 = CustomGuardrail(api_base="https://api.aim.security") + guardrail2 = AimGuardrail(api_base="https://api.aim.security") assert guardrail2.api_base == "https://api.aim.security" assert guardrail2.ws_api_base.startswith("wss://api.aim.security") # Test with environment variable os.environ["AIM_API_BASE"] = "https://api.aim.security/" - guardrail3 = CustomGuardrail(api_base=None) + guardrail3 = AimGuardrail(api_base=None) assert guardrail3.api_base == "https://api.aim.security" assert guardrail3.ws_api_base.startswith("wss://api.aim.security") del os.environ["AIM_API_BASE"]