From 0ef189e477c0832a90a3f87a72cd533017abc5f5 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 1 Nov 2024 18:06:33 +0800 Subject: [PATCH] test(decorators/assets): extend test_determine_kwargs to cover active asset --- tests/decorators/test_assets.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/decorators/test_assets.py b/tests/decorators/test_assets.py index c5d4edadd024f..9105a5cec520b 100644 --- a/tests/decorators/test_assets.py +++ b/tests/decorators/test_assets.py @@ -23,6 +23,7 @@ from airflow.assets import Asset from airflow.decorators.assets import AssetRef, _AssetMainOperator, asset +from airflow.models.asset import AssetActive, AssetModel pytestmark = pytest.mark.db_test @@ -145,14 +146,19 @@ def test__attrs_post_init__( class Test_AssetMainOperator: - def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset): + def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset, session): + example_asset = AssetModel(uri="s3://bucket/object1", name="inlet_asset_1") + session.add(example_asset) + session.add(AssetActive.for_asset(example_asset)) + session.commit() + asset_definition = asset(schedule=None, uri="s3://bucket/object", group="MLModel", extra={"k": "v"})( example_asset_func_with_valid_arg_as_inlet_asset ) op = _AssetMainOperator( task_id="__main__", - inlets=[], + inlets=[AssetRef(name="inlet_asset_1"), AssetRef(name="inlet_asset_2")], outlets=[asset_definition], python_callable=example_asset_func_with_valid_arg_as_inlet_asset, definition_name="example_asset_func", @@ -160,6 +166,6 @@ def test_determine_kwargs(self, example_asset_func_with_valid_arg_as_inlet_asset assert op.determine_kwargs(context={"k": "v"}) == { "_self": Asset(name="example_asset_func"), "context": {"k": "v"}, - "inlet_asset_1": Asset(name="inlet_asset_1"), + "inlet_asset_1": Asset(name="inlet_asset_1", uri="s3://bucket/object1"), "inlet_asset_2": Asset(name="inlet_asset_2"), }