From 02db959fdb6e2d6a49a244f258d44630c7da79b4 Mon Sep 17 00:00:00 2001
From: nanguage <nanguage@yahoo.com>
Date: Sat, 28 Sep 2024 10:53:30 +0800
Subject: [PATCH] try to fix CI

---
 tests/test_dask_job.py | 39 ++++++++++++++++++++++++---------------
 1 file changed, 24 insertions(+), 15 deletions(-)

diff --git a/tests/test_dask_job.py b/tests/test_dask_job.py
index 59a77c0..5df7f3e 100644
--- a/tests/test_dask_job.py
+++ b/tests/test_dask_job.py
@@ -80,18 +80,27 @@ async def main():
 
 @pytest.mark.asyncio
 async def test_dask_generator():
-    with Engine() as engine:
-        async def gen():
-            for i in range(10):
-                yield i
-
-        job = DaskJob(gen)
-        await engine.submit_async(job)
-        await job.wait_until_status("running")
-        assert job.status == "running"
-        g = job.result()
-        i = 0
-        async for x in g:
-            assert x == i
-            i += 1
-        assert job.status == "done"
+    port = PortManager.find_free_port()
+    cluster = LocalCluster(
+        dashboard_address=f":{port}",
+        asynchronous=True,
+        processes=False,
+    )
+    client = Client(cluster)
+    engine = Engine()
+    engine.dask_client = client
+
+    async def gen():
+        for i in range(10):
+            yield i
+
+    job = DaskJob(gen)
+    await engine.submit_async(job)
+    await job.wait_until_status("running")
+    assert job.status == "running"
+    g = job.result()
+    i = 0
+    async for x in g:
+        assert x == i
+        i += 1
+    assert job.status == "done"