amoghrajesh commented on code in PR #53791:
URL: https://github.com/apache/airflow/pull/53791#discussion_r2242480029
##########
providers/standard/tests/unit/standard/sensors/test_time_delta.py:
##########
@@ -34,12 +34,17 @@
WaitSensor,
)
from airflow.providers.standard.triggers.temporal import DateTimeTrigger
-from airflow.utils import timezone
-from airflow.utils.timezone import datetime
from airflow.utils.types import DagRunType
from tests_common.test_utils import db
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+ from airflow.sdk import timezone
+ from airflow.sdk.timezone import datetime
+else:
+ from airflow.utils import timezone # type: ignore[no-redef]
+ from airflow.utils.timezone import datetime # type: ignore[no-redef]
Review Comment:
Same comment as above.
Can merge the imports into one too
##########
providers/standard/tests/unit/standard/sensors/test_weekday.py:
##########
@@ -26,11 +26,16 @@
from airflow.models.dag import DAG
from airflow.providers.standard.sensors.weekday import DayOfWeekSensor
from airflow.providers.standard.utils.weekday import WeekDay
-from airflow.utils import timezone
-from airflow.utils.timezone import datetime
from tests_common.test_utils import db
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+ from airflow.sdk import timezone
+ from airflow.sdk.timezone import datetime
+else:
+ from airflow.utils import timezone # type: ignore[no-redef]
+ from airflow.utils.timezone import datetime # type: ignore[no-redef]
Review Comment:
Same comment as above
##########
providers/standard/tests/unit/standard/sensors/test_date_time.py:
##########
@@ -25,7 +25,13 @@
from airflow import macros
from airflow.models.dag import DAG
from airflow.providers.standard.sensors.date_time import DateTimeSensor
-from airflow.utils import timezone
+
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+ from airflow.sdk import timezone
+else:
+ from airflow.utils import timezone # type: ignore[no-redef]
Review Comment:
try / except here too please
##########
providers/standard/tests/unit/standard/sensors/test_python.py:
##########
@@ -24,10 +24,15 @@
from airflow.exceptions import AirflowSensorTimeout
from airflow.providers.standard.sensors.python import PythonSensor
-from airflow.sensors.base import PokeReturnValue
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
from unit.standard.operators.test_python import BasePythonTest
+if AIRFLOW_V_3_1_PLUS:
+ from airflow.sdk.bases.sensor import PokeReturnValue
+else:
+ from airflow.sensors.base import PokeReturnValue # type: ignore[no-redef]
Review Comment:
You could just add this in version_compat.py and import them like:
```
from tests_common.test_utils.version_compat import PokeReturnValue
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
