diff --git a/python_modules/dagster/dagster_tests/asset_defs_tests/test_decorators.py b/python_modules/dagster/dagster_tests/asset_defs_tests/test_decorators.py index 9f1e782a6cd71..f004971c0d724 100644 --- a/python_modules/dagster/dagster_tests/asset_defs_tests/test_decorators.py +++ b/python_modules/dagster/dagster_tests/asset_defs_tests/test_decorators.py @@ -23,6 +23,7 @@ TimeWindowPartitionMapping, _check as check, build_asset_context, + build_op_context, graph_asset, graph_multi_asset, io_manager, @@ -610,7 +611,9 @@ def my_asset(context, **kwargs): assert len(my_asset.op.input_defs) == 1 assert AssetKey("upstream") in my_asset.keys_by_input_name.values() assert my_asset(build_asset_context(), upstream=5) == 7 - assert my_asset.op(build_asset_context(), upstream=5) == 7 + assert ( + my_asset.op(build_op_context(), upstream=5) == 7 + ) # TODO - this test is odd now since my_asset should expect an AssetContext @asset def upstream(): ... @@ -649,7 +652,7 @@ def my_asset(context, **kwargs): assert len(my_asset.op.input_defs) == 1 assert AssetKey("upstream") in my_asset.keys_by_input_name.values() assert my_asset(build_asset_context(), upstream=5) == (7,) - assert my_asset.op(build_asset_context(), upstream=5) == (7,) + assert my_asset.op(build_op_context(), upstream=5) == (7,) @asset def upstream(): ...