1# coding: utf-8 2"""Mock objects for unit tests.""" 3 4from .nodes import Status 5from .tasks import AbinitTask 6from .flows import Flow 7 8 9def change_task_start(task, mocked_status="Error"): 10 """Return a AbinitTaskMockedStart object.""" 11 task.__class__ = AbinitTaskMockedStart 12 task.mocked_status = Status.as_status(mocked_status) 13 return task 14 15 16class AbinitTaskMockedStart(AbinitTask): 17 """A Task whose status is always self.mocked_status.""" 18 def start(self, **kwargs): 19 self.set_status(self.mocked_status, msg="Mocking status with %s" % self.mocked_status) 20 return 1 21 22 23def infinite_flow(flow): 24 """Return an InfiniteFlow.""" 25 flow.__class__ = InfiniteFlow 26 return flow 27 28 29class InfiniteFlow(Flow): 30 """A Flow that will never reach `all_ok`""" 31 def check_status(self, **kwargs): 32 super().check_status(**kwargs) 33 34 for task in self.iflat_tasks(status=self.S_OK): 35 task.set_status(task.S_INIT) 36 task.reset_from_scratch() 37