def test_handlers_when_success_flow(param):
def func(x):
return x * 2
queue_handler = MagicMock()
start_handler = MagicMock()
success_handler = MagicMock()
scheduler = ThreadPoolScheduler()
expected_result = func(param)
task = Task(
func,
on_start_handlers=[start_handler],
on_queue_handlers=[queue_handler],
on_completion_success_handlers=[success_handler],
)
scheduler.add_task(task=task, params=param)
assert queue_handler.called
assert task.state == State.QUEUED
assert not success_handler.called
assert not start_handler.called
result = scheduler.results()
assert next(result) == expected_result
assert start_handler.called
assert success_handler.called
handler_args = success_handler.call_args_list[0][0]
assert handler_args[0] == expected_result
assert task.state == State.SUCCESS
def test_task_context_kv_store(param):
key_1 = "key_1"
value_1 = "value_1"
key_2 = "key_2"
value_2 = "value_2"
def func(x, context: TaskContext = None):
assert context.state == State.RUNNING
context[key_1] = value_1
context[key_2] = value_2
return x * 2
expected_result = param * 2
task = Task(func)
assert task.state == State.CREATED
result = task(param)
assert result == expected_result
assert result == task.result
assert task.state == State.SUCCESS
assert task.context[key_1] == value_1
assert task.context[key_2] == value_2
def test_scheduler_add_task_accepts_task(param):
scheduler = SerialScheduler()
task = Task(some_func)
scheduler.add_task(task, param)
_task, _param = next(scheduler.tasks)
assert isinstance(_task, Task)
assert task == _task
assert _param == param
assert _task.state == State.QUEUED
def test_task_execution_for_error_flow(param):
def func(x):
raise Exception
task = Task(func)
assert task.state == State.CREATED
result = task(param)
assert result is None
assert isinstance(task.error, Exception)
assert task.state == State.FAILED
def test_task_execution_for_success_flow(param):
expected_result = some_func(param)
task = Task(some_func)
assert task.state == State.CREATED
result = task(param)
assert result == expected_result
assert result == task.result
assert task.state == State.SUCCESS
def add_task(self,
task: Union[Task, Callable],
params: Any = None) -> None:
if isinstance(task, Task):
task = task
elif isinstance(task, Callable):
task = Task(task)
else:
raise InvalidTask("Expected Callable or instance of Task")
self.tasks = itertools.chain(self.tasks, ((task, params), ))
task.state = State.QUEUED
def test_thread_pool_scheduler_returns_success_result(param):
expected_result = some_func(param)
scheduler = ThreadPoolScheduler()
task = Task(some_func)
scheduler.add_task(task, param)
assert task.state == State.QUEUED
results = scheduler.results()
_result = next(results)
assert task.state == State.SUCCESS
assert _result == expected_result
def test_thread_pool_scheduler_returns_error_result(param):
def error_func(x):
raise Exception
scheduler = ThreadPoolScheduler()
task = Task(error_func)
scheduler.add_task(task, param)
assert task.state == State.QUEUED
results = scheduler.results()
_result = next(results)
assert task.state == State.FAILED
assert _result is None
def test_task_creation_with_context(param):
def func(x, context: TaskContext = None):
assert context.state == State.RUNNING
return x * 2
expected_result = param * 2
task = Task(func)
assert task.state == State.CREATED
result = task(param)
assert result == expected_result
assert result == task.result
assert task.state == State.SUCCESS
def test_handlers_when_error_flow(param):
class CustomException(Exception):
pass
def func(x):
raise CustomException
queue_handler = MagicMock()
start_handler = MagicMock()
failure_handler = MagicMock()
scheduler = ThreadPoolScheduler()
task = Task(
func,
on_start_handlers=[start_handler],
on_queue_handlers=[queue_handler],
on_failure_handlers=[failure_handler],
)
scheduler.add_task(task=task, params=param)
assert queue_handler.called
assert task.state == State.QUEUED
assert not failure_handler.called
assert not start_handler.called
result = scheduler.results()
assert next(result) is None
assert start_handler.called
assert failure_handler.called
handler_args = failure_handler.call_args_list[0][0]
assert isinstance(handler_args[0], CustomException)
assert task.state == State.FAILED