Skip to content

Commit

Permalink
fix: fetch all variables, if job is present
Browse files Browse the repository at this point in the history
  • Loading branch information
dimastbk committed May 24, 2024
1 parent 796c5a5 commit 28a01c4
Show file tree
Hide file tree
Showing 7 changed files with 69 additions and 32 deletions.
11 changes: 9 additions & 2 deletions pyzeebe/function_tools/parameter_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down
10 changes: 9 additions & 1 deletion pyzeebe/task/task_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {}
Expand Down
6 changes: 3 additions & 3 deletions pyzeebe/task/task_config.py
Original file line number Diff line number Diff line change
@@ -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

Expand All @@ -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],
Expand Down
2 changes: 1 addition & 1 deletion pyzeebe/worker/job_poller.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
11 changes: 7 additions & 4 deletions tests/unit/function_tools/parameter_tools_test.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Callable, List
from typing import Callable, List, Optional

import pytest

Expand All @@ -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"]),
Expand All @@ -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


Expand Down
31 changes: 31 additions & 0 deletions tests/unit/task/task_builder_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
30 changes: 9 additions & 21 deletions tests/unit/utils/dummy_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 28a01c4

Please sign in to comment.