From 774532d45b0891c63d351f26a4d47bb71e6b3355 Mon Sep 17 00:00:00 2001 From: JamieDeMaria Date: Fri, 12 Jan 2024 15:46:46 -0500 Subject: [PATCH] update a test suite --- .../test_asset_partition_mappings.py | 22 ++++++------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/python_modules/dagster/dagster_tests/asset_defs_tests/partition_mapping_tests/test_asset_partition_mappings.py b/python_modules/dagster/dagster_tests/asset_defs_tests/partition_mapping_tests/test_asset_partition_mappings.py index 7c1a07f8916a7..6da99ce78abdd 100644 --- a/python_modules/dagster/dagster_tests/asset_defs_tests/partition_mapping_tests/test_asset_partition_mappings.py +++ b/python_modules/dagster/dagster_tests/asset_defs_tests/partition_mapping_tests/test_asset_partition_mappings.py @@ -600,9 +600,7 @@ def upstream(): ], ) def downstream(context: AssetExecutionContext): - upstream_key = datetime.strptime( - context.asset_partition_key_for_input("upstream"), "%Y-%m-%d" - ) + upstream_key = datetime.strptime(context.upstream_partition_key("upstream"), "%Y-%m-%d") current_partition_key = datetime.strptime(context.partition_key, "%Y-%m-%d") @@ -653,12 +651,8 @@ def multi_asset_1(): @multi_asset(specs=[asset_3, asset_4], partitions_def=partitions_def) def multi_asset_2(context: AssetExecutionContext): - asset_1_key = datetime.strptime( - context.asset_partition_key_for_input("asset_1"), "%Y-%m-%d" - ) - asset_2_key = datetime.strptime( - context.asset_partition_key_for_input("asset_2"), "%Y-%m-%d" - ) + asset_1_key = datetime.strptime(context.upstream_partition_key("asset_1"), "%Y-%m-%d") + asset_2_key = datetime.strptime(context.upstream_partition_key("asset_2"), "%Y-%m-%d") current_partition_key = datetime.strptime(context.partition_key, "%Y-%m-%d") @@ -760,7 +754,7 @@ def test_self_dependent_partition_mapping_with_asset_deps(): ) def self_dependent(context: AssetExecutionContext): upstream_key = datetime.strptime( - context.asset_partition_key_for_input("self_dependent"), "%Y-%m-%d" + context.upstream_partition_key("self_dependent"), "%Y-%m-%d" ) current_partition_key = datetime.strptime(context.partition_key, "%Y-%m-%d") @@ -786,9 +780,7 @@ def self_dependent(context: AssetExecutionContext): @multi_asset(specs=[asset_1], partitions_def=partitions_def) def the_multi_asset(context: AssetExecutionContext): - asset_1_key = datetime.strptime( - context.asset_partition_key_for_input("asset_1"), "%Y-%m-%d" - ) + asset_1_key = datetime.strptime(context.upstream_partition_key("asset_1"), "%Y-%m-%d") current_partition_key = datetime.strptime(context.partition_key, "%Y-%m-%d") @@ -810,7 +802,7 @@ def upstream(): deps=[AssetDep(upstream, partition_mapping=SpecificPartitionsPartitionMapping(["apple"]))], ) def downstream(context: AssetExecutionContext): - assert context.asset_partition_key_for_input("upstream") == "apple" + assert context.upstream_partition_key("upstream") == "apple" assert context.partition_key == "orange" with instance_for_test() as instance: @@ -840,7 +832,7 @@ def asset_1_multi_asset(): @multi_asset(specs=[asset_2], partitions_def=partitions_def) def asset_2_multi_asset(context: AssetExecutionContext): - assert context.asset_partition_key_for_input("asset_1") == "apple" + assert context.upstream_partition_key("asset_1") == "apple" assert context.partition_key == "orange" with instance_for_test() as instance: