fixup! Resolve lazy_object_proxy for PyVirtualEnvOperator

Author: kaxil

providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_adx.py

@@ -25,6 +25,7 @@
 from airflow.models import DAG
 from airflow.providers.microsoft.azure.hooks.adx import AzureDataExplorerHook
 from airflow.providers.microsoft.azure.operators.adx import AzureDataExplorerQueryOperator
+from airflow.providers.microsoft.azure.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils.timezone import datetime
 
 TEST_DAG_ID = "unit_tests"
@@ -88,12 +89,20 @@ def test_azure_data_explorer_query_operator_xcom_push_and_pull(
     mock_conn,
     mock_run_query,
     create_task_instance_of_operator,
+    request,
 ):
-    ti = create_task_instance_of_operator(
-        AzureDataExplorerQueryOperator,
-        dag_id="test_azure_data_explorer_query_operator_xcom_push_and_pull",
-        **MOCK_DATA,
-    )
-    ti.run()
-
-    assert ti.xcom_pull(task_ids=MOCK_DATA["task_id"]) == str(MOCK_RESULT)
+    if AIRFLOW_V_3_0_PLUS:
+        run_task = request.getfixturevalue("run_task")
+        task = AzureDataExplorerQueryOperator(**MOCK_DATA)
+        run_task(task=task)
+
+        assert run_task.xcom.get(key="return_value", task_id=task.task_id) == str(MOCK_RESULT)
+    else:
+        ti = create_task_instance_of_operator(
+            AzureDataExplorerQueryOperator,
+            dag_id="test_azure_data_explorer_query_operator_xcom_push_and_pull",
+            **MOCK_DATA,
+        )
+        ti.run()
+
+        assert ti.xcom_pull(task_ids=MOCK_DATA["task_id"]) == str(MOCK_RESULT)

providers/oracle/tests/unit/oracle/operators/test_oracle.py

@@ -27,6 +27,8 @@
 from airflow.providers.oracle.hooks.oracle import OracleHook
 from airflow.providers.oracle.operators.oracle import OracleStoredProcedureOperator
 
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
 
 class TestOracleStoredProcedureOperator:
     @mock.patch.object(OracleHook, "run", autospec=OracleHook.run)
@@ -65,12 +67,20 @@ def test_push_oracle_exit_to_xcom(self, mock_callproc, request, dag_maker):
         error = f"ORA-{ora_exit_code}: This is a five-digit ORA error code"
         mock_callproc.side_effect = oracledb.DatabaseError(error)
 
-        with dag_maker(dag_id=f"dag_{request.node.name}"):
+        if AIRFLOW_V_3_0_PLUS:
+            run_task = request.getfixturevalue("run_task")
             task = OracleStoredProcedureOperator(
                 procedure=procedure, oracle_conn_id=oracle_conn_id, parameters=parameters, task_id=task_id
             )
-        dr = dag_maker.create_dagrun(run_id=task_id)
-        ti = TaskInstance(task=task, run_id=dr.run_id)
-        with pytest.raises(oracledb.DatabaseError, match=re.escape(error)):
-            ti.run()
-        assert ti.xcom_pull(task_ids=task.task_id, key="ORA") == ora_exit_code
+            run_task(task=task)
+            assert run_task.xcom.get(task_id=task.task_id, key="ORA") == ora_exit_code
+        else:
+            with dag_maker(dag_id=f"dag_{request.node.name}"):
+                task = OracleStoredProcedureOperator(
+                    procedure=procedure, oracle_conn_id=oracle_conn_id, parameters=parameters, task_id=task_id
+                )
+            dr = dag_maker.create_dagrun(run_id=task_id)
+            ti = TaskInstance(task=task, run_id=dr.run_id)
+            with pytest.raises(oracledb.DatabaseError, match=re.escape(error)):
+                ti.run()
+            assert ti.xcom_pull(task_ids=task.task_id, key="ORA") == ora_exit_code

providers/snowflake/tests/unit/snowflake/decorators/test_snowpark.py

@@ -24,6 +24,7 @@
 import pytest
 
 from airflow.decorators import task
+from airflow.providers.snowflake.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils import timezone
 
 if TYPE_CHECKING:
@@ -156,7 +157,7 @@ def func(session: Session):
         mock_snowflake_hook.return_value.get_snowpark_session.assert_called_once()
 
     @mock.patch("airflow.providers.snowflake.operators.snowpark.SnowflakeHook")
-    def test_snowpark_decorator_multiple_output(self, mock_snowflake_hook, dag_maker):
+    def test_snowpark_decorator_multiple_output(self, mock_snowflake_hook, dag_maker, request):
         @task.snowpark(
             task_id=TASK_ID,
             snowflake_conn_id=CONN_ID,
@@ -171,15 +172,23 @@ def func(session: Session):
             assert session == mock_snowflake_hook.return_value.get_snowpark_session.return_value
             return {"a": 1, "b": "2"}
 
-        with dag_maker(dag_id=TEST_DAG_ID):
-            ret = func()
-
-        dr = dag_maker.create_dagrun()
-        ret.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
-        ti = dr.get_task_instances()[0]
-        assert ti.xcom_pull(key="a") == 1
-        assert ti.xcom_pull(key="b") == "2"
-        assert ti.xcom_pull() == {"a": 1, "b": "2"}
+        if AIRFLOW_V_3_0_PLUS:
+            run_task = request.getfixturevalue("run_task")
+            op = func().operator
+            run_task(task=op)
+            assert run_task.xcom.get(key="a") == 1
+            assert run_task.xcom.get(key="b") == "2"
+            assert run_task.xcom.get(key="return_value") == {"a": 1, "b": "2"}
+        else:
+            with dag_maker(dag_id=TEST_DAG_ID):
+                ret = func()
+
+            dr = dag_maker.create_dagrun()
+            ret.operator.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
+            ti = dr.get_task_instances()[0]
+            assert ti.xcom_pull(key="a") == 1
+            assert ti.xcom_pull(key="b") == "2"
+            assert ti.xcom_pull() == {"a": 1, "b": "2"}
         mock_snowflake_hook.assert_called_once()
         mock_snowflake_hook.return_value.get_snowpark_session.assert_called_once()
 

providers/standard/tests/unit/standard/decorators/test_python.py

@@ -215,7 +215,22 @@ def identity_notyping_with_decorator_call(x: int):
 
         assert identity_notyping_with_decorator_call(5).operator.multiple_outputs is False
 
-    def test_manual_multiple_outputs_false_with_typings(self):
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Different test for AF 2")
+    def test_manual_multiple_outputs_false_with_typings(self, run_task):
+        @task_decorator(multiple_outputs=False)
+        def identity2(x: int, y: int) -> tuple[int, int]:
+            return x, y
+
+        res = identity2(8, 4)
+        run_task(task=res.operator)
+
+        assert not res.operator.multiple_outputs
+        assert run_task.xcom.get(key=res.key) == (8, 4)
+        assert run_task.xcom.get(key="return_value_0") is None
+        assert run_task.xcom.get(key="return_value_1") is None
+
+    @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Different test for AF 3")
+    def test_manual_multiple_outputs_false_with_typings_af2(self):
         @task_decorator(multiple_outputs=False)
         def identity2(x: int, y: int) -> tuple[int, int]:
             return x, y
@@ -233,7 +248,22 @@ def identity2(x: int, y: int) -> tuple[int, int]:
         assert ti.xcom_pull(key="return_value_0") is None
         assert ti.xcom_pull(key="return_value_1") is None
 
-    def test_multiple_outputs_ignore_typing(self):
+    @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Different test for AF 2")
+    def test_multiple_outputs_ignore_typing(self, run_task):
+        @task_decorator
+        def identity_tuple(x: int, y: int) -> tuple[int, int]:
+            return x, y
+
+        ident = identity_tuple(35, 36)
+        run_task(task=ident.operator)
+
+        assert not ident.operator.multiple_outputs
+        assert run_task.xcom.get(key=ident.key) == (35, 36)
+        assert run_task.xcom.get(key="return_value_0") is None
+        assert run_task.xcom.get(key="return_value_1") is None
+
+    @pytest.mark.skipif(AIRFLOW_V_3_0_PLUS, reason="Different test for AF 3")
+    def test_multiple_outputs_ignore_typing_af2(self):
         @task_decorator
         def identity_tuple(x: int, y: int) -> tuple[int, int]:
             return x, y

providers/standard/tests/unit/standard/operators/test_datetime.py

@@ -214,7 +214,7 @@ def test_branch_datetime_operator_upper_comparison_outside_range(self, target_up
 
     @pytest.mark.parametrize("target_lower", [target_lower for (target_lower, _) in targets])
     @time_machine.travel("2020-07-07 09:00:00")
-    def test_branch_datetime_operator_lower_comparison_outside_range(self, target_lower, run_task):
+    def test_branch_datetime_operator_lower_comparison_outside_range(self, target_lower):
         """Check BranchDateTimeOperator branch operation"""
         self.branch_op.target_lower = target_lower
         self.branch_op.target_upper = None