fix tests

src/litserve/server.py

@@ -428,6 +428,13 @@ def generate_client_file(port: Union[str, int] = 8000):
         except Exception as e:
             logger.exception(f"Error copying file: {e}")
 
+    def verify_worker_status(self):
+        while not any(v == WorkerSetupStatus.READY for v in self.workers_setup_status.values()):
+            if any(v == WorkerSetupStatus.ERROR for v in self.workers_setup_status.values()):
+                raise RuntimeError("One or more workers failed to start. Shutting down LitServe")
+            time.sleep(0.05)
+        logger.debug("One or more workers are ready to serve requests")
+
     def run(
         self,
         host: str = "0.0.0.0",
@@ -472,11 +479,8 @@ def run(
             api_server_worker_type = "process"
 
         manager, litserve_workers = self.launch_inference_worker(num_api_servers)
-        while not any(v == WorkerSetupStatus.READY for v in self.workers_setup_status.values()):
-            if any(v == WorkerSetupStatus.ERROR for v in self.workers_setup_status.values()):
-                raise RuntimeError("One or more workers failed to start. Shutting down LitServe")
-            time.sleep(0.05)
-        logger.debug("One or more workers are ready to serve requests")
+
+        self.verify_worker_status()
         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")

tests/test_lit_server.py

@@ -213,6 +213,7 @@ def test_start_server(mock_uvicon):
 def test_server_run_with_api_server_worker_type(mock_uvicorn):
     api = ls.test_examples.SimpleLitAPI()
     server = ls.LitServer(api, devices=1)
+    server.verify_worker_status = MagicMock()
     with pytest.raises(ValueError, match=r"Must be 'process' or 'thread'"):
         server.run(api_server_worker_type="invalid")
 
@@ -258,6 +259,7 @@ def test_server_run_windows(mock_uvicorn):
 
 def test_server_terminate():
     server = LitServer(SimpleLitAPI())
+    server.verify_worker_status = MagicMock()
     mock_manager = MagicMock()
 
     with patch("litserve.server.LitServer._start_server", side_effect=Exception("mocked error")) as mock_start, patch(