Provider Migration: Update docker for Airflow 3.0 compatibility

providers/docker/src/airflow/providers/docker/operators/docker.py

@@ -34,12 +34,12 @@
 from docker.types import LogConfig, Mount, Ulimit
 from dotenv import dotenv_values
 
-from airflow.models import BaseOperator
 from airflow.providers.docker.exceptions import (
     DockerContainerFailedException,
     DockerContainerFailedSkipException,
 )
 from airflow.providers.docker.hooks.docker import DockerHook
+from airflow.providers.docker.version_compat import BaseOperator
 
 if TYPE_CHECKING:
     from logging import Logger

providers/docker/src/airflow/providers/docker/version_compat.py

@@ -33,3 +33,11 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
 
 
 AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
+
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseOperator
+else:
+    from airflow.models import BaseOperator
+
+__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator"]

providers/docker/tests/unit/docker/decorators/test_docker.py

@@ -53,11 +53,11 @@ def f():
             return [random.random() for _ in range(100)]
 
         with dag_maker(session=session):
-            ret = f()
+            f()
         session.commit()
         dr = dag_maker.create_dagrun(session=session)
         session.expunge_all()
-        ret.operator.run(start_date=dr.logical_date, end_date=dr.logical_date, session=session)
+        dag_maker.run_ti("f", dr)
         ti = dr.get_task_instances(session=session)[0]
         assert len(ti.xcom_pull()) == 100
 
@@ -70,11 +70,11 @@ def f(num_results):
             return [random.random() for _ in range(num_results)]
 
         with dag_maker(session=session):
-            ret = f(50)
+            f(50)
 
         dr = dag_maker.create_dagrun(session=session)
         session.expunge_all()
-        ret.operator.run(start_date=dr.logical_date, end_date=dr.logical_date)
+        dag_maker.run_ti("f", dr)
         ti = dr.get_task_instances(session=session)[0]
         result = ti.xcom_pull(session=session)
         assert isinstance(result, list)
@@ -102,13 +102,11 @@ def return_dict(number: int):
 
         test_number = 10
         with dag_maker(session=session):
-            ret = return_dict(test_number)
+            return_dict(test_number)
 
         dr = dag_maker.create_dagrun(session=session)
         session.expunge_all()
-
-        ret.operator.run(start_date=dr.logical_date, end_date=dr.logical_date, session=session)
-
+        dag_maker.run_ti("return_dict", dr)
         ti = dr.get_task_instances(session=session)[0]
         assert ti.xcom_pull(key="number", session=session) == test_number + 1
         assert ti.xcom_pull(key="43", session=session) == 43
@@ -121,12 +119,11 @@ def f():
             pass
 
         with dag_maker(session=session):
-            ret = f()
+            f()
 
         dr = dag_maker.create_dagrun(session=session)
         session.expunge_all()
-
-        ret.operator.run(start_date=dr.logical_date, end_date=dr.logical_date, session=session)
+        dag_maker.run_ti("f", dr)
         ti = dr.get_task_instances(session=session)[0]
         assert ti.xcom_pull(session=session) is None
 
@@ -174,15 +171,16 @@ def f(exit_code):
             raise SystemExit(exit_code)
 
         with dag_maker(session=session):
-            ret = f(actual_exit_code)
+            f(actual_exit_code)
 
         dr = dag_maker.create_dagrun(session=session)
         session.expunge_all()
         if expected_state == TaskInstanceState.FAILED:
             with pytest.raises(AirflowException):
-                ret.operator.run(start_date=dr.logical_date, end_date=dr.logical_date, session=session)
+                dag_maker.run_ti("f", dr)
+
         else:
-            ret.operator.run(start_date=dr.logical_date, end_date=dr.logical_date, session=session)
+            dag_maker.run_ti("f", dr)
             ti = dr.get_task_instances(session=session)[0]
             assert ti.state == expected_state
 
@@ -343,13 +341,13 @@ def f():
         ch = logging.StreamHandler(log_capture_string)
         docker_operator_logger.addHandler(ch)
         with dag_maker(session=session):
-            ret = f()
+            f()
 
         dr = dag_maker.create_dagrun(session=session)
         session.expunge_all()
 
         with pytest.raises(AirflowException):
-            ret.operator.run(start_date=dr.logical_date, end_date=dr.logical_date, session=session)
+            dag_maker.run_ti("f", dr)
         ti = dr.get_task_instances(session=session)[0]
         assert ti.state == TaskInstanceState.FAILED