diff --git a/src/hera/workflows/_meta_mixins.py b/src/hera/workflows/_meta_mixins.py index 1bca8002e..c4c096170 100644 --- a/src/hera/workflows/_meta_mixins.py +++ b/src/hera/workflows/_meta_mixins.py @@ -673,7 +673,8 @@ def script_call_wrapper(*args, **kwargs) -> Union[FuncR, Step, Task, None]: if _context.pieces: return script_template.__call__(*args, **kwargs) - return func(*args, **kwargs) + + return func(*args) # Set the wrapped function to the original function so that we can use it later script_call_wrapper.wrapped_function = func # type: ignore @@ -732,7 +733,7 @@ def container_call_wrapper(*args, **kwargs) -> Union[FuncR, Step, Task, None]: if _context.pieces: return container_template.__call__(*args, **kwargs) - return func(*args, **kwargs) + return func(*args) # Set the template name to the inferred name container_call_wrapper.template_name = name # type: ignore @@ -817,7 +818,7 @@ def call_wrapper(*args, **kwargs): return self._create_subnode(subnode_name, func, template, *args, **kwargs) - return func(*args, **kwargs) + return func(*args) call_wrapper.template_name = name # type: ignore diff --git a/tests/test_unit/test_decorators.py b/tests/test_unit/test_decorators.py index 42c76bfd5..c1ef7bcea 100644 --- a/tests/test_unit/test_decorators.py +++ b/tests/test_unit/test_decorators.py @@ -170,7 +170,7 @@ def test_dag_task_auto_depends(): task_b = next(iter([t for t in dag_template.tasks if t.name == "task-b"]), None) assert task_b.depends == "setup-task" - final_task = next(iter([t for t in dag_template.tasks if t.name == "final-task"]), None) + final_task = next(iter([t for t in dag_template.tasks if t.name == "final-task-name"]), None) assert final_task.depends == "task-a && task-b" diff --git a/tests/workflow_decorators/dag.py b/tests/workflow_decorators/dag.py index 0e395ba83..feba26b45 100644 --- a/tests/workflow_decorators/dag.py +++ b/tests/workflow_decorators/dag.py @@ -43,6 +43,6 @@ def worker(worker_input: WorkerInput) -> WorkerOutput: setup_task = setup() task_a = concat(ConcatInput(word_a=worker_input.value_a, word_b=setup_task.environment_parameter)) task_b = concat(ConcatInput(word_a=worker_input.value_b, word_b=setup_task.result)) - final_task = concat(ConcatInput(word_a=task_a.result, word_b=task_b.result)) + final_task = concat(ConcatInput(word_a=task_a.result, word_b=task_b.result), name="final-task-name") return WorkerOutput(value=final_task.result)