Use non-deprecated context in tests for Airflow 3

airflow-core/src/airflow/example_dags/tutorial_taskflow_templates.py

@@ -21,8 +21,7 @@
 # [START import_module]
 import pendulum
 
-from airflow.providers.standard.operators.python import get_current_context
-from airflow.sdk import dag, task
+from airflow.sdk import dag, get_current_context, task
 
 # [END import_module]
 

airflow-core/tests/unit/models/test_mappedoperator.py

@@ -334,7 +334,7 @@ def execute(self, context):
 
 
 def _create_mapped_with_name_template_taskflow(*, task_id, map_names, template):
-    from airflow.providers.standard.operators.python import get_current_context
+    from airflow.sdk import get_current_context
 
     @task(task_id=task_id, map_index_template=template)
     def task1(map_name):
@@ -360,7 +360,7 @@ def execute(self, context):
 
 
 def _create_named_map_index_renders_on_failure_taskflow(*, task_id, map_names, template):
-    from airflow.providers.standard.operators.python import get_current_context
+    from airflow.sdk import get_current_context
 
     @task(task_id=task_id, map_index_template=template)
     def task1(map_name):

airflow-core/tests/unit/models/test_taskinstance.py

@@ -3713,7 +3713,7 @@ def f(*args, **kwargs):
 
     def test_get_current_context_works_in_template(self, dag_maker):
         def user_defined_macro():
-            from airflow.providers.standard.operators.python import get_current_context
+            from airflow.sdk import get_current_context
 
             get_current_context()
 

providers/amazon/tests/system/amazon/aws/example_sagemaker.py

@@ -52,7 +52,6 @@
     SageMakerTransformSensor,
     SageMakerTuningSensor,
 )
-from airflow.providers.standard.operators.python import get_current_context
 from airflow.utils.trigger_rule import TriggerRule
 
 from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder, prune_logs
@@ -420,6 +419,13 @@ def set_up(env_id, role_arn):
     _install_aws_cli_if_needed()
     _build_and_upload_docker_image(preprocess_script, ecr_repository_uri)
 
+    from airflow.providers.amazon.version_compat import AIRFLOW_V_3_0_PLUS
+
+    if AIRFLOW_V_3_0_PLUS:
+        from airflow.sdk import get_current_context
+    else:
+        from airflow.providers.standard.operators.python import get_current_context
+
     ti = get_current_context()["ti"]
     ti.xcom_push(key="docker_image", value=ecr_repository_uri)
     ti.xcom_push(key="bucket_name", value=bucket_name)

providers/snowflake/src/airflow/providers/snowflake/operators/snowpark.py

@@ -20,7 +20,7 @@
 from collections.abc import Collection, Mapping, Sequence
 from typing import Any, Callable
 
-from airflow.providers.common.compat.standard.operators import PythonOperator, get_current_context
+from airflow.providers.common.compat.standard.operators import PythonOperator
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 from airflow.providers.snowflake.utils.snowpark import inject_session_into_op_kwargs
 
@@ -115,6 +115,14 @@ def execute_callable(self):
             session_parameters=self.session_parameters,
         )
         session = hook.get_snowpark_session()
+
+        from airflow.providers.snowflake.version_compat import AIRFLOW_V_3_0_PLUS
+
+        if AIRFLOW_V_3_0_PLUS:
+            from airflow.sdk import get_current_context
+        else:
+            from airflow.providers.standard.operators.python import get_current_context
+
         context = get_current_context()
         session.update_query_tag(
             {