From 55c25458c14e41dae6b95fb85cfb0aa7ccd9576f Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU <68415893+jason810496@users.noreply.github.com> Date: Wed, 27 Nov 2024 21:42:46 +0800 Subject: [PATCH] AIP-84 Refactor test cases with `datetime_zulu_format` (#44323) * Refactor test_dag_run * Refactor test_assets * Rename datetime_zulu_format to from_datetime_to_zulu * Fix TestListDagRunsBatch --- .../core_api/routes/public/test_assets.py | 95 +++++++++---------- .../core_api/routes/public/test_dag_run.py | 51 +++++----- .../core_api/routes/public/test_event_logs.py | 6 +- .../routes/public/test_import_error.py | 4 +- .../core_api/routes/public/test_job.py | 6 +- tests_common/test_utils/format_datetime.py | 4 +- 6 files changed, 80 insertions(+), 86 deletions(-) diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index 0e7db3910531d..d4b4172c33147 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -18,6 +18,7 @@ import urllib from collections.abc import Generator +from datetime import datetime from unittest import mock import pytest @@ -39,12 +40,14 @@ from tests_common.test_utils.asserts import assert_queries_count from tests_common.test_utils.db import clear_db_assets, clear_db_runs +from tests_common.test_utils.format_datetime import from_datetime_to_zulu_without_ms + +DEFAULT_DATE = datetime(2020, 6, 11, 18, 0, 0, tzinfo=timezone.utc) pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] def _create_assets(session, num: int = 2) -> None: - default_time = "2020-06-11T18:00:00+00:00" assets = [ AssetModel( id=i, @@ -52,8 +55,8 @@ def _create_assets(session, num: int = 2) -> None: uri=f"s3://bucket/key/{i}", group="asset", extra={"foo": "bar"}, - created_at=timezone.parse(default_time), - updated_at=timezone.parse(default_time), + created_at=DEFAULT_DATE, + updated_at=DEFAULT_DATE, ) for i in range(1, 1 + num) ] @@ -62,7 +65,6 @@ def _create_assets(session, num: int = 2) -> None: def _create_assets_with_sensitive_extra(session, num: int = 2) -> None: - default_time = "2020-06-11T18:00:00+00:00" assets = [ AssetModel( id=i, @@ -70,8 +72,8 @@ def _create_assets_with_sensitive_extra(session, num: int = 2) -> None: uri=f"s3://bucket/key/{i}", group="asset", extra={"password": "bar"}, - created_at=timezone.parse(default_time), - updated_at=timezone.parse(default_time), + created_at=DEFAULT_DATE, + updated_at=DEFAULT_DATE, ) for i in range(1, 1 + num) ] @@ -85,7 +87,6 @@ def _create_provided_asset(session, asset: AssetModel) -> None: def _create_assets_events(session, num: int = 2) -> None: - default_time = "2020-06-11T18:00:00+00:00" assets_events = [ AssetEvent( id=i, @@ -94,7 +95,7 @@ def _create_assets_events(session, num: int = 2) -> None: source_task_id="source_task_id", source_dag_id="source_dag_id", source_run_id=f"source_run_id_{i}", - timestamp=timezone.parse(default_time), + timestamp=DEFAULT_DATE, ) for i in range(1, 1 + num) ] @@ -103,7 +104,6 @@ def _create_assets_events(session, num: int = 2) -> None: def _create_assets_events_with_sensitive_extra(session, num: int = 2) -> None: - default_time = "2020-06-11T18:00:00+00:00" assets_events = [ AssetEvent( id=i, @@ -112,7 +112,7 @@ def _create_assets_events_with_sensitive_extra(session, num: int = 2) -> None: source_task_id="source_task_id", source_dag_id="source_dag_id", source_run_id=f"source_run_id_{i}", - timestamp=timezone.parse(default_time), + timestamp=DEFAULT_DATE, ) for i in range(1, 1 + num) ] @@ -126,22 +126,21 @@ def _create_provided_asset_event(session, asset_event: AssetEvent) -> None: def _create_dag_run(session, num: int = 2): - default_time = "2020-06-11T18:00:00+00:00" dag_runs = [ DagRun( dag_id="source_dag_id", run_id=f"source_run_id_{i}", run_type=DagRunType.MANUAL, - logical_date=timezone.parse(default_time), - start_date=timezone.parse(default_time), - data_interval=(timezone.parse(default_time), timezone.parse(default_time)), + logical_date=DEFAULT_DATE, + start_date=DEFAULT_DATE, + data_interval=(DEFAULT_DATE, DEFAULT_DATE), external_trigger=True, state=DagRunState.SUCCESS, ) for i in range(1, 1 + num) ] for dag_run in dag_runs: - dag_run.end_date = timezone.parse(default_time) + dag_run.end_date = DEFAULT_DATE session.add_all(dag_runs) session.commit() @@ -156,11 +155,9 @@ def _create_asset_dag_run(session, num: int = 2): class TestAssets: - default_time = "2020-06-11T18:00:00+00:00" - @pytest.fixture def time_freezer(self) -> Generator: - freezer = time_machine.travel(self.default_time, tick=False) + freezer = time_machine.travel(DEFAULT_DATE, tick=False) freezer.start() yield @@ -218,7 +215,7 @@ def test_should_respond_200(self, test_client, session): response = test_client.get("/public/assets") assert response.status_code == 200 response_data = response.json() - tz_datetime_format = self.default_time.replace("+00:00", "Z") + tz_datetime_format = from_datetime_to_zulu_without_ms(DEFAULT_DATE) assert response_data == { "assets": [ { @@ -402,15 +399,15 @@ def test_should_respond_200(self, test_client, session): { "run_id": "source_run_id_1", "dag_id": "source_dag_id", - "logical_date": "2020-06-11T18:00:00Z", - "start_date": "2020-06-11T18:00:00Z", - "end_date": "2020-06-11T18:00:00Z", + "logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), "state": "success", - "data_interval_start": "2020-06-11T18:00:00Z", - "data_interval_end": "2020-06-11T18:00:00Z", + "data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE), } ], - "timestamp": "2020-06-11T18:00:00Z", + "timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE), }, { "id": 2, @@ -425,15 +422,15 @@ def test_should_respond_200(self, test_client, session): { "run_id": "source_run_id_2", "dag_id": "source_dag_id", - "logical_date": "2020-06-11T18:00:00Z", - "start_date": "2020-06-11T18:00:00Z", - "end_date": "2020-06-11T18:00:00Z", + "logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), "state": "success", - "data_interval_start": "2020-06-11T18:00:00Z", - "data_interval_end": "2020-06-11T18:00:00Z", + "data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE), } ], - "timestamp": "2020-06-11T18:00:00Z", + "timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE), }, ], "total_entries": 2, @@ -514,15 +511,15 @@ def test_should_mask_sensitive_extra(self, test_client, session): { "run_id": "source_run_id_1", "dag_id": "source_dag_id", - "logical_date": "2020-06-11T18:00:00Z", - "start_date": "2020-06-11T18:00:00Z", - "end_date": "2020-06-11T18:00:00Z", + "logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), "state": "success", - "data_interval_start": "2020-06-11T18:00:00Z", - "data_interval_end": "2020-06-11T18:00:00Z", + "data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE), } ], - "timestamp": "2020-06-11T18:00:00Z", + "timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE), }, { "id": 2, @@ -537,15 +534,15 @@ def test_should_mask_sensitive_extra(self, test_client, session): { "run_id": "source_run_id_2", "dag_id": "source_dag_id", - "logical_date": "2020-06-11T18:00:00Z", - "start_date": "2020-06-11T18:00:00Z", - "end_date": "2020-06-11T18:00:00Z", + "logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE), "state": "success", - "data_interval_start": "2020-06-11T18:00:00Z", - "data_interval_end": "2020-06-11T18:00:00Z", + "data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE), + "data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE), } ], - "timestamp": "2020-06-11T18:00:00Z", + "timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE), }, ], "total_entries": 2, @@ -566,7 +563,7 @@ class TestGetAssetEndpoint(TestAssets): def test_should_respond_200(self, test_client, url, session): self.create_assets(num=1) assert session.query(AssetModel).count() == 1 - tz_datetime_format = self.default_time.replace("+00:00", "Z") + tz_datetime_format = from_datetime_to_zulu_without_ms(DEFAULT_DATE) with assert_queries_count(6): response = test_client.get( f"/public/assets/{url}", @@ -596,7 +593,7 @@ def test_should_respond_404(self, test_client): @pytest.mark.enable_redact def test_should_mask_sensitive_extra(self, test_client, session): self.create_assets_with_sensitive_extra() - tz_datetime_format = self.default_time.replace("+00:00", "Z") + tz_datetime_format = from_datetime_to_zulu_without_ms(DEFAULT_DATE) uri = "s3://bucket/key/1" response = test_client.get( f"/public/assets/{uri}", @@ -641,7 +638,7 @@ def test_should_respond_200(self, test_client, session, create_dummy_dag): assert response.json() == { "queued_events": [ { - "created_at": self.default_time.replace("+00:00", "Z"), + "created_at": from_datetime_to_zulu_without_ms(DEFAULT_DATE), "uri": "s3://bucket/key/1", "dag_id": "dag", } @@ -721,7 +718,7 @@ def test_should_respond_200(self, test_client, session): "source_run_id": None, "source_map_index": -1, "created_dagruns": [], - "timestamp": self.default_time.replace("+00:00", "Z"), + "timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE), } def test_invalid_attr_not_allowed(self, test_client, session): @@ -748,7 +745,7 @@ def test_should_mask_sensitive_extra(self, test_client, session): "source_run_id": None, "source_map_index": -1, "created_dagruns": [], - "timestamp": self.default_time.replace("+00:00", "Z"), + "timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE), } @@ -769,7 +766,7 @@ def test_should_respond_200(self, test_client, session, create_dummy_dag): assert response.json() == { "queued_events": [ { - "created_at": self.default_time.replace("+00:00", "Z"), + "created_at": from_datetime_to_zulu_without_ms(DEFAULT_DATE), "uri": "s3://bucket/key/1", "dag_id": "dag", } diff --git a/tests/api_fastapi/core_api/routes/public/test_dag_run.py b/tests/api_fastapi/core_api/routes/public/test_dag_run.py index d453c973c8d0f..bd0657a9e79cb 100644 --- a/tests/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/tests/api_fastapi/core_api/routes/public/test_dag_run.py @@ -39,6 +39,7 @@ clear_db_runs, clear_db_serialized_dags, ) +from tests_common.test_utils.format_datetime import from_datetime_to_zulu, from_datetime_to_zulu_without_ms pytestmark = pytest.mark.db_test @@ -192,22 +193,20 @@ def test_get_dag_run_not_found(self, test_client): class TestGetDagRuns: - @staticmethod - def parse_datetime(datetime_str): - return datetime_str.isoformat().replace("+00:00", "Z") if datetime_str else None - @staticmethod def get_dag_run_dict(run: DagRun): return { "dag_run_id": run.run_id, "dag_id": run.dag_id, - "logical_date": TestGetDagRuns.parse_datetime(run.logical_date), - "queued_at": TestGetDagRuns.parse_datetime(run.queued_at), - "start_date": TestGetDagRuns.parse_datetime(run.start_date), - "end_date": TestGetDagRuns.parse_datetime(run.end_date), - "data_interval_start": TestGetDagRuns.parse_datetime(run.data_interval_start), - "data_interval_end": TestGetDagRuns.parse_datetime(run.data_interval_end), - "last_scheduling_decision": TestGetDagRuns.parse_datetime(run.last_scheduling_decision), + "logical_date": from_datetime_to_zulu_without_ms(run.logical_date), + "queued_at": from_datetime_to_zulu(run.queued_at) if run.queued_at else None, + "start_date": from_datetime_to_zulu_without_ms(run.start_date), + "end_date": from_datetime_to_zulu(run.end_date), + "data_interval_start": from_datetime_to_zulu_without_ms(run.data_interval_start), + "data_interval_end": from_datetime_to_zulu_without_ms(run.data_interval_end), + "last_scheduling_decision": from_datetime_to_zulu(run.last_scheduling_decision) + if run.last_scheduling_decision + else None, "run_type": run.run_type, "state": run.state, "external_trigger": run.external_trigger, @@ -482,22 +481,20 @@ def test_invalid_state(self, test_client): class TestListDagRunsBatch: - @staticmethod - def parse_datetime(datetime_str): - return datetime_str.isoformat().replace("+00:00", "Z") if datetime_str else None - @staticmethod def get_dag_run_dict(run: DagRun): return { "dag_run_id": run.run_id, "dag_id": run.dag_id, - "logical_date": TestGetDagRuns.parse_datetime(run.logical_date), - "queued_at": TestGetDagRuns.parse_datetime(run.queued_at), - "start_date": TestGetDagRuns.parse_datetime(run.start_date), - "end_date": TestGetDagRuns.parse_datetime(run.end_date), - "data_interval_start": TestGetDagRuns.parse_datetime(run.data_interval_start), - "data_interval_end": TestGetDagRuns.parse_datetime(run.data_interval_end), - "last_scheduling_decision": TestGetDagRuns.parse_datetime(run.last_scheduling_decision), + "logical_date": from_datetime_to_zulu_without_ms(run.logical_date), + "queued_at": from_datetime_to_zulu_without_ms(run.queued_at) if run.queued_at else None, + "start_date": from_datetime_to_zulu_without_ms(run.start_date), + "end_date": from_datetime_to_zulu(run.end_date), + "data_interval_start": from_datetime_to_zulu_without_ms(run.data_interval_start), + "data_interval_end": from_datetime_to_zulu_without_ms(run.data_interval_end), + "last_scheduling_decision": from_datetime_to_zulu_without_ms(run.last_scheduling_decision) + if run.last_scheduling_decision + else None, "run_type": run.run_type, "state": run.state, "external_trigger": run.external_trigger, @@ -976,7 +973,7 @@ def test_should_respond_200(self, test_client, dag_maker, session): expected_response = { "asset_events": [ { - "timestamp": event.timestamp.isoformat().replace("+00:00", "Z"), + "timestamp": from_datetime_to_zulu(event.timestamp), "asset_id": asset1_id, "uri": asset1.uri, "extra": {}, @@ -989,11 +986,11 @@ def test_should_respond_200(self, test_client, dag_maker, session): { "dag_id": "TEST_DAG_ID", "run_id": "TEST_DAG_RUN_ID", - "data_interval_end": dr.data_interval_end.isoformat().replace("+00:00", "Z"), - "data_interval_start": dr.data_interval_start.isoformat().replace("+00:00", "Z"), + "data_interval_end": from_datetime_to_zulu_without_ms(dr.data_interval_end), + "data_interval_start": from_datetime_to_zulu_without_ms(dr.data_interval_start), "end_date": None, - "logical_date": dr.logical_date.isoformat().replace("+00:00", "Z"), - "start_date": dr.start_date.isoformat().replace("+00:00", "Z"), + "logical_date": from_datetime_to_zulu_without_ms(dr.logical_date), + "start_date": from_datetime_to_zulu_without_ms(dr.start_date), "state": "running", } ], diff --git a/tests/api_fastapi/core_api/routes/public/test_event_logs.py b/tests/api_fastapi/core_api/routes/public/test_event_logs.py index 84ab5d168c9d0..090fe1ea01cbb 100644 --- a/tests/api_fastapi/core_api/routes/public/test_event_logs.py +++ b/tests/api_fastapi/core_api/routes/public/test_event_logs.py @@ -24,7 +24,7 @@ from airflow.utils.session import provide_session from tests_common.test_utils.db import clear_db_logs, clear_db_runs -from tests_common.test_utils.format_datetime import datetime_zulu_format, datetime_zulu_format_without_ms +from tests_common.test_utils.format_datetime import from_datetime_to_zulu, from_datetime_to_zulu_without_ms pytestmark = pytest.mark.db_test @@ -161,14 +161,14 @@ def test_get_event_log(self, test_client, setup, event_log_key, expected_status_ expected_json = { "event_log_id": event_log_id, - "when": datetime_zulu_format(event_log.dttm) if event_log.dttm else None, + "when": from_datetime_to_zulu(event_log.dttm) if event_log.dttm else None, "dag_id": expected_body.get("dag_id"), "task_id": expected_body.get("task_id"), "run_id": expected_body.get("run_id"), "map_index": event_log.map_index, "try_number": event_log.try_number, "event": expected_body.get("event"), - "logical_date": datetime_zulu_format_without_ms(event_log.logical_date) + "logical_date": from_datetime_to_zulu_without_ms(event_log.logical_date) if event_log.logical_date else None, "owner": expected_body.get("owner"), diff --git a/tests/api_fastapi/core_api/routes/public/test_import_error.py b/tests/api_fastapi/core_api/routes/public/test_import_error.py index 91c01d282d3e2..33a3fbaac745c 100644 --- a/tests/api_fastapi/core_api/routes/public/test_import_error.py +++ b/tests/api_fastapi/core_api/routes/public/test_import_error.py @@ -24,7 +24,7 @@ from airflow.utils.session import provide_session from tests_common.test_utils.db import clear_db_import_errors -from tests_common.test_utils.format_datetime import datetime_zulu_format_without_ms +from tests_common.test_utils.format_datetime import from_datetime_to_zulu_without_ms pytestmark = pytest.mark.db_test @@ -116,7 +116,7 @@ def test_get_import_error( return expected_json = { "import_error_id": import_error_id, - "timestamp": datetime_zulu_format_without_ms(expected_body["timestamp"]), + "timestamp": from_datetime_to_zulu_without_ms(expected_body["timestamp"]), "filename": expected_body["filename"], "stack_trace": expected_body["stack_trace"], } diff --git a/tests/api_fastapi/core_api/routes/public/test_job.py b/tests/api_fastapi/core_api/routes/public/test_job.py index 7a085187e33e0..c4c78836819df 100644 --- a/tests/api_fastapi/core_api/routes/public/test_job.py +++ b/tests/api_fastapi/core_api/routes/public/test_job.py @@ -26,7 +26,7 @@ from airflow.utils.state import JobState, State from tests_common.test_utils.db import clear_db_jobs -from tests_common.test_utils.format_datetime import datetime_zulu_format +from tests_common.test_utils.format_datetime import from_datetime_to_zulu pytestmark = pytest.mark.db_test @@ -155,9 +155,9 @@ def test_get_jobs( "dag_id": None, "state": "running", "job_type": "SchedulerJob", - "start_date": datetime_zulu_format(self.scheduler_jobs[idx].start_date), + "start_date": from_datetime_to_zulu(self.scheduler_jobs[idx].start_date), "end_date": None, - "latest_heartbeat": datetime_zulu_format(self.scheduler_jobs[idx].latest_heartbeat), + "latest_heartbeat": from_datetime_to_zulu(self.scheduler_jobs[idx].latest_heartbeat), "executor_class": None, "hostname": self.scheduler_jobs[idx].hostname, "unixname": self.scheduler_jobs[idx].unixname, diff --git a/tests_common/test_utils/format_datetime.py b/tests_common/test_utils/format_datetime.py index be26d4351e52d..04f0120ed36b8 100644 --- a/tests_common/test_utils/format_datetime.py +++ b/tests_common/test_utils/format_datetime.py @@ -19,11 +19,11 @@ from datetime import datetime -def datetime_zulu_format(dt: datetime) -> str: +def from_datetime_to_zulu(dt: datetime) -> str: """Format a datetime object to a string in Zulu time.""" return dt.strftime("%Y-%m-%dT%H:%M:%S.%fZ") -def datetime_zulu_format_without_ms(dt: datetime) -> str: +def from_datetime_to_zulu_without_ms(dt: datetime) -> str: """Format a datetime object to a string in Zulu time without milliseconds.""" return dt.strftime("%Y-%m-%dT%H:%M:%SZ")