From 2d44d5ae39cd08ca5da4a25769257eeee19b7346 Mon Sep 17 00:00:00 2001 From: briantu Date: Thu, 3 Oct 2024 14:44:20 -0400 Subject: [PATCH] Fix tests --- .../dagster/dagster/_core/definitions/base_asset_graph.py | 2 +- .../automation_condition_evaluator.py | 2 +- .../dagster_tests/asset_defs_tests/test_asset_graph.py | 6 +++++- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py b/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py index 1195628c73813..37ccca15b1fab 100644 --- a/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py +++ b/python_modules/dagster/dagster/_core/definitions/base_asset_graph.py @@ -301,7 +301,7 @@ def toposorted_asset_keys(self) -> Sequence[AssetKey]: ] @cached_property - def toposorted_entity_keys(self) -> Sequence[Sequence[EntityKey]]: + def toposorted_entity_keys_by_level(self) -> Sequence[Sequence[EntityKey]]: """Return topologically sorted levels for entity keys in graph. Keys with the same topological level are sorted alphabetically to provide stability. """ diff --git a/python_modules/dagster/dagster/_core/definitions/declarative_automation/automation_condition_evaluator.py b/python_modules/dagster/dagster/_core/definitions/declarative_automation/automation_condition_evaluator.py index 00fb4e99f7ceb..624536186b480 100644 --- a/python_modules/dagster/dagster/_core/definitions/declarative_automation/automation_condition_evaluator.py +++ b/python_modules/dagster/dagster/_core/definitions/declarative_automation/automation_condition_evaluator.py @@ -152,7 +152,7 @@ async def _evaluate_entity_async(entity_key: EntityKey) -> None: ) num_evaluated += 1 - for topo_level in self.asset_graph.toposorted_entity_keys: + for topo_level in self.asset_graph.toposorted_entity_keys_by_level: coroutines = [ _evaluate_entity_async(entity_key) for entity_key in topo_level diff --git a/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_graph.py b/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_graph.py index f94f92bdd8f7a..0d8ffe5b6b5f9 100644 --- a/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_graph.py +++ b/python_modules/dagster/dagster_tests/asset_defs_tests/test_asset_graph.py @@ -753,7 +753,11 @@ def Bc(): ... asset_graph = asset_graph_from_assets([A, B, Ac, Bc]) assert asset_graph.toposorted_asset_keys == [A.key, B.key] - assert asset_graph.toposorted_entity_keys == [A.key, Ac.check_key, B.key, Bc.check_key] + assert asset_graph.toposorted_entity_keys_by_level == [ + [A.key], + [Ac.check_key, B.key], + [Bc.check_key], + ] def test_required_assets_and_checks_by_key_asset_decorator(