test(task-sdk): extend test_run_with_asset_outlets with asset subclasses

apache · Mar 11, 2025 · c298093 · c298093
1 parent b309297
commit c298093
Showing 1 changed file with 40 additions and 1 deletion.
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -52,7 +52,9 @@
     TaskInstance,
     TerminalTIState,
 )
-from airflow.sdk.definitions.asset import Asset, AssetAlias
+from airflow.sdk.definitions.asset import Asset, AssetAlias, Dataset, Model
+
+# from airflow.sdk.definitions.asset.decorators import AssetDefinition
 from airflow.sdk.definitions.param import DagParam
 from airflow.sdk.definitions.variable import Variable
 from airflow.sdk.execution_time.comms import (
@@ -703,6 +705,43 @@ def test_dag_parsing_context(make_ti_context, mock_supervisor_comms, monkeypatch
             ),
             id="asset",
         ),
+        pytest.param(
+            [Dataset(name="s3://bucket/my-task", uri="s3://bucket/my-task")],
+            SucceedTask(
+                state="success",
+                end_date=timezone.datetime(2024, 12, 3, 10, 0),
+                task_outlets=[
+                    AssetProfile(name="s3://bucket/my-task", uri="s3://bucket/my-task", asset_type="Asset")
+                ],
+                outlet_events=[
+                    {
+                        "key": {"name": "s3://bucket/my-task", "uri": "s3://bucket/my-task"},
+                        "extra": {},
+                        "asset_alias_events": [],
+                    }
+                ],
+            ),
+            id="dataset",
+        ),
+        pytest.param(
+            [Model(name="s3://bucket/my-task", uri="s3://bucket/my-task")],
+            SucceedTask(
+                state="success",
+                end_date=timezone.datetime(2024, 12, 3, 10, 0),
+                task_outlets=[
+                    AssetProfile(name="s3://bucket/my-task", uri="s3://bucket/my-task", asset_type="Asset")
+                ],
+                outlet_events=[
+                    {
+                        "key": {"name": "s3://bucket/my-task", "uri": "s3://bucket/my-task"},
+                        "extra": {},
+                        "asset_alias_events": [],
+                    }
+                ],
+            ),
+            id="model",
+        ),
+        # TODO: asset name ref and asset uri ref are not checked and should be added later
         pytest.param(
             [AssetAlias(name="example-alias", group="asset")],
             SucceedTask(