diff --git a/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py b/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py index 1ebaee9f24a9b..ca3c35136d311 100644 --- a/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py +++ b/python_modules/dagster/dagster_tests/core_tests/execution_tests/test_asset_execution_context.py @@ -1,12 +1,21 @@ -import dagster._check as check from dagster import OpExecutionContext, AssetExecutionContext, asset, materialize +import inspect +from dagster._core.execution.context.compute import _get_deprecation_kwargs +import warnings +import pytest def test_deprecation_warnings(): # Test that every method on OpExecutionContext is either reimplemented by AssetExecutionContext # or throws a deprecation warning on AssetExecutionContext - asset_execution_context_methods = [ + # If this test fails, it is likely because you added a method to OpExecutionContext without + # also adding that method to AssetExecutionContext. Please add the same method to AssetExecutionContext + # with an appropriate deprecation warning (see existing deprecated methods for examples). + # If the method should not be deprecated for AssetExecutionContext, please still add the same method + # to AssetExecutionContext and add the method name to asset_execution_context_attrs + + asset_execution_context_attrs = [ "op_execution_context", "is_asset_step", "asset_key", @@ -31,5 +40,78 @@ def test_deprecation_warnings(): "get_asset_provenance", "get_assets_code_version", "asset_check_spec", - "partition_kgey_range_for_asset_key", + "partition_key_range_for_asset_key", + "instance", + # TODO - potentially remove below if they get deprecated + "resources", + "run_config" + ] + + other_ignores = [ + "_abc_impl", + # TODO - what to do about instance vars for OpExecutionContext? listed below + "_events", + "_output_metadata", + "_pdb", + "_step_execution_context" ] + + + + @asset + def test_context(context: AssetExecutionContext): + asset_context = context + op_context = context.op_execution_context + + op_context_properties = [] + + for attr in dir(OpExecutionContext): + if isinstance(getattr(OpExecutionContext, attr), property): + op_context_properties.append(attr) + if attr in asset_execution_context_attrs: + continue + try: + deprecation_info = getattr(AssetExecutionContext, attr).fget._deprecated + except Exception: + raise Exception(f"Property {attr} on OpExecutionContext does not have an implementation on AssetExecutionContext. All methods on OpExecutionContext must be re-implemented on AssetExecutionContext with appropriate deprecation warnings. See the class implementation of AssetExecutionContext for more details.") + + expected_deprecation_args = _get_deprecation_kwargs(attr) + assert deprecation_info.breaking_version == expected_deprecation_args["breaking_version"] + assert deprecation_info.additional_warn_text == expected_deprecation_args["additional_warn_text"] + assert deprecation_info.subject == expected_deprecation_args["subject"] + + + + for method in dir(op_context): + if method in asset_execution_context_attrs or method[:2] == "__" or method in op_context_properties or method in other_ignores: + continue + + if inspect.ismethod(getattr(op_context, method)): + assert method in dir(asset_context) + try: + deprecation_info = getattr(asset_context, method)._deprecated + except Exception: + raise Exception(f"Method {method} on OpExecutionContext does not have an implementation on AssetExecutionContext. All methods on OpExecutionContext must be re-implemented on AssetExecutionContext with appropriate deprecation warnings. See the class implementation of AssetExecutionContext for more details.") + + expected_deprecation_args = _get_deprecation_kwargs(method) + assert deprecation_info.breaking_version == expected_deprecation_args["breaking_version"] + assert deprecation_info.additional_warn_text == expected_deprecation_args["additional_warn_text"] + assert deprecation_info.subject == expected_deprecation_args["subject"] + else: + raise Exception(f"Attribute {method} not accounted for in AssetExecutionContext deprecation test") + + + materialize([test_context]) + + +def test_instance_check(): + # turn off any outer warnings filters, e.g. ignores that are set in pyproject.toml + warnings.resetwarnings() + warnings.filterwarnings("error") + + @asset + def test_instance_check(context: AssetExecutionContext): + isinstance(context, OpExecutionContext) + + with pytest.raises(DeprecationWarning): + materialize([test_instance_check]) \ No newline at end of file