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