Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 4 additions & 3 deletions task-sdk/src/airflow/sdk/definitions/asset/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,14 +87,15 @@ def _fetch_asset(name: str) -> Asset:

value: Any
for key, param in inspect.signature(self.python_callable).parameters.items():
if param.default is not inspect.Parameter.empty:
value = param.default
elif key == "self":
# ``self``/``context``/``outlet_events`` are reserved keys.
if key == "self":
value = _fetch_asset(self._definition_name)
elif key == "context":
value = context
elif key == "outlet_events":
value = context["outlet_events"]
elif param.default is not inspect.Parameter.empty:
value = param.default
else:
value = _fetch_asset(key)
yield key, value
Expand Down
38 changes: 38 additions & 0 deletions task-sdk/tests/task_sdk/definitions/test_asset_decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -449,6 +449,44 @@ def test_determine_kwargs_defaults(
mock.call.send(GetAssetByName(name="inlet_asset_1")),
]

@pytest.mark.parametrize(
("python_callable", "reserved_key"),
[
pytest.param(lambda self=None: None, "self", id="self"),
pytest.param(lambda context=None: None, "context", id="context"),
pytest.param(lambda outlet_events=None: None, "outlet_events", id="outlet_events"),
],
)
def test_determine_kwargs_injects_reserved_key_with_default(
self, mock_supervisor_comms, func_fixer, python_callable, reserved_key
):
"""A default on a reserved arg (self/context/outlet_events) must not shadow the injected value."""
mock_supervisor_comms.send.side_effect = [
AssetResult(name="example_asset_func", uri="example_asset_func", group="asset"),
]
fixed_callable = func_fixer(python_callable)
definition = asset(schedule=None)(fixed_callable)
outlet_events = OutletEventAccessors()
context = {"outlet_events": outlet_events}

op = _AssetMainOperator(
task_id="example_asset_func",
inlets=[],
outlets=[definition],
python_callable=fixed_callable,
definition_name="example_asset_func",
)

injected = op.determine_kwargs(context=context)[reserved_key]

assert injected is not None
if reserved_key == "self":
assert isinstance(injected, Asset)
elif reserved_key == "context":
assert injected is context
else:
assert injected is outlet_events

def test_from_definition_custom_name(self, mock_supervisor_comms, func_fixer):
@func_fixer
def example_asset_func(self):
Expand Down