Skip to content

Commit 2e195a6

Browse files
jason810496ArshiaZr
jason810496
authored and
ArshiaZr
committed
AIP-84 Refactor test cases with datetime_zulu_format (apache#44323)
* Refactor test_dag_run * Refactor test_assets * Rename datetime_zulu_format to from_datetime_to_zulu * Fix TestListDagRunsBatch
1 parent 0ea46d4 commit 2e195a6

6 files changed

Lines changed: 80 additions & 86 deletions

File tree

‎tests/api_fastapi/core_api/routes/public/test_assets.py‎

Lines changed: 46 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import urllib
2020
from collections.abc import Generator
21+
from datetime import datetime
2122
from unittest import mock
2223

2324
import pytest
@@ -39,21 +40,23 @@
3940

4041
from tests_common.test_utils.asserts import assert_queries_count
4142
from tests_common.test_utils.db import clear_db_assets, clear_db_runs
43+
from tests_common.test_utils.format_datetime import from_datetime_to_zulu_without_ms
44+
45+
DEFAULT_DATE = datetime(2020, 6, 11, 18, 0, 0, tzinfo=timezone.utc)
4246

4347
pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode]
4448

4549

4650
def _create_assets(session, num: int = 2) -> None:
47-
default_time = "2020-06-11T18:00:00+00:00"
4851
assets = [
4952
AssetModel(
5053
id=i,
5154
name=f"simple{i}",
5255
uri=f"s3://bucket/key/{i}",
5356
group="asset",
5457
extra={"foo": "bar"},
55-
created_at=timezone.parse(default_time),
56-
updated_at=timezone.parse(default_time),
58+
created_at=DEFAULT_DATE,
59+
updated_at=DEFAULT_DATE,
5760
)
5861
for i in range(1, 1 + num)
5962
]
@@ -62,16 +65,15 @@ def _create_assets(session, num: int = 2) -> None:
6265

6366

6467
def _create_assets_with_sensitive_extra(session, num: int = 2) -> None:
65-
default_time = "2020-06-11T18:00:00+00:00"
6668
assets = [
6769
AssetModel(
6870
id=i,
6971
name=f"sensitive{i}",
7072
uri=f"s3://bucket/key/{i}",
7173
group="asset",
7274
extra={"password": "bar"},
73-
created_at=timezone.parse(default_time),
74-
updated_at=timezone.parse(default_time),
75+
created_at=DEFAULT_DATE,
76+
updated_at=DEFAULT_DATE,
7577
)
7678
for i in range(1, 1 + num)
7779
]
@@ -85,7 +87,6 @@ def _create_provided_asset(session, asset: AssetModel) -> None:
8587

8688

8789
def _create_assets_events(session, num: int = 2) -> None:
88-
default_time = "2020-06-11T18:00:00+00:00"
8990
assets_events = [
9091
AssetEvent(
9192
id=i,
@@ -94,7 +95,7 @@ def _create_assets_events(session, num: int = 2) -> None:
9495
source_task_id="source_task_id",
9596
source_dag_id="source_dag_id",
9697
source_run_id=f"source_run_id_{i}",
97-
timestamp=timezone.parse(default_time),
98+
timestamp=DEFAULT_DATE,
9899
)
99100
for i in range(1, 1 + num)
100101
]
@@ -103,7 +104,6 @@ def _create_assets_events(session, num: int = 2) -> None:
103104

104105

105106
def _create_assets_events_with_sensitive_extra(session, num: int = 2) -> None:
106-
default_time = "2020-06-11T18:00:00+00:00"
107107
assets_events = [
108108
AssetEvent(
109109
id=i,
@@ -112,7 +112,7 @@ def _create_assets_events_with_sensitive_extra(session, num: int = 2) -> None:
112112
source_task_id="source_task_id",
113113
source_dag_id="source_dag_id",
114114
source_run_id=f"source_run_id_{i}",
115-
timestamp=timezone.parse(default_time),
115+
timestamp=DEFAULT_DATE,
116116
)
117117
for i in range(1, 1 + num)
118118
]
@@ -126,22 +126,21 @@ def _create_provided_asset_event(session, asset_event: AssetEvent) -> None:
126126

127127

128128
def _create_dag_run(session, num: int = 2):
129-
default_time = "2020-06-11T18:00:00+00:00"
130129
dag_runs = [
131130
DagRun(
132131
dag_id="source_dag_id",
133132
run_id=f"source_run_id_{i}",
134133
run_type=DagRunType.MANUAL,
135-
logical_date=timezone.parse(default_time),
136-
start_date=timezone.parse(default_time),
137-
data_interval=(timezone.parse(default_time), timezone.parse(default_time)),
134+
logical_date=DEFAULT_DATE,
135+
start_date=DEFAULT_DATE,
136+
data_interval=(DEFAULT_DATE, DEFAULT_DATE),
138137
external_trigger=True,
139138
state=DagRunState.SUCCESS,
140139
)
141140
for i in range(1, 1 + num)
142141
]
143142
for dag_run in dag_runs:
144-
dag_run.end_date = timezone.parse(default_time)
143+
dag_run.end_date = DEFAULT_DATE
145144
session.add_all(dag_runs)
146145
session.commit()
147146

@@ -156,11 +155,9 @@ def _create_asset_dag_run(session, num: int = 2):
156155

157156

158157
class TestAssets:
159-
default_time = "2020-06-11T18:00:00+00:00"
160-
161158
@pytest.fixture
162159
def time_freezer(self) -> Generator:
163-
freezer = time_machine.travel(self.default_time, tick=False)
160+
freezer = time_machine.travel(DEFAULT_DATE, tick=False)
164161
freezer.start()
165162

166163
yield
@@ -218,7 +215,7 @@ def test_should_respond_200(self, test_client, session):
218215
response = test_client.get("/public/assets")
219216
assert response.status_code == 200
220217
response_data = response.json()
221-
tz_datetime_format = self.default_time.replace("+00:00", "Z")
218+
tz_datetime_format = from_datetime_to_zulu_without_ms(DEFAULT_DATE)
222219
assert response_data == {
223220
"assets": [
224221
{
@@ -402,15 +399,15 @@ def test_should_respond_200(self, test_client, session):
402399
{
403400
"run_id": "source_run_id_1",
404401
"dag_id": "source_dag_id",
405-
"logical_date": "2020-06-11T18:00:00Z",
406-
"start_date": "2020-06-11T18:00:00Z",
407-
"end_date": "2020-06-11T18:00:00Z",
402+
"logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
403+
"start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
404+
"end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
408405
"state": "success",
409-
"data_interval_start": "2020-06-11T18:00:00Z",
410-
"data_interval_end": "2020-06-11T18:00:00Z",
406+
"data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
407+
"data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
411408
}
412409
],
413-
"timestamp": "2020-06-11T18:00:00Z",
410+
"timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
414411
},
415412
{
416413
"id": 2,
@@ -425,15 +422,15 @@ def test_should_respond_200(self, test_client, session):
425422
{
426423
"run_id": "source_run_id_2",
427424
"dag_id": "source_dag_id",
428-
"logical_date": "2020-06-11T18:00:00Z",
429-
"start_date": "2020-06-11T18:00:00Z",
430-
"end_date": "2020-06-11T18:00:00Z",
425+
"logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
426+
"start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
427+
"end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
431428
"state": "success",
432-
"data_interval_start": "2020-06-11T18:00:00Z",
433-
"data_interval_end": "2020-06-11T18:00:00Z",
429+
"data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
430+
"data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
434431
}
435432
],
436-
"timestamp": "2020-06-11T18:00:00Z",
433+
"timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
437434
},
438435
],
439436
"total_entries": 2,
@@ -514,15 +511,15 @@ def test_should_mask_sensitive_extra(self, test_client, session):
514511
{
515512
"run_id": "source_run_id_1",
516513
"dag_id": "source_dag_id",
517-
"logical_date": "2020-06-11T18:00:00Z",
518-
"start_date": "2020-06-11T18:00:00Z",
519-
"end_date": "2020-06-11T18:00:00Z",
514+
"logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
515+
"start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
516+
"end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
520517
"state": "success",
521-
"data_interval_start": "2020-06-11T18:00:00Z",
522-
"data_interval_end": "2020-06-11T18:00:00Z",
518+
"data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
519+
"data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
523520
}
524521
],
525-
"timestamp": "2020-06-11T18:00:00Z",
522+
"timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
526523
},
527524
{
528525
"id": 2,
@@ -537,15 +534,15 @@ def test_should_mask_sensitive_extra(self, test_client, session):
537534
{
538535
"run_id": "source_run_id_2",
539536
"dag_id": "source_dag_id",
540-
"logical_date": "2020-06-11T18:00:00Z",
541-
"start_date": "2020-06-11T18:00:00Z",
542-
"end_date": "2020-06-11T18:00:00Z",
537+
"logical_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
538+
"start_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
539+
"end_date": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
543540
"state": "success",
544-
"data_interval_start": "2020-06-11T18:00:00Z",
545-
"data_interval_end": "2020-06-11T18:00:00Z",
541+
"data_interval_start": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
542+
"data_interval_end": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
546543
}
547544
],
548-
"timestamp": "2020-06-11T18:00:00Z",
545+
"timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
549546
},
550547
],
551548
"total_entries": 2,
@@ -566,7 +563,7 @@ class TestGetAssetEndpoint(TestAssets):
566563
def test_should_respond_200(self, test_client, url, session):
567564
self.create_assets(num=1)
568565
assert session.query(AssetModel).count() == 1
569-
tz_datetime_format = self.default_time.replace("+00:00", "Z")
566+
tz_datetime_format = from_datetime_to_zulu_without_ms(DEFAULT_DATE)
570567
with assert_queries_count(6):
571568
response = test_client.get(
572569
f"/public/assets/{url}",
@@ -596,7 +593,7 @@ def test_should_respond_404(self, test_client):
596593
@pytest.mark.enable_redact
597594
def test_should_mask_sensitive_extra(self, test_client, session):
598595
self.create_assets_with_sensitive_extra()
599-
tz_datetime_format = self.default_time.replace("+00:00", "Z")
596+
tz_datetime_format = from_datetime_to_zulu_without_ms(DEFAULT_DATE)
600597
uri = "s3://bucket/key/1"
601598
response = test_client.get(
602599
f"/public/assets/{uri}",
@@ -641,7 +638,7 @@ def test_should_respond_200(self, test_client, session, create_dummy_dag):
641638
assert response.json() == {
642639
"queued_events": [
643640
{
644-
"created_at": self.default_time.replace("+00:00", "Z"),
641+
"created_at": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
645642
"uri": "s3://bucket/key/1",
646643
"dag_id": "dag",
647644
}
@@ -721,7 +718,7 @@ def test_should_respond_200(self, test_client, session):
721718
"source_run_id": None,
722719
"source_map_index": -1,
723720
"created_dagruns": [],
724-
"timestamp": self.default_time.replace("+00:00", "Z"),
721+
"timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
725722
}
726723

727724
def test_invalid_attr_not_allowed(self, test_client, session):
@@ -748,7 +745,7 @@ def test_should_mask_sensitive_extra(self, test_client, session):
748745
"source_run_id": None,
749746
"source_map_index": -1,
750747
"created_dagruns": [],
751-
"timestamp": self.default_time.replace("+00:00", "Z"),
748+
"timestamp": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
752749
}
753750

754751

@@ -769,7 +766,7 @@ def test_should_respond_200(self, test_client, session, create_dummy_dag):
769766
assert response.json() == {
770767
"queued_events": [
771768
{
772-
"created_at": self.default_time.replace("+00:00", "Z"),
769+
"created_at": from_datetime_to_zulu_without_ms(DEFAULT_DATE),
773770
"uri": "s3://bucket/key/1",
774771
"dag_id": "dag",
775772
}

‎tests/api_fastapi/core_api/routes/public/test_dag_run.py‎

Lines changed: 24 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
clear_db_runs,
4040
clear_db_serialized_dags,
4141
)
42+
from tests_common.test_utils.format_datetime import from_datetime_to_zulu, from_datetime_to_zulu_without_ms
4243

4344
pytestmark = pytest.mark.db_test
4445

@@ -192,22 +193,20 @@ def test_get_dag_run_not_found(self, test_client):
192193

193194

194195
class TestGetDagRuns:
195-
@staticmethod
196-
def parse_datetime(datetime_str):
197-
return datetime_str.isoformat().replace("+00:00", "Z") if datetime_str else None
198-
199196
@staticmethod
200197
def get_dag_run_dict(run: DagRun):
201198
return {
202199
"dag_run_id": run.run_id,
203200
"dag_id": run.dag_id,
204-
"logical_date": TestGetDagRuns.parse_datetime(run.logical_date),
205-
"queued_at": TestGetDagRuns.parse_datetime(run.queued_at),
206-
"start_date": TestGetDagRuns.parse_datetime(run.start_date),
207-
"end_date": TestGetDagRuns.parse_datetime(run.end_date),
208-
"data_interval_start": TestGetDagRuns.parse_datetime(run.data_interval_start),
209-
"data_interval_end": TestGetDagRuns.parse_datetime(run.data_interval_end),
210-
"last_scheduling_decision": TestGetDagRuns.parse_datetime(run.last_scheduling_decision),
201+
"logical_date": from_datetime_to_zulu_without_ms(run.logical_date),
202+
"queued_at": from_datetime_to_zulu(run.queued_at) if run.queued_at else None,
203+
"start_date": from_datetime_to_zulu_without_ms(run.start_date),
204+
"end_date": from_datetime_to_zulu(run.end_date),
205+
"data_interval_start": from_datetime_to_zulu_without_ms(run.data_interval_start),
206+
"data_interval_end": from_datetime_to_zulu_without_ms(run.data_interval_end),
207+
"last_scheduling_decision": from_datetime_to_zulu(run.last_scheduling_decision)
208+
if run.last_scheduling_decision
209+
else None,
211210
"run_type": run.run_type,
212211
"state": run.state,
213212
"external_trigger": run.external_trigger,
@@ -482,22 +481,20 @@ def test_invalid_state(self, test_client):
482481

483482

484483
class TestListDagRunsBatch:
485-
@staticmethod
486-
def parse_datetime(datetime_str):
487-
return datetime_str.isoformat().replace("+00:00", "Z") if datetime_str else None
488-
489484
@staticmethod
490485
def get_dag_run_dict(run: DagRun):
491486
return {
492487
"dag_run_id": run.run_id,
493488
"dag_id": run.dag_id,
494-
"logical_date": TestGetDagRuns.parse_datetime(run.logical_date),
495-
"queued_at": TestGetDagRuns.parse_datetime(run.queued_at),
496-
"start_date": TestGetDagRuns.parse_datetime(run.start_date),
497-
"end_date": TestGetDagRuns.parse_datetime(run.end_date),
498-
"data_interval_start": TestGetDagRuns.parse_datetime(run.data_interval_start),
499-
"data_interval_end": TestGetDagRuns.parse_datetime(run.data_interval_end),
500-
"last_scheduling_decision": TestGetDagRuns.parse_datetime(run.last_scheduling_decision),
489+
"logical_date": from_datetime_to_zulu_without_ms(run.logical_date),
490+
"queued_at": from_datetime_to_zulu_without_ms(run.queued_at) if run.queued_at else None,
491+
"start_date": from_datetime_to_zulu_without_ms(run.start_date),
492+
"end_date": from_datetime_to_zulu(run.end_date),
493+
"data_interval_start": from_datetime_to_zulu_without_ms(run.data_interval_start),
494+
"data_interval_end": from_datetime_to_zulu_without_ms(run.data_interval_end),
495+
"last_scheduling_decision": from_datetime_to_zulu_without_ms(run.last_scheduling_decision)
496+
if run.last_scheduling_decision
497+
else None,
501498
"run_type": run.run_type,
502499
"state": run.state,
503500
"external_trigger": run.external_trigger,
@@ -976,7 +973,7 @@ def test_should_respond_200(self, test_client, dag_maker, session):
976973
expected_response = {
977974
"asset_events": [
978975
{
979-
"timestamp": event.timestamp.isoformat().replace("+00:00", "Z"),
976+
"timestamp": from_datetime_to_zulu(event.timestamp),
980977
"asset_id": asset1_id,
981978
"uri": asset1.uri,
982979
"extra": {},
@@ -989,11 +986,11 @@ def test_should_respond_200(self, test_client, dag_maker, session):
989986
{
990987
"dag_id": "TEST_DAG_ID",
991988
"run_id": "TEST_DAG_RUN_ID",
992-
"data_interval_end": dr.data_interval_end.isoformat().replace("+00:00", "Z"),
993-
"data_interval_start": dr.data_interval_start.isoformat().replace("+00:00", "Z"),
989+
"data_interval_end": from_datetime_to_zulu_without_ms(dr.data_interval_end),
990+
"data_interval_start": from_datetime_to_zulu_without_ms(dr.data_interval_start),
994991
"end_date": None,
995-
"logical_date": dr.logical_date.isoformat().replace("+00:00", "Z"),
996-
"start_date": dr.start_date.isoformat().replace("+00:00", "Z"),
992+
"logical_date": from_datetime_to_zulu_without_ms(dr.logical_date),
993+
"start_date": from_datetime_to_zulu_without_ms(dr.start_date),
997994
"state": "running",
998995
}
999996
],

0 commit comments

Comments
 (0)