test(dataset): add test case test_dataset_iter_dataset_aliases

apache · Aug 6, 2024 · f3375a2 · f3375a2
1 parent 31be76a
commit f3375a2
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/tests/datasets/test_dataset.py b/tests/datasets/test_dataset.py
@@ -28,6 +28,7 @@
 from airflow.datasets import (
     BaseDataset,
     Dataset,
+    DatasetAlias,
     DatasetAll,
     DatasetAny,
     _DatasetAliasCondition,
@@ -136,6 +137,24 @@ def test_dataset_iter_datasets():
     assert list(dataset1.iter_datasets()) == [("s3://bucket1/data1", dataset1)]
 
 
+@pytest.mark.db_test
+def test_dataset_iter_dataset_aliases():
+    base_dataset = DatasetAll(
+        DatasetAlias("example-alias-1"),
+        Dataset("1"),
+        DatasetAny(
+            Dataset("2"),
+            DatasetAlias("example-alias-2"),
+            Dataset("3"),
+            DatasetAll(DatasetAlias("example-alias-3"), Dataset("4"), DatasetAlias("example-alias-4")),
+        ),
+        DatasetAll(DatasetAlias("example-alias-5"), Dataset("5")),
+    )
+    assert list(base_dataset.iter_dataset_aliases()) == [
+        DatasetAlias(f"example-alias-{i}") for i in range(1, 6)
+    ]
+
+
 def test_dataset_evaluate():
     assert dataset1.evaluate({"s3://bucket1/data1": True}) is True
     assert dataset1.evaluate({"s3://bucket1/data1": False}) is False