test(test_router_timeout.py): fix test

This commit is contained in:
Krrish Dholakia 2024-07-19 19:00:55 -07:00
parent 335cf65452
commit 4de19cb116

View file

@ -1,8 +1,12 @@
#### What this tests #### #### What this tests ####
# This tests if the router timeout error handling during fallbacks # This tests if the router timeout error handling during fallbacks
import sys, os, time import asyncio
import traceback, asyncio import os
import sys
import time
import traceback
import pytest import pytest
sys.path.insert( sys.path.insert(
@ -12,9 +16,10 @@ sys.path.insert(
import os import os
from dotenv import load_dotenv
import litellm import litellm
from litellm import Router from litellm import Router
from dotenv import load_dotenv
load_dotenv() load_dotenv()
@ -37,6 +42,7 @@ def test_router_timeouts():
"litellm_params": { "litellm_params": {
"model": "claude-instant-1.2", "model": "claude-instant-1.2",
"api_key": "os.environ/ANTHROPIC_API_KEY", "api_key": "os.environ/ANTHROPIC_API_KEY",
"mock_response": "hello world",
}, },
"tpm": 20000, "tpm": 20000,
}, },
@ -90,7 +96,9 @@ def test_router_timeouts():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_router_timeouts_bedrock(): async def test_router_timeouts_bedrock():
import openai, uuid import uuid
import openai
# Model list for OpenAI and Anthropic models # Model list for OpenAI and Anthropic models
_model_list = [ _model_list = [