Merge pull request #404 from dimastbk/fix-fetch-variables

fix: fetch all variables, if job is present

pyzeebe/function_tools/parameter_tools.py

@@ -5,12 +5,19 @@
 from pyzeebe.job.job import Job
 
 
-def get_parameters_from_function(task_function: Function) -> List[str]:
+def get_parameters_from_function(task_function: Function) -> Optional[List[str]]:
     function_signature = inspect.signature(task_function)
     for _, parameter in function_signature.parameters.items():
         if parameter.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
             return []
-    return list(function_signature.parameters)
+
+    if not function_signature.parameters:
+        return None
+
+    if all(param.annotation == Job for param in function_signature.parameters.values()):
+        return []
+
+    return [param.name for param in function_signature.parameters.values() if param.annotation != Job]
 
 
 def get_job_parameter_name(function: Function) -> Optional[str]:

pyzeebe/task/task_builder.py

@@ -58,7 +58,15 @@ async def run_original_task_function(
     task_function: DictFunction, task_config: TaskConfig, job: Job
 ) -> Tuple[Dict, bool]:
     try:
-        returned_value = await task_function(**job.variables)  # type: ignore
+        if task_config.variables_to_fetch is None:
+            variables = {}
+        else:
+            variables = {
+                k: v
+                for k, v in job.variables.items()
+                if k in task_config.variables_to_fetch or k == task_config.job_parameter_name
+            }
+        returned_value = await task_function(**variables)  # type: ignore
 
         if returned_value is None:
             returned_value = {}

pyzeebe/task/task_config.py

@@ -17,7 +17,7 @@ def __init__(
         timeout_ms: int,
         max_jobs_to_activate: int,
         max_running_jobs: int,
-        variables_to_fetch: List[str],
+        variables_to_fetch: Optional[List[str]],
         single_value: bool,
         variable_name: str,
         before: List[TaskDecorator],

pyzeebe/worker/job_poller.py

@@ -59,7 +59,7 @@ async def poll_once(self):
                 worker=self.worker_name,
                 timeout=self.task.config.timeout_ms,
                 max_jobs_to_activate=self.calculate_max_jobs_to_activate(),
-                variables_to_fetch=self.task.config.variables_to_fetch,
+                variables_to_fetch=self.task.config.variables_to_fetch or [],
                 request_timeout=self.request_timeout,
                 tenant_ids=self.tenant_ids,
             )

tests/unit/function_tools/parameter_tools_test.py

@@ -1,4 +1,4 @@
-from typing import Callable, List
+from typing import Callable, List, Optional
 
 import pytest
 
@@ -10,7 +10,7 @@ class TestGetFunctionParameters:
     @pytest.mark.parametrize(
         "fn,expected",
         [
-            (dummy_functions.no_param, []),
+            (dummy_functions.no_param, None),
             (dummy_functions.one_param, ["x"]),
             (dummy_functions.multiple_params, ["x", "y", "z"]),
             (dummy_functions.one_keyword_param, ["x"]),
@@ -19,15 +19,18 @@ class TestGetFunctionParameters:
             (dummy_functions.args_param, []),
             (dummy_functions.kwargs_param, []),
             (dummy_functions.standard_named_params, ["args", "kwargs"]),
-            (dummy_functions.lambda_no_params, []),
+            (dummy_functions.with_job_parameter, []),
+            (dummy_functions.with_job_parameter_and_param, ["x"]),
+            (dummy_functions.with_multiple_job_parameters, []),
+            (dummy_functions.lambda_no_params, None),
             (dummy_functions.lambda_one_param, ["x"]),
             (dummy_functions.lambda_multiple_params, ["x", "y", "z"]),
             (dummy_functions.lambda_one_keyword_param, ["x"]),
             (dummy_functions.lambda_multiple_keyword_params, ["x", "y", "z"]),
             (dummy_functions.lambda_positional_and_keyword_params, ["x", "y"]),
         ],
     )
-    def test_get_params(self, fn: Callable, expected: List[str]):
+    def test_get_params(self, fn: Callable, expected: Optional[List[str]]):
         assert parameter_tools.get_parameters_from_function(fn) == expected
 
 

tests/unit/task/task_builder_test.py

@@ -35,6 +35,7 @@ async def test_no_additional_variables_are_added_to_result(
         self, single_value_task_config: TaskConfig, mocked_job_with_adapter: Job
     ):
         mocked_job_with_adapter.variables = {"x": 1}
+        single_value_task_config.variables_to_fetch = ["x"]
 
         task = task_builder.build_task(lambda x: x, single_value_task_config)
         job = await task.job_handler(mocked_job_with_adapter)
@@ -87,12 +88,42 @@ async def test_parameters_are_provided_to_task(
         self, original_task_function: Callable, task_config: TaskConfig, mocked_job_with_adapter: Job
     ):
         mocked_job_with_adapter.variables = {"x": 1}
+        task_config.variables_to_fetch = ["x"]
+
         job_handler = task_builder.build_job_handler(original_task_function, task_config)
 
         await job_handler(mocked_job_with_adapter)
 
         original_task_function.assert_called_with(x=1)
 
+    @pytest.mark.asyncio
+    async def test_parameters_are_provided_to_task_with_only_job(
+        self, original_task_function: Callable, task_config: TaskConfig, mocked_job_with_adapter: Job
+    ):
+        mocked_job_with_adapter.variables = {"x": 1}
+        task_config.job_parameter_name = "job"
+        task_config.variables_to_fetch = []
+
+        job_handler = task_builder.build_job_handler(original_task_function, task_config)
+
+        await job_handler(mocked_job_with_adapter)
+
+        original_task_function.assert_called_with(job=mocked_job_with_adapter)
+
+    @pytest.mark.asyncio
+    async def test_parameters_are_provided_to_task_with_arg_and_job(
+        self, original_task_function: Callable, task_config: TaskConfig, mocked_job_with_adapter: Job
+    ):
+        mocked_job_with_adapter.variables = {"x": 1}
+        task_config.job_parameter_name = "job"
+        task_config.variables_to_fetch = ["x"]
+
+        job_handler = task_builder.build_job_handler(original_task_function, task_config)
+
+        await job_handler(mocked_job_with_adapter)
+
+        original_task_function.assert_called_with(job=mocked_job_with_adapter, x=1)
+
     @pytest.mark.asyncio
     async def test_variables_are_added_to_result(
         self, original_task_function: Callable, task_config: TaskConfig, mocked_job_with_adapter: Job

tests/unit/utils/dummy_functions.py

@@ -41,29 +41,17 @@ def with_job_parameter(job: Job):
     pass
 
 
-def with_multiple_job_parameters(job: Job, job2: Job):
+def with_job_parameter_and_param(x, job: Job):
     pass
 
 
-def lambda_no_params():
-    return None
-
-
-def lambda_one_param(x):
-    return None
-
-
-def lambda_multiple_params(x, y, z):
-    return None
-
-
-def lambda_one_keyword_param(x=0):
-    return None
-
-
-def lambda_multiple_keyword_params(x=0, y=0, z=0):
-    return None
+def with_multiple_job_parameters(job: Job, job2: Job):
+    pass
 
 
-def lambda_positional_and_keyword_params(x, y=0):
-    return None
+lambda_no_params = lambda: None
+lambda_one_param = lambda x: None
+lambda_multiple_params = lambda x, y, z: None
+lambda_one_keyword_param = lambda x=0: None
+lambda_multiple_keyword_params = lambda x=0, y=0, z=0: None
+lambda_positional_and_keyword_params = lambda x, y=0: None