fix test_unpack_is_non_rootish

dask · Dec 22, 2023 · 570af51 · 570af51
1 parent 2140f88
commit 570af51
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/distributed/shuffle/tests/test_shuffle.py b/distributed/shuffle/tests/test_shuffle.py
@@ -2486,10 +2486,12 @@ def __init__(self, scheduler: Scheduler):
         self.in_barrier = asyncio.Event()
         self.block_barrier = asyncio.Event()
 
-    async def barrier(self, id: ShuffleId, run_id: int, consistent: bool) -> None:
+    async def barrier(
+        self, id: ShuffleId, run_id: int, consistent: bool, worker: None
+    ) -> None:
         self.in_barrier.set()
         await self.block_barrier.wait()
-        return await super().barrier(id, run_id, consistent)
+        return await super().barrier(id, run_id, consistent, worker)
 
 
 @gen_cluster(client=True)