From c413686eadb0f63d2aca4d2cc68f62dfdfbb6a31 Mon Sep 17 00:00:00 2001
From: Sunny Wan <SunnyWan2020@gmail.com>
Date: Mon, 3 Mar 2025 17:49:11 -0500
Subject: [PATCH] wrote tests for snowflake

---
 snowflake_testing.py                    |  9 ---
 tests/llm_translation/test_snowflake.py | 78 +++++++++++++++++++++++++
 2 files changed, 78 insertions(+), 9 deletions(-)
 delete mode 100644 snowflake_testing.py
 create mode 100644 tests/llm_translation/test_snowflake.py

diff --git a/snowflake_testing.py b/snowflake_testing.py
deleted file mode 100644
index 7e779b5803..0000000000
--- a/snowflake_testing.py
+++ /dev/null
@@ -1,9 +0,0 @@
-import os
-from litellm import completion 
-
-os.environ["SNOWFLAKE_ACCOUNT_ID"] = "YOUR ACCOUNT"
-os.environ["SNOWFLAKE_JWT"] = "YOUR JWT"
-
-messages = [{"role": "user", "content": "Write me a poem about the blue sky"}]
-
-completion(model="snowflake/mistral-7b", messages=messages)
\ No newline at end of file
diff --git a/tests/llm_translation/test_snowflake.py b/tests/llm_translation/test_snowflake.py
new file mode 100644
index 0000000000..e64fb8d6d1
--- /dev/null
+++ b/tests/llm_translation/test_snowflake.py
@@ -0,0 +1,78 @@
+import os
+import sys
+import traceback
+from dotenv import load_dotenv
+
+load_dotenv()
+import pytest
+
+from litellm import completion, acompletion, set_verbose
+
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.asyncio
+async def test_chat_completion_snowflake(sync_mode):
+    try:
+        messages = [
+            {
+                "role": "user",
+                "content": "Write me a poem about the blue sky",
+            },
+        ]
+
+        if sync_mode:
+            response = completion(
+                model="snowflake/mistral-7b",
+                messages=messages,
+            )
+            print(response)
+            assert response is not None
+        else:
+            response = await acompletion(
+                model="snowflake/mistral-7b",
+                messages=messages,
+            )
+            print(response)
+            assert response is not None
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync_mode", [True, False])
+async def test_chat_completion_snowflake_stream(sync_mode):
+    try:
+        set_verbose = True
+        messages = [
+            {
+                "role": "user",
+                "content": "Write me a poem about the blue sky",
+            },
+        ]
+        
+        if sync_mode is False:
+            response = await acompletion(
+                model="snowflake/mistral-7b",
+                messages=messages,
+                max_tokens=100,
+                stream=True,
+            )
+            
+            chunk_count = 0
+            async for chunk in response:
+                print(chunk)
+                chunk_count += 1
+            assert chunk_count > 0
+        else:
+            response = completion(
+                model="snowflake/mistral-7b",
+                messages=messages,
+                max_tokens=100,
+                stream=True,
+            )
+                
+            chunk_count = 0
+            for chunk in response:
+                print(chunk)
+                chunk_count += 1
+            assert chunk_count > 0
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")