Add unit tests for async setup

python/tests/server/fixtures/setup_async.py

@@ -0,0 +1,12 @@
+class Predictor:
+    async def download(self) -> None:
+        print("download complete!")
+
+    async def setup(self) -> None:
+        print("setup starting...")
+        await self.download()
+        print("setup complete!")
+
+    async def predict(self) -> str:
+        print("running prediction")
+        return "output"

python/tests/server/fixtures/setup_async_with_sync_predict.py

@@ -0,0 +1,9 @@
+class Predictor:
+    async def download(self) -> None:
+        print("setup used asyncio.run! it's not very effective...")
+
+    async def setup(self) -> None:
+        await self.download()
+
+    def predict(self) -> str:
+        return "output"

python/tests/server/test_worker.py

@@ -351,6 +351,57 @@ def test_async_predictor_on_python_3_10_or_older_raises_error(worker):
     )
 
 
+@uses_worker(
+    "setup_async", max_concurrency=1, min_python=(3, 11), is_async=True, setup=False
+)
+def test_setup_async(worker: Worker):
+    fut = worker.setup()
+    setup_result = Result()
+    setup_sid = worker.subscribe(setup_result.handle_event)
+
+    # with pytest.raises(FatalWorkerException):
+    fut.result()
+    worker.unsubscribe(setup_sid)
+
+    assert setup_result.stdout_lines == [
+        "setup starting...\n",
+        "download complete!\n",
+        "setup complete!\n",
+    ]
+
+    predict_result = Result()
+    predict_sid = worker.subscribe(predict_result.handle_event, tag="p1")
+    worker.predict({}, tag="p1").result()
+
+    assert predict_result.done
+    assert predict_result.output == "output"
+    assert predict_result.stdout_lines == ["running prediction\n"]
+
+    worker.unsubscribe(predict_sid)
+
+
+@uses_worker(
+    "setup_async_with_sync_predict",
+    max_concurrency=1,
+    min_python=(3, 11),
+    is_async=False,
+    setup=False,
+)
+def test_setup_async_with_sync_predict_raises_error(worker: Worker):
+    fut = worker.setup()
+    result = Result()
+    worker.subscribe(result.handle_event)
+
+    with pytest.raises(FatalWorkerException):
+        fut.result()
+    assert result.done
+    assert result.done.error
+    assert (
+        result.done.error_detail
+        == "Invalid predictor: to use an async setup method you must use an async predict method"
+    )
+
+
 @uses_worker("simple", max_concurrency=5, setup=False)
 def test_concurrency_with_sync_predictor_raises_error(worker):
     fut = worker.setup()