Separate out tests

setup.py

@@ -12,6 +12,7 @@ extras_require = {
         "text_generation @ git+https://github.com/ml-energy/text_generation_energy@master",
     ],
     "benchmark": ["zeus-ml", "fschat==0.2.23", "tyro", "rich"],
+    "dev": ["pytest"],
 }
 
 extras_require["all"] = list(set(sum(extras_require.values(), [])))

spitfight/colosseum/client.py

@@ -1,10 +1,8 @@
 from __future__ import annotations
 
 import json
-import unittest
 import contextlib
 from uuid import uuid4, UUID
-from copy import deepcopy
 from typing import Generator, Literal
 
 import requests
@@ -92,15 +90,3 @@ def _check_response(response: requests.Response) -> None:
         raise gr.Error(response.json()["detail"])
     elif response.status_code >= 500:
         raise gr.Error("Failed to talk to our backend server. Please try again later.")
-
-
-class TestControllerClient(unittest.TestCase):
-    def test_new_uuid_on_deepcopy(self):
-        client = ControllerClient("http://localhost:8000")
-        clients = [client.fork() for _ in range(50)]
-        request_ids = [client.request_id for client in clients]
-        assert len(set(request_ids)) == len(request_ids)
-
-
-if __name__ == "__main__":
-    unittest.main()

spitfight/utils.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import time
 import heapq
 import asyncio
-import unittest
 from typing import TypeVar, Generic, AsyncGenerator, Any, Coroutine
 
 from fastapi.logger import logger
@@ -178,128 +177,3 @@ class TokenGenerationBuffer:
                 break
 
         return return_buffer or None
-
-
-
-class TestTokenGenerationBuffer(unittest.TestCase):
-    def test_basic1(self):
-        buffer = TokenGenerationBuffer(stop_str="stop")
-
-        buffer.append("hello")
-        self.assertEqual(buffer.pop(), "hello")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("world")
-        self.assertEqual(buffer.pop(), "world")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("stop")
-        self.assertEqual(buffer.pop(), None)
-        self.assertTrue(buffer.matched_stop_str)
-        self.assertEqual(buffer.pop(), None)
-        self.assertTrue(buffer.matched_stop_str)
-        self.assertEqual(buffer.pop(), None)
-        self.assertTrue(buffer.matched_stop_str)
-        self.assertEqual(buffer.pop(), None)
-        self.assertTrue(buffer.matched_stop_str)
-
-    def test_basic2(self):
-        buffer = TokenGenerationBuffer(stop_str="stop")
-
-        buffer.append("hi")
-        self.assertEqual(buffer.pop(), "hi")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("stole")
-        self.assertEqual(buffer.pop(), "stole")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("sto")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("ic")
-        self.assertEqual(buffer.pop(), "stoic")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("st")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("opper")
-        self.assertEqual(buffer.pop(), "stopper")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("sto")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("p")
-        self.assertEqual(buffer.pop(), None)
-        self.assertTrue(buffer.matched_stop_str)
-
-    def test_falcon1(self):
-        buffer = TokenGenerationBuffer(stop_str="\nUser")
-
-        buffer.append("Hi")
-        self.assertEqual(buffer.pop(), "Hi")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("!")
-        self.assertEqual(buffer.pop(), "!")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("\n")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("User")
-        self.assertEqual(buffer.pop(), None)
-        self.assertTrue(buffer.matched_stop_str)
-
-    def test_falcon2(self):
-        buffer = TokenGenerationBuffer(stop_str="\nUser")
-
-        buffer.append("\n")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("\n")
-        self.assertEqual(buffer.pop(), "\n")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("\n")
-        self.assertEqual(buffer.pop(), "\n")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("\n")
-        self.assertEqual(buffer.pop(), "\n")
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("User")
-        self.assertEqual(buffer.pop(), None)
-        self.assertEqual(buffer.pop(), None)
-        self.assertTrue(buffer.matched_stop_str)
-
-    def test_no_stop_str(self):
-        buffer = TokenGenerationBuffer(stop_str=None)
-
-        buffer.append("hello")
-        self.assertEqual(buffer.pop(), "hello")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("world")
-        self.assertEqual(buffer.pop(), "world")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-        buffer.append("\n")
-        self.assertEqual(buffer.pop(), "\n")
-        self.assertEqual(buffer.pop(), None)
-        self.assertFalse(buffer.matched_stop_str)
-
-
-if __name__ == "__main__":
-    unittest.main()

tests/colosseum/test_client.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from spitfight.colosseum.client import ControllerClient
+
+
+def test_new_uuid_on_deepcopy():
+    client = ControllerClient("http://localhost:8000")
+    clients = [client.fork() for _ in range(50)]
+    request_ids = [client.request_id for client in clients]
+    assert len(set(request_ids)) == len(request_ids)

tests/test_utils.py

@@ -0,0 +1,122 @@
+from __future__ import annotations
+
+from spitfight.utils import TokenGenerationBuffer
+
+
+def test_basic1():
+    buffer = TokenGenerationBuffer(stop_str="stop")
+
+    buffer.append("hello")
+    assert buffer.pop() == "hello"
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("world")
+    assert buffer.pop() == "world"
+    assert not buffer.matched_stop_str
+
+    buffer.append("stop")
+    assert buffer.pop() == None
+    assert buffer.matched_stop_str
+    assert buffer.pop() == None
+    assert buffer.matched_stop_str
+    assert buffer.pop() == None
+    assert buffer.matched_stop_str
+    assert buffer.pop() == None
+    assert buffer.matched_stop_str
+
+def test_basic2():
+    buffer = TokenGenerationBuffer(stop_str="stop")
+
+    buffer.append("hi")
+    assert buffer.pop() == "hi"
+    assert not buffer.matched_stop_str
+
+    buffer.append("stole")
+    assert buffer.pop() == "stole"
+    assert not buffer.matched_stop_str
+
+    buffer.append("sto")
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("ic")
+    assert buffer.pop() == "stoic"
+    assert not buffer.matched_stop_str
+
+    buffer.append("st")
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("opper")
+    assert buffer.pop() == "stopper"
+    assert not buffer.matched_stop_str
+
+    buffer.append("sto")
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("p")
+    assert buffer.pop() == None
+    assert buffer.matched_stop_str
+
+def test_falcon1():
+    buffer = TokenGenerationBuffer(stop_str="\nUser")
+
+    buffer.append("Hi")
+    assert buffer.pop() == "Hi"
+    assert not buffer.matched_stop_str
+
+    buffer.append("!")
+    assert buffer.pop() == "!"
+    assert not buffer.matched_stop_str
+
+    buffer.append("\n")
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("User")
+    assert buffer.pop() == None
+    assert buffer.matched_stop_str
+
+def test_falcon2():
+    buffer = TokenGenerationBuffer(stop_str="\nUser")
+
+    buffer.append("\n")
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("\n")
+    assert buffer.pop() == "\n"
+    assert not buffer.matched_stop_str
+
+    buffer.append("\n")
+    assert buffer.pop() == "\n"
+    assert not buffer.matched_stop_str
+
+    buffer.append("\n")
+    assert buffer.pop() == "\n"
+    assert not buffer.matched_stop_str
+
+    buffer.append("User")
+    assert buffer.pop() == None
+    assert buffer.pop() == None
+    assert buffer.matched_stop_str
+
+def test_no_stop_str():
+    buffer = TokenGenerationBuffer(stop_str=None)
+
+    buffer.append("hello")
+    assert buffer.pop() == "hello"
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("world")
+    assert buffer.pop() == "world"
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str
+
+    buffer.append("\n")
+    assert buffer.pop() == "\n"
+    assert buffer.pop() == None
+    assert not buffer.matched_stop_str