From d83f0b02da52aa135b91b1c1e145ec313626b227 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Mon, 6 May 2024 07:14:33 -0700
Subject: [PATCH] test: fix local tests

---
 litellm/tests/test_completion.py | 74 ++++++++++++++++++++++++--------
 1 file changed, 55 insertions(+), 19 deletions(-)

diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py
index 4c24caf7173..c8a09180590 100644
--- a/litellm/tests/test_completion.py
+++ b/litellm/tests/test_completion.py
@@ -1441,13 +1441,18 @@ def test_completion_ollama_hosted():
 # test_completion_ollama_hosted()
 
 
-@pytest.mark.parametrize(("model"), [
+@pytest.mark.skip(reason="Local test")
+@pytest.mark.parametrize(
+    ("model"),
+    [
         "ollama/llama2",
         "ollama_chat/llama2",
-    ]
+    ],
 )
 def test_completion_ollama_function_call(model):
-    messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
+    messages = [
+        {"role": "user", "content": "What's the weather like in San Francisco?"}
+    ]
     tools = [
         {
             "type": "function",
@@ -1473,19 +1478,27 @@ def test_completion_ollama_function_call(model):
         response = litellm.completion(model=model, messages=messages, tools=tools)
         print(response)
         assert response.choices[0].message.tool_calls
-        assert response.choices[0].message.tool_calls[0].function.name == "get_current_weather"
+        assert (
+            response.choices[0].message.tool_calls[0].function.name
+            == "get_current_weather"
+        )
         assert response.choices[0].finish_reason == "tool_calls"
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
 
-@pytest.mark.parametrize(("model"), [
+@pytest.mark.skip(reason="Local test")
+@pytest.mark.parametrize(
+    ("model"),
+    [
         "ollama/llama2",
         "ollama_chat/llama2",
-    ]
+    ],
 )
 def test_completion_ollama_function_call_stream(model):
-    messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
+    messages = [
+        {"role": "user", "content": "What's the weather like in San Francisco?"}
+    ]
     tools = [
         {
             "type": "function",
@@ -1508,24 +1521,33 @@ def test_completion_ollama_function_call_stream(model):
     ]
     try:
         litellm.set_verbose = True
-        response = litellm.completion(model=model, messages=messages, tools=tools, stream=True)
+        response = litellm.completion(
+            model=model, messages=messages, tools=tools, stream=True
+        )
         print(response)
         first_chunk = next(response)
         assert first_chunk.choices[0].delta.tool_calls
-        assert first_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_weather"
+        assert (
+            first_chunk.choices[0].delta.tool_calls[0].function.name
+            == "get_current_weather"
+        )
         assert first_chunk.choices[0].finish_reason == "tool_calls"
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
 
-@pytest.mark.parametrize(("model"), [
+@pytest.mark.parametrize(
+    ("model"),
+    [
         "ollama/llama2",
         "ollama_chat/llama2",
-    ]
+    ],
 )
 @pytest.mark.asyncio
 async def test_acompletion_ollama_function_call(model):
-    messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
+    messages = [
+        {"role": "user", "content": "What's the weather like in San Francisco?"}
+    ]
     tools = [
         {
             "type": "function",
@@ -1548,23 +1570,32 @@ async def test_acompletion_ollama_function_call(model):
     ]
     try:
         litellm.set_verbose = True
-        response = await litellm.acompletion(model=model, messages=messages, tools=tools)
+        response = await litellm.acompletion(
+            model=model, messages=messages, tools=tools
+        )
         print(response)
         assert response.choices[0].message.tool_calls
-        assert response.choices[0].message.tool_calls[0].function.name == "get_current_weather"
+        assert (
+            response.choices[0].message.tool_calls[0].function.name
+            == "get_current_weather"
+        )
         assert response.choices[0].finish_reason == "tool_calls"
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
 
-@pytest.mark.parametrize(("model"), [
+@pytest.mark.parametrize(
+    ("model"),
+    [
         "ollama/llama2",
         "ollama_chat/llama2",
-    ]
+    ],
 )
 @pytest.mark.asyncio
 async def test_acompletion_ollama_function_call_stream(model):
-    messages = [{"role": "user", "content": "What's the weather like in San Francisco?"}]
+    messages = [
+        {"role": "user", "content": "What's the weather like in San Francisco?"}
+    ]
     tools = [
         {
             "type": "function",
@@ -1587,11 +1618,16 @@ async def test_acompletion_ollama_function_call_stream(model):
     ]
     try:
         litellm.set_verbose = True
-        response = await litellm.acompletion(model=model, messages=messages, tools=tools, stream=True)
+        response = await litellm.acompletion(
+            model=model, messages=messages, tools=tools, stream=True
+        )
         print(response)
         first_chunk = await anext(response)
         assert first_chunk.choices[0].delta.tool_calls
-        assert first_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_weather"
+        assert (
+            first_chunk.choices[0].delta.tool_calls[0].function.name
+            == "get_current_weather"
+        )
         assert first_chunk.choices[0].finish_reason == "tool_calls"
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")