Skip to content

Commit

Permalink
AIP-84 Refactor test cases with datetime_zulu_format (apache#44323)
Browse files Browse the repository at this point in the history 
* Refactor test_dag_run

* Refactor test_assets

* Rename datetime_zulu_format to from_datetime_to_zulu

* Fix TestListDagRunsBatch
  • Loading branch information
@jason810496
jason810496 authored Nov 27, 2024
1 parent 1de5a96 commit 55c2545
Show file tree
Hide file tree
Showing 6 changed files with 80 additions and 86 deletions.
95 changes: 46 additions & 49 deletions tests/api_fastapi/core_api/routes/public/test_assets.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import urllib
from collections.abc import Generator
from datetime import datetime
from unittest import mock

import pytest
Expand All @@ -39,21 +40,23 @@

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,
name=f"simple{i}",
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)
]
Expand All @@ -62,16 +65,15 @@ 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,
name=f"sensitive{i}",
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)
]
Expand All @@ -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,
Expand All @@ -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)
]
Expand All @@ -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,
Expand All @@ -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)
]
Expand All @@ -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()

Expand All @@ -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
Expand Down Expand Up @@ -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": [
{
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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}",
Expand Down Expand Up @@ -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}",
Expand Down Expand Up @@ -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",
}
Expand Down Expand Up @@ -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):
Expand All @@ -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),
}


Expand All @@ -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",
}
Expand Down
51 changes: 24 additions & 27 deletions tests/api_fastapi/core_api/routes/public/test_dag_run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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": {},
Expand All @@ -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",
}
],
Expand Down
Loading

0 comments on commit 55c2545

Please sign in to comment.