Fix serialization tests

apache · Aug 26, 2024 · c1b2265 · c1b2265
1 parent 4686e04
commit c1b2265
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/tests/api_connexion/schemas/test_dag_schema.py b/tests/api_connexion/schemas/test_dag_schema.py
@@ -245,6 +245,7 @@ def test_serialize_test_dag_with_dataset_schedule_detail_schema(url_safe_seriali
         "start_date": "2020-06-19T00:00:00+00:00",
         "tags": [{"name": "example1"}, {"name": "example2"}],
         "template_searchpath": None,
+        "timetable_summary": "Dataset",
         "timezone": UTC_JSON_REPR,
         "max_active_runs": 16,
         "max_consecutive_failed_dag_runs": 0,

diff --git a/tests/serialization/test_dag_serialization.py b/tests/serialization/test_dag_serialization.py
@@ -142,6 +142,12 @@ def detect_task_dependencies(task: Operator) -> DagDependency | None:  # type: i
             },
         },
         "start_date": 1564617600.0,
+        "timetable": {
+            "__type": "airflow.timetables.interval.DeltaDataIntervalTimetable",
+            "__var": {
+                "delta": 86400.0,
+            },
+        },
         "_task_group": {
             "_group_id": None,
             "prefix_group_id": True,

diff --git a/tests/serialization/test_pydantic_models.py b/tests/serialization/test_pydantic_models.py
@@ -27,7 +27,7 @@
 from airflow.jobs.job import Job
 from airflow.jobs.local_task_job_runner import LocalTaskJobRunner
 from airflow.models import MappedOperator
-from airflow.models.dag import DAG, DagModel
+from airflow.models.dag import DAG, DagModel, create_timetable
 from airflow.models.dataset import (
     DagScheduleDatasetReference,
     DatasetEvent,
@@ -175,10 +175,12 @@ def test_serializing_pydantic_dagrun(session, create_task_instance):
     ],
 )
 def test_serializing_pydantic_dagmodel(schedule):
+    timetable = create_timetable(schedule, timezone.utc)
     dag_model = DagModel(
         dag_id="test-dag",
         fileloc="/tmp/dag_1.py",
-        timetable_summary=schedule,
+        timetable_summary=timetable.summary,
+        timetable_description=timetable.description,
         is_active=True,
         is_paused=False,
     )
@@ -189,7 +191,8 @@ def test_serializing_pydantic_dagmodel(schedule):
     deserialized_model = DagModelPydantic.model_validate_json(json_string)
     assert deserialized_model.dag_id == "test-dag"
     assert deserialized_model.fileloc == "/tmp/dag_1.py"
-    assert deserialized_model.timetable_summary == schedule
+    assert deserialized_model.timetable_summary == timetable.summary
+    assert deserialized_model.timetable_description == timetable.description
     assert deserialized_model.is_active is True
     assert deserialized_model.is_paused is False