mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 11:14:04 +00:00
63 lines
1.8 KiB
Python
63 lines
1.8 KiB
Python
# What is this?
|
|
## This tests the blocked user pre call hook for the proxy server
|
|
|
|
|
|
import sys, os, asyncio, time, random
|
|
from datetime import datetime
|
|
import traceback
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import os
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
import pytest
|
|
import litellm
|
|
from litellm.proxy.enterprise.enterprise_hooks.blocked_user_list import (
|
|
_ENTERPRISE_BlockedUserList,
|
|
)
|
|
from litellm import Router, mock_completion
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.proxy._types import UserAPIKeyAuth
|
|
from litellm.caching import DualCache
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_block_user_check():
|
|
"""
|
|
- Set a blocked user as a litellm module value
|
|
- Test to see if a call with that user id is made, an error is raised
|
|
- Test to see if a call without that user is passes
|
|
"""
|
|
litellm.blocked_user_list = ["user_id_1"]
|
|
|
|
blocked_user_obj = _ENTERPRISE_BlockedUserList()
|
|
|
|
_api_key = "sk-12345"
|
|
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
|
|
local_cache = DualCache()
|
|
|
|
## Case 1: blocked user id passed
|
|
try:
|
|
await blocked_user_obj.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=local_cache,
|
|
call_type="completion",
|
|
data={"user_id": "user_id_1"},
|
|
)
|
|
pytest.fail(f"Expected call to fail")
|
|
except Exception as e:
|
|
pass
|
|
|
|
## Case 2: normal user id passed
|
|
try:
|
|
await blocked_user_obj.async_pre_call_hook(
|
|
user_api_key_dict=user_api_key_dict,
|
|
cache=local_cache,
|
|
call_type="completion",
|
|
data={"user_id": "user_id_2"},
|
|
)
|
|
except Exception as e:
|
|
pytest.fail(f"An error occurred - {str(e)}")
|