test fix
This commit is contained in:
parent
3441b7b737
commit
bbfddd00e4
@ -23,65 +23,6 @@ from litellm import APIConnectionError, Router
|
||||
from unittest.mock import ANY
|
||||
|
||||
|
||||
async def test_router_init():
|
||||
"""
|
||||
1. Initializes clients on the router with 0
|
||||
2. Checks if client is still valid
|
||||
3. Checks if new client was initialized
|
||||
"""
|
||||
model_list = [
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
},
|
||||
"model_info": {"id": "1234"},
|
||||
"tpm": 100000,
|
||||
"rpm": 10000,
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4.1-mini",
|
||||
"api_key": os.getenv("OPENAI_API_KEY"),
|
||||
},
|
||||
"tpm": 100000,
|
||||
"rpm": 10000,
|
||||
},
|
||||
]
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": f"write a one sentence poem {time.time()}?"}
|
||||
]
|
||||
client_ttl_time = 2
|
||||
router = Router(
|
||||
model_list=model_list,
|
||||
redis_host=os.environ["REDIS_HOST"],
|
||||
redis_password=os.environ["REDIS_PASSWORD"],
|
||||
redis_port=os.environ["REDIS_PORT"],
|
||||
cache_responses=True,
|
||||
timeout=30,
|
||||
routing_strategy="simple-shuffle",
|
||||
client_ttl=client_ttl_time,
|
||||
)
|
||||
model = "gpt-3.5-turbo"
|
||||
cache_key = f"1234_async_client"
|
||||
## ASSERT IT EXISTS AT THE START ##
|
||||
assert router.cache.get_cache(key=cache_key) is not None
|
||||
response1 = await router.acompletion(model=model, messages=messages, temperature=1)
|
||||
await asyncio.sleep(client_ttl_time)
|
||||
## ASSERT IT'S CLEARED FROM CACHE ##
|
||||
assert router.cache.get_cache(key=cache_key, local_only=True) is None
|
||||
## ASSERT IT EXISTS AFTER RUNNING __GET_CLIENT() ##
|
||||
assert (
|
||||
router._get_client(
|
||||
deployment=model_list[0], client_type="async", kwargs={"stream": False}
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skip(
|
||||
reason="This test is not relevant to the current codebase. The default Azure AD workflow is used."
|
||||
)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user