Refactor/task builder

pyzeebe/__init__.py

@@ -7,7 +7,7 @@
 from pyzeebe.job.job import Job
 from pyzeebe.job.job_status import JobStatus
 from pyzeebe.task.exception_handler import ExceptionHandler
-from pyzeebe.task.task_config import TaskConfig, default_exception_handler
+from pyzeebe.task.task_config import TaskConfig
 from pyzeebe.task.types import TaskDecorator
-from pyzeebe.worker.task_router import ZeebeTaskRouter
+from pyzeebe.worker.task_router import ZeebeTaskRouter, default_exception_handler
 from pyzeebe.worker.worker import ZeebeWorker

pyzeebe/task/task_config.py

@@ -1,37 +1,24 @@
-import logging
-from typing import List, Optional
+from typing import List
 
-from pyzeebe.errors import NoVariableNameGivenError
-from pyzeebe.job.job import Job
 from pyzeebe.task.exception_handler import ExceptionHandler
 from pyzeebe.task.types import TaskDecorator
 
-logger = logging.getLogger(__name__)
-
-
-def default_exception_handler(e: Exception, job: Job) -> None:
-    logger.warning(f"Task type: {job.type} - failed job {job}. Error: {e}.")
-    job.set_failure_status(f"Failed job. Error: {e}")
-
 
 class TaskConfig:
-    def __init__(self, type: str, exception_handler: ExceptionHandler = default_exception_handler,
-                 timeout_ms: int = 10000, max_jobs_to_activate: int = 32,
-                 variables_to_fetch: Optional[List[str]] = None,
-                 single_value: bool = False, variable_name: Optional[str] = None, before: List[TaskDecorator] = None,
-                 after: List[TaskDecorator] = None):
-        if single_value and not variable_name:
-            raise NoVariableNameGivenError(type)
-
+    def __init__(self, type: str, exception_handler: ExceptionHandler,
+                 timeout_ms: int, max_jobs_to_activate: int,
+                 variables_to_fetch: List[str],
+                 single_value: bool, variable_name: str, before: List[TaskDecorator],
+                 after: List[TaskDecorator]):
         self.type = type
         self.exception_handler = exception_handler
         self.timeout_ms = timeout_ms
         self.max_jobs_to_activate = max_jobs_to_activate
         self.variables_to_fetch = variables_to_fetch
         self.single_value = single_value
         self.variable_name = variable_name
-        self.before = before or []
-        self.after = after or []
+        self.before = before
+        self.after = after
 
     def __repr__(self):
         return f"TaskConfig(type={self.type}, exception_handler={self.exception_handler}, " \

pyzeebe/worker/task_router.py

@@ -1,15 +1,23 @@
 import logging
-from typing import Tuple, List, Callable, Union
+from typing import Callable, List, Tuple, Optional
 
 from pyzeebe import TaskDecorator
-from pyzeebe.errors import TaskNotFoundError, DuplicateTaskTypeError
+from pyzeebe.errors import (DuplicateTaskTypeError, NoVariableNameGivenError,
+                            TaskNotFoundError)
+from pyzeebe.job.job import Job
 from pyzeebe.task import task_builder
+from pyzeebe.task.exception_handler import ExceptionHandler
 from pyzeebe.task.task import Task
 from pyzeebe.task.task_config import TaskConfig
 
 logger = logging.getLogger(__name__)
 
 
+def default_exception_handler(e: Exception, job: Job) -> None:
+    logger.warning(f"Task type: {job.type} - failed job {job}. Error: {e}.")
+    job.set_failure_status(f"Failed job. Error: {e}")
+
+
 class ZeebeTaskRouter:
     def __init__(self, before: List[TaskDecorator] = None, after: List[TaskDecorator] = None):
         """
@@ -21,25 +29,60 @@ def __init__(self, before: List[TaskDecorator] = None, after: List[TaskDecorator
         self._after: List[TaskDecorator] = after or []
         self.tasks: List[Task] = []
 
-    def task(self, task_config: Union[TaskConfig, str]):
+    def task(self, task_type: str, exception_handler: ExceptionHandler = default_exception_handler,
+             variables_to_fetch: Optional[List[str]] = None, timeout_ms: int = 10000, max_jobs_to_activate: int = 32,
+             before: List[TaskDecorator] = None, after: List[TaskDecorator] = None, single_value: bool = False,
+             variable_name: str = None):
         """
         Decorator to create a task
-
         Args:
-            task_config (Union[str, TaskConfig]): Either the task type or a task configuration object
+            task_type (str): The task type
+
+            exception_handler (ExceptionHandler): Handler that will be called when a job fails. 
+
+            variables_to_fetch (Optional[List[str]]): The variables to request from Zeebe when activating jobs.
+
+            timeout_ms (int): Maximum duration of the task in milliseconds. If the timeout is surpassed Zeebe will give up 
+                                on the worker and retry it. Default: 10000 (10 seconds).
 
+            max_jobs_to_activate (int):  Maximum jobs the worker will execute in parallel (of this task). Default: 32
+
+            before (List[TaskDecorator]): All decorators which should be performed before the task.
+
+            after (List[TaskDecorator]): All decorators which should be performed after the task.
+
+            single_value (bool): If the function returns a single value (int, string, list) and not a dictionary set
+                                 this to True. Default: False
+
+            variable_name (str): If single_value then this will be the variable name given to zeebe:
+                                        { <variable_name>: <function_return_value> }
         Raises:
             DuplicateTaskTypeError: If a task from the router already exists in the worker
-
+            NoVariableNameGivenError: When single_value is set, but no variable_name is given
         """
-        if isinstance(task_config, str):
-            task_config = TaskConfig(task_config)
-        config_with_decorators = self._add_decorators_to_config(task_config)
-
-        def task_wrapper(fn: Callable):
-            task = task_builder.build_task(fn, config_with_decorators)
+        if single_value and not variable_name:
+            raise NoVariableNameGivenError(task_type)
+
+        def task_wrapper(task_function: Callable):
+            config = TaskConfig(
+                task_type,
+                exception_handler,
+                timeout_ms,
+                max_jobs_to_activate,
+                variables_to_fetch or task_builder.get_parameters_from_function(
+                    task_function),
+                single_value,
+                variable_name or "",
+                before or [],
+                after or []
+            )
+            config_with_decorators = self._add_decorators_to_config(config)
+
+            task = task_builder.build_task(
+                task_function, config_with_decorators
+            )
             self._add_task(task)
-            return fn
+            return task_function
 
         return task_wrapper
 

tests/conftest.py

@@ -69,7 +69,17 @@ def first_active_job(task, job_from_task, grpc_servicer) -> str:
 
 @pytest.fixture
 def task_config(task_type):
-    return TaskConfig(task_type, MagicMock())
+    return TaskConfig(
+        type=task_type,
+        exception_handler=MagicMock(),
+        timeout_ms=10000,
+        max_jobs_to_activate=32,
+        variables_to_fetch=[],
+        single_value=False,
+        variable_name="",
+        before=[],
+        after=[]
+    )
 
 
 @pytest.fixture

tests/unit/utils/random_utils.py

@@ -10,7 +10,7 @@
 RANDOM_RANGE = 1000000000
 
 
-def random_job(task: Task = task_builder.build_task(lambda x: {"x": x}, TaskConfig("test")),
+def random_job(task: Task = task_builder.build_task(lambda x: {"x": x}, TaskConfig("test", lambda: None, 10000, 32, [], False, "", [], [])),
                zeebe_adapter: ZeebeAdapter = None) -> Job:
     return Job(_type=task.type, key=randint(0, RANDOM_RANGE), worker=str(uuid4()),
                retries=randint(0, 10), workflow_instance_key=randint(0, RANDOM_RANGE),

tests/unit/worker/task_router_test.py

@@ -1,11 +1,13 @@
+from unittest.mock import patch
 from uuid import uuid4
 
 import pytest
-
 from pyzeebe import TaskDecorator
-from pyzeebe.errors import TaskNotFoundError, DuplicateTaskTypeError
+from pyzeebe.errors import DuplicateTaskTypeError, TaskNotFoundError
+from pyzeebe.job.job import Job
 from pyzeebe.task.task import Task
-from pyzeebe.worker.task_router import ZeebeTaskRouter
+from pyzeebe.worker.task_router import (ZeebeTaskRouter,
+                                        default_exception_handler)
 from tests.unit.utils.random_utils import randint
 
 
@@ -97,3 +99,11 @@ def test_add_after_decorator_through_constructor(decorator: TaskDecorator):
     router = ZeebeTaskRouter(after=[decorator])
 
     assert len(router._after) == 1
+
+
+def test_default_exception_handler_logs_a_warning(mocked_job_with_adapter: Job):
+    with patch("pyzeebe.worker.task_router.logger.warning") as logging_mock:
+        default_exception_handler(Exception(), mocked_job_with_adapter)
+
+        mocked_job_with_adapter.set_failure_status.assert_called()
+        logging_mock.assert_called()

tests/unit/worker/worker_test.py

@@ -43,7 +43,8 @@ def test_variables_to_fetch_match_function_parameters(self, zeebe_worker: ZeebeW
         def dummy_function(x):
             pass
 
-        assert zeebe_worker.get_task(task_type).config.variables_to_fetch == expected_variables_to_fetch
+        assert zeebe_worker.get_task(
+            task_type).config.variables_to_fetch == expected_variables_to_fetch
 
 
 class TestDecorator:
@@ -67,21 +68,6 @@ def test_add_constructor_after_decorator(self, decorator: TaskDecorator):
         assert len(zeebe_worker._after) == 1
         assert decorator in zeebe_worker._after
 
-    def test_decorator_failed(self, zeebe_worker: ZeebeWorker, task: Task, decorator: TaskDecorator,
-                              job_from_task: Job):
-        decorator.side_effect = Exception()
-        zeebe_worker.before(decorator)
-        zeebe_worker.after(decorator)
-
-        @zeebe_worker.task(task.config)
-        def dummy_function():
-            pass
-
-        task = zeebe_worker.get_task(task.type)
-        task.job_handler(job_from_task)
-
-        assert decorator.call_count == 2
-
 
 class TestHandleJobs:
     @pytest.fixture(autouse=True)