Source code for doot.util.mock_gen

  1#!/usr/bin/env python3
  2"""
  3Some utility functions to more easily setup mocks
  4
  5
  6"""
  7# Imports:
  8from __future__ import annotations
  9
 10# ##-- stdlib imports
 11import datetime
 12import enum
 13import functools as ftz
 14import itertools as itz
 15import logging as logmod
 16import pathlib as pl
 17import re
 18import time
 19import types
 20import weakref
 21from importlib.metadata import EntryPoint
 22from unittest.mock import MagicMock, PropertyMock, create_autospec
 23from uuid import UUID, uuid1
 24
 25# ##-- end stdlib imports
 26
 27# ##-- 3rd party imports
 28from jgdv.structs.chainguard import ChainGuard
 29from jgdv.structs.strang import CodeReference
 30from jgdv.cli import ParamSpec
 31
 32# ##-- end 3rd party imports
 33
 34# ##-- 1st party imports
 35from doot.workflow._interface import QueueMeta_e
 36from doot.workflow import TaskName, TaskSpec, ActionSpec
 37
 38# ##-- end 1st party imports
 39
 40# ##-- types
 41# isort: off
 42import abc
 43import collections.abc
 44from typing import TYPE_CHECKING, cast, assert_type, assert_never
 45from typing import Generic, NewType, Any
 46# Protocols:
 47from typing import Protocol, runtime_checkable
 48# Typing Decorators:
 49from typing import no_type_check, final, override, overload
 50
 51if TYPE_CHECKING:
 52    from jgdv import Maybe
 53    from typing import Final
 54    from typing import ClassVar, LiteralString
 55    from typing import Never, Self, Literal
 56    from typing import TypeGuard
 57    from collections.abc import Iterable, Iterator, Callable, Generator
 58    from collections.abc import Sequence, Mapping, MutableMapping, Hashable
 59
 60##--|
 61from doot.workflow._interface import (Job_p, Task_p)
 62from doot.control.tracker._interface import WorkflowTracker_p
 63# isort: on
 64# ##-- end types
 65
 66##-- logging
 67logging = logmod.getLogger(__name__)
 68##-- end logging
 69
[docs] 70def _add_prop(m, name, val): 71 setattr(type(m), name, PropertyMock(return_value=val))
72
[docs] 73def task_network(tasks:dict): 74 built = [] 75 for name, [pre, post] in tasks.items(): 76 current = mock_task_spec(name, pre=pre, post=post) 77 built.append(current) 78 79 return built
80
[docs] 81def mock_task(name, spec=None, actions:int=1, **kwargs): 82 task_m = MagicMock(spec=Task_p, 83 name=name, 84 state={}, 85 **kwargs) 86 task_m.spec = spec or mock_task_spec(name=name, action_count=actions) 87 _add_prop(task_m, "name", TaskName(name)) 88 _add_prop(task_m, "actions", task_m.spec.actions) 89 return task_m
90
[docs] 91def mock_job(name, pre=None, post=None, spec=None, **kwargs): 92 task_m = MagicMock(spec=Job_p, 93 name=name, 94 state={}, 95 **kwargs) 96 _add_prop(task_m, "name", name) 97 task_m.spec = spec or mock_task_spec(name=name) 98 return task_m
99
[docs] 100def mock_task_spec(name="agroup::mockSpec", pre=None, post=None, action_count=1, extra=None, **kwargs): 101 extra = extra or {} 102 if "sleep" not in extra: 103 extra['sleep'] = 0.1 104 spec_m = MagicMock(TaskSpec(name=name), 105 actions=mock_action_specs(num=action_count), 106 extra=ChainGuard(extra), 107 priority=10, 108 queue_behaviour=QueueMeta_e.default, 109 depends_on=pre or [], 110 required_for=post or [], 111 setup=[], 112 cleanup=[], 113 ) 114 spec_m.name = TaskName(name) 115 return spec_m
116
[docs] 117def mock_action_specs(num=1) -> list: 118 results = [] 119 for x in range(num): 120 action_spec_m = MagicMock(spec=ActionSpec(), 121 args=[], 122 kwargs=ChainGuard()) 123 type(action_spec_m).__call__ = MagicMock(return_value=None) 124 results.append(action_spec_m) 125 126 return results
127
[docs] 128def mock_entry_point(name="basic", value=None): 129 m = MagicMock(spec=EntryPoint) 130 _add_prop(m, "name", name) 131 _add_prop(m, "value", value) 132 m.load = MagicMock(return_value=value) 133 return m
134
[docs] 135def mock_task_ctor(name="APretendClass", module="pretend", params=None, mocker=None): 136 137 class MockedSubClass(Task_p): 138 139 def __new__(cls, *args, **kwargs): 140 m = mocker.Mock(spec=cls) 141 _add_prop(mock_ctor, "name", name) 142 _add_prop(mock_ctor, "param_specs", params or []) 143 return m 144 145 @classmethod 146 def param_specs(cls): 147 return params or [] 148 149 mock_ctor = MockedSubClass 150 mock_ctor.__module__ = module 151 mock_ctor.__name__ = name 152 return mock_ctor
153
[docs] 154def mock_code_ref(returns=None): 155 code_ref_m = MagicMock(spec=CodeReference()) 156 code_ref_m.__call__ = MagicMock(return_value=returns) 157 return code_ref_m
158
[docs] 159def mock_param_spec(name, val, type=Any): 160 m = MagicMock(spec=ParamSpec(name=name, type=type), default=val, prefix="-") 161 162 return m
163
[docs] 164def mock_tracker(tasks): 165 tracker_m = MagicMock(spec=WorkflowTracker_p) 166 local_tasks = tasks[:] 167 168 def simple_pop(): 169 if bool(local_tasks): 170 return local_tasks.pop() 171 return None 172 173 tracker_m.next_for = simple_pop 174 return tracker_m