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"