From 2e94df30bf067bf4b3433c79f92ffd3b21237a2e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 19:26:10 -0700 Subject: [PATCH] test: fix tests --- tests/local_testing/test_add_update_models.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/local_testing/test_add_update_models.py b/tests/local_testing/test_add_update_models.py index 4a04fcdf4a..1ea714d9b7 100644 --- a/tests/local_testing/test_add_update_models.py +++ b/tests/local_testing/test_add_update_models.py @@ -15,8 +15,12 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest, logging, asyncio -import litellm, asyncio -from litellm.proxy.proxy_server import add_new_model, update_model, LitellmUserRoles +import litellm +from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + update_model, +) +from litellm.proxy._types import LitellmUserRoles from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.proxy.management_endpoints.team_endpoints import new_team