scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

cleanup

Browse source
This commit is contained in:
hxlir 2026-03-01 11:44:11 +02:00
parent 5ac1f92ea2
commit 8a997bbda8
1 changed files with 15 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
tests/guardrails_tests/test_aim_guardrail_base_url.py
View file

@ -1,24 +1,21 @@
import pytest
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
import os
# 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 = AimGuardrail(api_base="https://api.aim.security/")
@pytest.mark.parametrize("api_base", [
"https://api.aim.security/",
"https://api.aim.security",
])
def test_aim_base_url_trailing_slash(monkeypatch, api_base):
monkeypatch.setenv("AIM_API_KEY", "test-key")
guardrail = AimGuardrail(api_base=api_base)
assert guardrail.api_base == "https://api.aim.security"
assert guardrail.ws_api_base.startswith("wss://api.aim.security")
assert guardrail.ws_api_base == "wss://api.aim.security"
# Test without trailing slash
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 = 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"]
def test_aim_base_url_from_env(monkeypatch):
monkeypatch.setenv("AIM_API_KEY", "test-key")
monkeypatch.setenv("AIM_API_BASE", "https://api.aim.security/")
guardrail = AimGuardrail(api_base=None)
assert guardrail.api_base == "https://api.aim.security"
assert guardrail.ws_api_base == "wss://api.aim.security"