Skip to content

Commit

Permalink
try to fix CI
Browse files Browse the repository at this point in the history
  • Loading branch information
Nanguage committed Sep 28, 2024
1 parent 35d1a16 commit 02db959
Showing 1 changed file with 24 additions and 15 deletions.
39 changes: 24 additions & 15 deletions tests/test_dask_job.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

0 comments on commit 02db959

Please sign in to comment.