diff --git a/src/plumpy/message.py b/src/plumpy/message.py index 024d008f..7b2922eb 100644 --- a/src/plumpy/message.py +++ b/src/plumpy/message.py @@ -192,17 +192,16 @@ async def __call__(self, communicator: kiwipy.Communicator, task: Dict[str, Any] """ task_type = task[TASK_KEY] if task_type == LAUNCH_TASK: - return await self._launch(communicator, **task.get(TASK_ARGS, {})) + return await self._launch(**task.get(TASK_ARGS, {})) if task_type == CONTINUE_TASK: - return await self._continue(communicator, **task.get(TASK_ARGS, {})) + return await self._continue(**task.get(TASK_ARGS, {})) if task_type == CREATE_TASK: - return await self._create(communicator, **task.get(TASK_ARGS, {})) + return await self._create(**task.get(TASK_ARGS, {})) raise TaskRejectedError async def _launch( self, - _communicator: kiwipy.Communicator, process_class: str, persist: bool, nowait: bool, @@ -242,9 +241,7 @@ async def _launch( return proc.future().result() - async def _continue( - self, _communicator: kiwipy.Communicator, pid: 'PID_TYPE', nowait: bool, tag: Optional[str] = None - ) -> Union[PID_TYPE, Any]: + async def _continue(self, pid: 'PID_TYPE', nowait: bool, tag: Optional[str] = None) -> Union[PID_TYPE, Any]: """ Continue the process @@ -272,7 +269,6 @@ async def _continue( async def _create( self, - _communicator: kiwipy.Communicator, process_class: str, persist: bool, init_args: Optional[Sequence[Any]] = None, diff --git a/tests/test_message.py b/tests/test_message.py index 82951afd..0a6ee96c 100644 --- a/tests/test_message.py +++ b/tests/test_message.py @@ -37,7 +37,7 @@ async def test_continue(): del process process = None - result = await launcher._continue(None, **plumpy.create_continue_body(pid)[message.TASK_ARGS]) + result = await launcher._continue(**plumpy.create_continue_body(pid)[message.TASK_ARGS]) assert result == utils.DummyProcess.EXPECTED_OUTPUTS @@ -51,5 +51,5 @@ async def test_loader_is_used(): launcher = plumpy.ProcessLauncher(persister=persister, loader=loader) continue_task = plumpy.create_continue_body(proc.pid) - result = await launcher._continue(None, **continue_task[message.TASK_ARGS]) + result = await launcher._continue(**continue_task[message.TASK_ARGS]) assert result == utils.DummyProcess.EXPECTED_OUTPUTS