From 39a82611325b11e5e0d0b05679fd6ef73c74bb65 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Thu, 11 Apr 2024 22:35:44 +0000
Subject: [PATCH] fix: split chat and tool tests

---
 .../test_flash_llama_simple.json}             |  2 +-
 integration-tests/models/test_chat_llama.py   | 42 +++++++++++++++++++
 integration-tests/models/test_tools_llama.py  | 26 ------------
 3 files changed, 43 insertions(+), 27 deletions(-)
 rename integration-tests/models/__snapshots__/{test_tools_llama/test_flash_llama_grammar_no_tools.json => test_chat_llama/test_flash_llama_simple.json} (97%)
 create mode 100644 integration-tests/models/test_chat_llama.py

diff --git a/integration-tests/models/__snapshots__/test_tools_llama/test_flash_llama_grammar_no_tools.json b/integration-tests/models/__snapshots__/test_chat_llama/test_flash_llama_simple.json
similarity index 97%
rename from integration-tests/models/__snapshots__/test_tools_llama/test_flash_llama_grammar_no_tools.json
rename to integration-tests/models/__snapshots__/test_chat_llama/test_flash_llama_simple.json
index 153a508d..0ff874f1 100644
--- a/integration-tests/models/__snapshots__/test_tools_llama/test_flash_llama_grammar_no_tools.json
+++ b/integration-tests/models/__snapshots__/test_chat_llama/test_flash_llama_simple.json
@@ -13,7 +13,7 @@
       "usage": null
     }
   ],
-  "created": 1710795556,
+  "created": 1712874856,
   "id": "",
   "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
   "object": "text_completion",
diff --git a/integration-tests/models/test_chat_llama.py b/integration-tests/models/test_chat_llama.py
new file mode 100644
index 00000000..11419a0e
--- /dev/null
+++ b/integration-tests/models/test_chat_llama.py
@@ -0,0 +1,42 @@
+import pytest
+import json
+
+from text_generation.types import GrammarType
+
+
+@pytest.fixture(scope="module")
+def flash_llama_chat_handle(launcher):
+    with launcher(
+        "TinyLlama/TinyLlama-1.1B-Chat-v1.0", num_shard=2, disable_grammar_support=False
+    ) as handle:
+        yield handle
+
+
+@pytest.fixture(scope="module")
+async def flash_llama_chat(flash_llama_chat_handle):
+    await flash_llama_chat_handle.health(300)
+    return flash_llama_chat_handle.client
+
+
+@pytest.mark.private
+async def test_flash_llama_simple(flash_llama_chat, response_snapshot):
+    response = await flash_llama_chat.chat(
+        max_tokens=100,
+        seed=1,
+        messages=[
+            {
+                "role": "system",
+                "content": "Youre a helpful assistant! Answer the users question best you can.",
+            },
+            {
+                "role": "user",
+                "content": "What is the weather like in Brooklyn, New York?",
+            },
+        ],
+    )
+
+    assert (
+        response.choices[0].message.content
+        == "As of today, there is a Update available for the Brooklyn, New York, area. According to the latest forecast, it's warm with high temperatures throughout the day. It's forecasted at 75°F for today and 77°F for tomorrow. However, in autumn, the weather typically changes drastically, becoming cooler and wetter. You can find the current weather forecast for the area through your local weather service. Additionally"
+    )
+    assert response == response_snapshot
diff --git a/integration-tests/models/test_tools_llama.py b/integration-tests/models/test_tools_llama.py
index ffcf232e..f18cb132 100644
--- a/integration-tests/models/test_tools_llama.py
+++ b/integration-tests/models/test_tools_llama.py
@@ -71,32 +71,6 @@ tools = [
 ]
 
 
-@pytest.mark.private
-async def test_flash_llama_grammar_no_tools(
-    flash_llama_grammar_tools, response_snapshot
-):
-    response = await flash_llama_grammar_tools.chat(
-        max_tokens=100,
-        seed=1,
-        messages=[
-            {
-                "role": "system",
-                "content": "Youre a helpful assistant! Answer the users question best you can.",
-            },
-            {
-                "role": "user",
-                "content": "What is the weather like in Brooklyn, New York?",
-            },
-        ],
-    )
-
-    assert (
-        response.choices[0].message.content
-        == "As of today, there is a Update available for the Brooklyn, New York, area. According to the latest forecast, it's warm with high temperatures throughout the day. It's forecasted at 75°F for today and 77°F for tomorrow. However, in autumn, the weather typically changes drastically, becoming cooler and wetter. You can find the current weather forecast for the area through your local weather service. Additionally"
-    )
-    assert response == response_snapshot
-
-
 @pytest.mark.asyncio
 @pytest.mark.private
 async def test_flash_llama_grammar_tools(flash_llama_grammar_tools, response_snapshot):