use correct context type annotations in tests (#16507)

dagster-io · Sep 25, 2023 · 88396f1 · 88396f1
1 parent 0cab9ca
commit 88396f1
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/python_modules/dagster-graphql/dagster_graphql_tests/graphql/repo.py b/python_modules/dagster-graphql/dagster_graphql_tests/graphql/repo.py
@@ -41,6 +41,7 @@
     Map,
     Noneable,
     Nothing,
+    OpExecutionContext,
     Out,
     Output,
     PythonObjectDagsterType,
@@ -733,7 +734,7 @@ def spawn() -> int:
     @op(
         required_resource_keys={"retry_count"},
     )
-    def fail(context: AssetExecutionContext, depth: int) -> int:
+    def fail(context: OpExecutionContext, depth: int) -> int:
         if context.resources.retry_count <= depth:
             raise Exception("fail")
 

diff --git a/python_modules/dagster/dagster_tests/execution_tests/test_execute_job.py b/python_modules/dagster/dagster_tests/execution_tests/test_execute_job.py
@@ -1,12 +1,12 @@
 import dagster._check as check
 import pytest
 from dagster import (
-    AssetExecutionContext,
     AssetKey,
     DagsterExecutionStepNotFoundError,
     DagsterInvalidConfigError,
     DagsterInvariantViolationError,
     Field,
+    OpExecutionContext,
     Out,
     Output,
     ReexecutionOptions,
@@ -377,7 +377,7 @@ def echo(x):
 
 
 @op
-def fail_once(context: AssetExecutionContext, x):
+def fail_once(context: OpExecutionContext, x):
     key = context.op_handle.name
     if context.instance.run_storage.get_cursor_values({key}).get(key):
         return x

diff --git a/python_modules/libraries/dagster-dbt/dagster_dbt_tests/core/test_resources_v2.py b/python_modules/libraries/dagster-dbt/dagster_dbt_tests/core/test_resources_v2.py
@@ -16,7 +16,7 @@
     materialize,
     op,
 )
-from dagster._core.execution.context.compute import OpExecutionContext
+from dagster._core.execution.context.compute import AssetExecutionContext
 from dagster_dbt import dbt_assets
 from dagster_dbt.asset_utils import build_dbt_asset_selection
 from dagster_dbt.core.resources_v2 import (
@@ -209,7 +209,7 @@ def test_dbt_with_partial_parse() -> None:
 
 def test_dbt_cli_debug_execution() -> None:
     @dbt_assets(manifest=manifest)
-    def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource):
+    def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource):
         yield from dbt.cli(["--debug", "run"], context=context).stream()
 
     result = materialize(
@@ -230,7 +230,7 @@ def test_dbt_cli_subsetted_execution() -> None:
     )
 
     @dbt_assets(manifest=manifest, select=dbt_select)
-    def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource):
+    def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource):
         dbt_cli_invocation = dbt.cli(["run"], context=context)
 
         assert dbt_cli_invocation.process.args == ["dbt", "run", "--select", dbt_select]
@@ -253,7 +253,7 @@ def test_dbt_cli_asset_selection() -> None:
     ]
 
     @dbt_assets(manifest=manifest)
-    def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource):
+    def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource):
         dbt_cli_invocation = dbt.cli(["run"], context=context)
 
         dbt_cli_args: List[str] = list(dbt_cli_invocation.process.args)  # type: ignore
@@ -283,7 +283,7 @@ def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource):
 @pytest.mark.parametrize("exclude", [None, "fqn:dagster_dbt_test_project.subdir.least_caloric"])
 def test_dbt_cli_default_selection(exclude: Optional[str]) -> None:
     @dbt_assets(manifest=manifest, exclude=exclude)
-    def my_dbt_assets(context: OpExecutionContext, dbt: DbtCliResource):
+    def my_dbt_assets(context: AssetExecutionContext, dbt: DbtCliResource):
         dbt_cli_invocation = dbt.cli(["run"], context=context)
 
         expected_args = ["dbt", "run", "--select", "fqn:*"]