diff --git a/tests/api/common/test_delete_dag.py b/tests/api/common/test_delete_dag.py index 2f89170a217c1..cf51856409cd2 100644 --- a/tests/api/common/test_delete_dag.py +++ b/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) diff --git a/tests/api/common/test_mark_tasks.py b/tests/api/common/test_mark_tasks.py index 89ad399d6f753..fee33a50e88e6 100644 --- a/tests/api/common/test_mark_tasks.py +++ b/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): diff --git a/tests/api_connexion/endpoints/test_extra_link_endpoint.py b/tests/api_connexion/endpoints/test_extra_link_endpoint.py index 6cb29987f0af5..c209f661a8f66 100644 --- a/tests/api_connexion/endpoints/test_extra_link_endpoint.py +++ b/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() diff --git a/tests/dag_processing/test_processor.py b/tests/dag_processing/test_processor.py index 2886de6b65c58..e2f9165131f2f 100644 --- a/tests/dag_processing/test_processor.py +++ b/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 = 'test1@test.com' 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',