Use aware datetime when needed

tests/api/common/test_delete_dag.py

@@ -16,14 +16,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import datetime
-
 import pytest
 
 from airflow import models
 from airflow.api.common.delete_dag import delete_dag
 from airflow.exceptions import AirflowException, DagNotFound
 from airflow.operators.empty import EmptyOperator
+from airflow.utils import timezone
 from airflow.utils.session import create_session
 from airflow.utils.state import State
 from airflow.utils.types import DagRunType
@@ -73,11 +72,11 @@ def setup_dag_models(self, for_sub_dag=False):
 
         task = EmptyOperator(
             task_id='dummy',
-            dag=models.DAG(dag_id=self.key, default_args={'start_date': datetime.datetime(2022, 1, 1)}),
+            dag=models.DAG(dag_id=self.key, default_args={'start_date': timezone.datetime(2022, 1, 1)}),
             owner='airflow',
         )
 
-        test_date = datetime.datetime(2022, 1, 1)
+        test_date = timezone.datetime(2022, 1, 1)
         with create_session() as session:
             session.add(DM(dag_id=self.key, fileloc=self.dag_file_path, is_subdag=for_sub_dag))
             dr = DR(dag_id=self.key, run_type=DagRunType.MANUAL, run_id="test", execution_date=test_date)

tests/api/common/test_mark_tasks.py

@@ -61,7 +61,7 @@ def create_dags(cls, dagbag):
         cls.dag2 = dagbag.get_dag('example_subdag_operator')
         cls.dag3 = dagbag.get_dag('example_trigger_target_dag')
         cls.dag4 = dagbag.get_dag('test_mapped_classic')
-        cls.execution_dates = [datetime.datetime(2022, 1, 1), datetime.datetime(2022, 1, 2)]
+        cls.execution_dates = [timezone.datetime(2022, 1, 1), timezone.datetime(2022, 1, 2)]
         start_date3 = cls.dag3.start_date
         cls.dag3_execution_dates = [
             start_date3,
@@ -482,9 +482,9 @@ def setup_class(cls):
         cls.dag2 = dagbag.dags['example_subdag_operator']
         cls.dag2.sync_to_db()
         cls.execution_dates = [
-            datetime.datetime(2022, 1, 1),
-            datetime.datetime(2022, 1, 1),
-            datetime.datetime(2022, 1, 1),
+            timezone.datetime(2022, 1, 1),
+            timezone.datetime(2022, 1, 2),
+            timezone.datetime(2022, 1, 3),
         ]
 
     def setup_method(self):

tests/api_connexion/endpoints/test_extra_link_endpoint.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import datetime
 import os
 from urllib.parse import quote_plus
 
@@ -29,6 +28,7 @@
 from airflow.plugins_manager import AirflowPlugin
 from airflow.providers.google.cloud.operators.bigquery import BigQueryExecuteQueryOperator
 from airflow.security import permissions
+from airflow.utils import timezone
 from airflow.utils.state import DagRunState
 from airflow.utils.types import DagRunType
 from tests.test_utils.api_connexion_utils import create_user, delete_user
@@ -61,7 +61,7 @@ def configured_app(minimal_app_for_api):
 class TestGetExtraLinks:
     @pytest.fixture(autouse=True)
     def setup_attrs(self, configured_app, session) -> None:
-        self.default_time = datetime.datetime(2020, 1, 1)
+        self.default_time = timezone.datetime(2020, 1, 1)
 
         clear_db_runs()
         clear_db_xcom()

tests/dag_processing/test_processor.py

@@ -111,7 +111,7 @@ def test_dag_file_processor_sla_miss_callback(self, create_dummy_dag):
 
         # Create dag with a start of 1 day ago, but an sla of 0
         # so we'll already have an sla_miss on the books.
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=1)
         dag, task = create_dummy_dag(
             dag_id='test_sla_miss',
             task_id='dummy',
@@ -140,7 +140,7 @@ def test_dag_file_processor_sla_miss_callback_invalid_sla(self, create_dummy_dag
         # Create dag with a start of 1 day ago, but an sla of 0
         # so we'll already have an sla_miss on the books.
         # Pass anything besides a timedelta object to the sla argument.
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=1)
         dag, task = create_dummy_dag(
             dag_id='test_sla_miss',
             task_id='dummy',
@@ -168,7 +168,7 @@ def test_dag_file_processor_sla_miss_callback_sent_notification(self, create_dum
 
         # Create dag with a start of 2 days ago, but an sla of 1 day
         # ago so we'll already have an sla_miss on the books
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=2)
         dag, task = create_dummy_dag(
             dag_id='test_sla_miss',
             task_id='dummy',
@@ -204,7 +204,7 @@ def test_dag_file_processor_sla_miss_doesnot_raise_integrity_error(self, dag_mak
 
         # Create dag with a start of 2 days ago, but an sla of 1 day
         # ago so we'll already have an sla_miss on the books
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=2)
         with dag_maker(
             dag_id='test_sla_miss',
             default_args={'start_date': test_start_date, 'sla': datetime.timedelta(days=1)},
@@ -245,7 +245,7 @@ def test_dag_file_processor_sla_miss_callback_exception(self, mock_stats_incr, c
 
         sla_callback = MagicMock(side_effect=RuntimeError('Could not call function'))
 
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=1)
         dag, task = create_dummy_dag(
             dag_id='test_sla_miss',
             task_id='dummy',
@@ -275,7 +275,7 @@ def test_dag_file_processor_only_collect_emails_from_sla_missed_tasks(
     ):
         session = settings.Session()
 
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=1)
         email1 = '[email protected]'
         dag, task = create_dummy_dag(
             dag_id='test_sla_miss',
@@ -315,7 +315,7 @@ def test_dag_file_processor_sla_miss_email_exception(
         # Mock the callback function so we can verify that it was not called
         mock_send_email.side_effect = RuntimeError('Could not send an email')
 
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=1)
         dag, task = create_dummy_dag(
             dag_id='test_sla_miss',
             task_id='dummy',
@@ -345,7 +345,7 @@ def test_dag_file_processor_sla_miss_deleted_task(self, create_dummy_dag):
         """
         session = settings.Session()
 
-        test_start_date = timezone.datetime(2022, 1, 1)
+        test_start_date = timezone.utcnow() - datetime.timedelta(days=1)
         dag, task = create_dummy_dag(
             dag_id='test_sla_miss',
             task_id='dummy',