Add explicit iens to init mock

xjules · Jan 2, 2024 · 48ac288 · 48ac288
1 parent 513c3e3
commit 48ac288
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/unit_tests/scheduler/conftest.py b/tests/unit_tests/scheduler/conftest.py
@@ -12,8 +12,11 @@ def __init__(self, init=None, wait=None, kill=None):
 
     async def _init(self, iens, *args, **kwargs):
         if self._mock_init is not None:
-            await self._mock_init(*args, **kwargs)
-        return iens
+            if self._mock_init.__code__.co_argcount > 0:
+                await self._mock_init(iens, *args, **kwargs)
+                return iens
+            else:
+                await self._mock_init(*args, **kwargs)
 
     async def _wait(self, iens):
         if self._mock_wait is not None: