diff --git a/pyzeebe/function_tools/parameter_tools.py b/pyzeebe/function_tools/parameter_tools.py index 09349aac..da61fcbd 100644 --- a/pyzeebe/function_tools/parameter_tools.py +++ b/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]: diff --git a/pyzeebe/task/task_builder.py b/pyzeebe/task/task_builder.py index c684ff3d..dceef6ed 100644 --- a/pyzeebe/task/task_builder.py +++ b/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 = {} diff --git a/pyzeebe/task/task_config.py b/pyzeebe/task/task_config.py index 1c8c4b43..2748a384 100644 --- a/pyzeebe/task/task_config.py +++ b/pyzeebe/task/task_config.py @@ -1,7 +1,7 @@ -from typing import List, Optional +from typing import List, Optional, Type, Union from pyzeebe.errors import NoVariableNameGivenError -from pyzeebe.function_tools import async_tools +from pyzeebe.function_tools import async_tools, parameter_tools from pyzeebe.task.exception_handler import ExceptionHandler from pyzeebe.task.types import AsyncTaskDecorator, TaskDecorator @@ -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], diff --git a/pyzeebe/worker/job_poller.py b/pyzeebe/worker/job_poller.py index fbb634b3..1bba63b1 100644 --- a/pyzeebe/worker/job_poller.py +++ b/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, ) diff --git a/tests/unit/function_tools/parameter_tools_test.py b/tests/unit/function_tools/parameter_tools_test.py index 08ddbe04..16fae394 100644 --- a/tests/unit/function_tools/parameter_tools_test.py +++ b/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,7 +19,10 @@ 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"]), @@ -27,7 +30,7 @@ class TestGetFunctionParameters: (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 diff --git a/tests/unit/task/task_builder_test.py b/tests/unit/task/task_builder_test.py index 8345ab60..754c14ea 100644 --- a/tests/unit/task/task_builder_test.py +++ b/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 diff --git a/tests/unit/utils/dummy_functions.py b/tests/unit/utils/dummy_functions.py index 56240dfa..3958e8f4 100644 --- a/tests/unit/utils/dummy_functions.py +++ b/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