Back to Repositories

Testing Custom Timetable Serialization in Apache Airflow

This test suite validates the timetable functionality in Apache Airflow, focusing on custom serialization and different timetable implementations including cron and delta-based scheduling.

Test Coverage Overview

The test suite provides comprehensive coverage of Airflow’s timetable components, including:

CronDataIntervalTimetable implementation and timezone handling
DeltaDataIntervalTimetable functionality
Custom timetable serialization and deserialization
Data interval inference and scheduling logic

Implementation Analysis

The testing approach utilizes Python’s class-based structure to validate timetable behaviors. The implementation focuses on two main patterns:

Factory functions for creating standard timetable instances
Custom timetable class with serialization capabilities for testing complex scenarios

Technical Details

Key technical components include:

Airflow’s base Timetable class implementation
Integration with settings.TIMEZONE configuration
DataInterval handling for scheduling
Serialization/deserialization methods for custom timetables

Best Practices Demonstrated

The test suite exemplifies several testing best practices:

Clear separation of concerns between different timetable types
Proper inheritance and interface implementation
Comprehensive equality comparison for custom objects
Explicit handling of edge cases in manual data interval inference

apache/airflow

tests_common/test_utils/timetables.py

            
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

from airflow import settings
from airflow.timetables.base import DataInterval, Timetable
from airflow.timetables.interval import CronDataIntervalTimetable, DeltaDataIntervalTimetable


def cron_timetable(expr: str) -> CronDataIntervalTimetable:
    return CronDataIntervalTimetable(expr, settings.TIMEZONE)


def delta_timetable(delta) -> DeltaDataIntervalTimetable:
    return DeltaDataIntervalTimetable(delta)


class CustomSerializationTimetable(Timetable):
    """Custom timetable for testing serialization."""

    def __init__(self, value: str):
        self.value = value

    @classmethod
    def deserialize(cls, data):
        return cls(data["value"])

    def __eq__(self, other) -> bool:
        """Only for testing purposes."""
        if not isinstance(other, CustomSerializationTimetable):
            return False
        return self.value == other.value

    def serialize(self):
        return {"value": self.value}

    @property
    def summary(self):
        return f"{type(self).__name__}({self.value!r})"

    def infer_manual_data_interval(self, *, run_after):
        raise DataInterval.exact(run_after)

    def next_dagrun_info(self, *, last_automated_data_interval, restriction):
        return None