test LitServer.run (#248)

* remove dead code

* add tests

* update test

* update

* update

* merge master

* update

* add start_server test

* update msg

* remove dead code

* skip windows

* revert deadcode removal

Author: aniketmaurya

src/litserve/server.py

@@ -412,13 +412,17 @@ def run(
         if num_api_servers is None:
             num_api_servers = len(self.workers)
 
-        manager, litserve_workers = self.launch_inference_worker(num_api_servers)
+        if num_api_servers < 1:
+            raise ValueError("num_api_servers must be greater than 0")
 
         if sys.platform == "win32":
+            print("Windows does not support forking. Using threads api_server_worker_type will be set to 'thread'")
             api_server_worker_type = "thread"
         elif api_server_worker_type is None:
             api_server_worker_type = "process"
 
+        manager, litserve_workers = self.launch_inference_worker(num_api_servers)
+
         try:
             servers = self._start_server(port, num_api_servers, log_level, sockets, api_server_worker_type, **kwargs)
             print(f"Swagger UI is available at http://0.0.0.0:{port}/docs")
@@ -447,7 +451,7 @@ def _start_server(self, port, num_uvicorn_servers, log_level, sockets, uvicorn_w
             elif uvicorn_worker_type == "thread":
                 w = threading.Thread(target=server.run, args=(sockets,))
             else:
-                raise ValueError("Invalid value for uvicorn_worker_type. Must be 'process' or 'thread'")
+                raise ValueError("Invalid value for api_server_worker_type. Must be 'process' or 'thread'")
             w.start()
             servers.append(w)
         return servers

tests/test_lit_server.py

@@ -14,6 +14,8 @@
 import asyncio
 import pickle
 import re
+import sys
+
 from asgi_lifespan import LifespanManager
 from litserve import LitAPI
 from fastapi import Request, Response, HTTPException
@@ -196,6 +198,64 @@ def test_server_run(mock_uvicorn):
     mock_uvicorn.Config.assert_called()
 
 
+@pytest.mark.skipif(sys.platform == "win32", reason="Test is only for Unix")
+@patch("litserve.server.uvicorn")
+def test_start_server(mock_uvicon):
+    server = LitServer(ls.examples.TestAPI(), spec=ls.OpenAISpec())
+    sockets = MagicMock()
+    server._start_server(8000, 1, "info", sockets, "process")
+    mock_uvicon.Server.assert_called()
+    assert server.lit_spec.response_queue_id is not None, "response_queue_id must be generated"
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="Test is only for Unix")
+@patch("litserve.server.uvicorn")
+def test_server_run_with_api_server_worker_type(mock_uvicorn):
+    api = ls.examples.SimpleLitAPI()
+    server = ls.LitServer(api, devices=1)
+    with pytest.raises(ValueError, match=r"Must be 'process' or 'thread'"):
+        server.run(api_server_worker_type="invalid")
+
+    with pytest.raises(ValueError, match=r"must be greater than 0"):
+        server.run(num_api_servers=0)
+
+    server.launch_inference_worker = MagicMock(return_value=[MagicMock(), [MagicMock()]])
+    server._start_server = MagicMock()
+
+    # Running the method to test
+    server.run(api_server_worker_type=None)
+    server.launch_inference_worker.assert_called_with(1)
+    actual = server._start_server.call_args
+    assert actual[0][4] == "process", "Server should run in process mode"
+
+    server.run(api_server_worker_type="thread")
+    server.launch_inference_worker.assert_called_with(1)
+    actual = server._start_server.call_args
+    assert actual[0][4] == "thread", "Server should run in thread mode"
+
+    server.run(api_server_worker_type="process")
+    server.launch_inference_worker.assert_called_with(1)
+    actual = server._start_server.call_args
+    assert actual[0][4] == "process", "Server should run in process mode"
+
+    server.run(api_server_worker_type="process", num_api_servers=10)
+    server.launch_inference_worker.assert_called_with(10)
+
+
+@pytest.mark.skipif(sys.platform != "win32", reason="Test is only for Windows")
+@patch("litserve.server.uvicorn")
+def test_server_run_windows(mock_uvicorn):
+    api = ls.examples.SimpleLitAPI()
+    server = ls.LitServer(api)
+    server.launch_inference_worker = MagicMock(return_value=[MagicMock(), [MagicMock()]])
+    server._start_server = MagicMock()
+
+    # Running the method to test
+    server.run(api_server_worker_type=None)
+    actual = server._start_server.call_args
+    assert actual[0][4] == "thread", "Windows only supports thread mode"
+
+
 def test_server_terminate():
     server = LitServer(SimpleLitAPI())
     mock_manager = MagicMock()