test_serialized_objects.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

import json
import math
import sys
from collections.abc import Iterator
from datetime import datetime, timedelta
from typing import TYPE_CHECKING

import pendulum
import pytest
from dateutil import relativedelta
from kubernetes.client import models as k8s
from pendulum.tz.timezone import FixedTimezone, Timezone
from uuid6 import uuid7

from airflow._shared.timezones import timezone
from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest
from airflow.exceptions import (
    AirflowException,
    AirflowFailException,
    AirflowRescheduleException,
    SerializationError,
    TaskDeferred,
)
from airflow.models.connection import Connection
from airflow.models.dag import DAG
from airflow.models.dagrun import DagRun
from airflow.models.taskinstance import TaskInstance
from airflow.models.xcom_arg import XComArg
from airflow.partition_mappers.identity import IdentityMapper as CoreIdentityMapper
from airflow.partition_mappers.temporal import (
    DailyMapper as CoreDailyMapper,
    HourlyMapper as CoureHourlyMapper,
    MonthlyMapper as CoreMonthlyMapper,
    QuarterlyMapper as CoreQuarterlyMapper,
    WeeklyMapper as CoreWeeklyMapper,
    YearlyMapper as CoreYearlyMapper,
)
from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator
from airflow.providers.standard.triggers.file import FileDeleteTrigger
from airflow.sdk import (
    BaseOperator,
    DailyMapper,
    HourlyMapper,
    IdentityMapper,
    MonthlyMapper,
    QuarterlyMapper,
    WeeklyMapper,
    YearlyMapper,
)
from airflow.sdk.definitions.asset import (
    Asset,
    AssetAlias,
    AssetAliasEvent,
    AssetUniqueKey,
    AssetWatcher,
    BaseAsset,
)
from airflow.sdk.definitions.deadline import (
    AsyncCallback,
    DeadlineAlert,
    DeadlineReference,
)
from airflow.sdk.definitions.decorators import task
from airflow.sdk.definitions.operator_resources import Resources
from airflow.sdk.definitions.param import Param
from airflow.sdk.definitions.taskgroup import TaskGroup
from airflow.sdk.execution_time.context import OutletEventAccessor, OutletEventAccessors
from airflow.serialization.definitions.assets import (
    SerializedAsset,
    SerializedAssetAlias,
    SerializedAssetAll,
    SerializedAssetAny,
    SerializedAssetBase,
    SerializedAssetRef,
)
from airflow.serialization.definitions.deadline import DeadlineAlertFields, SerializedDeadlineAlert
from airflow.serialization.encoders import ensure_serialized_asset, ensure_serialized_deadline_alert
from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding
from airflow.serialization.helpers import PartitionMapperNotFound
from airflow.serialization.serialized_objects import (
    BaseSerialization,
    DagSerialization,
    LazyDeserializedDAG,
    _has_kubernetes,
    create_scheduler_operator,
)
from airflow.triggers.base import BaseTrigger
from airflow.utils.db import LazySelectSequence
from airflow.utils.state import DagRunState, State
from airflow.utils.types import DagRunType

from unit.models import DEFAULT_DATE

if TYPE_CHECKING:
    from pydantic.types import JsonValue

DAG_ID = "dag_id_1"

TEST_CALLBACK_PATH = f"{__name__}.empty_callback_for_deadline"
TEST_CALLBACK_KWARGS = {"arg1": "value1"}

REFERENCE_TYPES = [
    pytest.param(DeadlineReference.DAGRUN_LOGICAL_DATE, id="logical_date"),
    pytest.param(DeadlineReference.DAGRUN_QUEUED_AT, id="queued_at"),
    pytest.param(DeadlineReference.FIXED_DATETIME(DEFAULT_DATE), id="fixed_deadline"),
]


async def empty_callback_for_deadline():
    """Used in a number of tests to confirm that Deadlines and DeadlineAlerts function correctly."""
    pass


def test_recursive_serialize_calls_must_forward_kwargs():
    """Any time we recurse cls.serialize, we must forward all kwargs."""
    import ast
    from pathlib import Path

    import airflow.serialization

    valid_recursive_call_count = 0
    skipped_recursive_calls = 0  # when another serialize method called
    file = Path(airflow.serialization.__path__[0]) / "serialized_objects.py"
    content = file.read_text()
    tree = ast.parse(content)

    class_def = None
    for stmt in ast.walk(tree):
        if isinstance(stmt, ast.ClassDef) and stmt.name == "BaseSerialization":
            class_def = stmt

    method_def = None
    for elem in ast.walk(class_def):
        if isinstance(elem, ast.FunctionDef) and elem.name == "serialize":
            method_def = elem
            break
    kwonly_args = [x.arg for x in method_def.args.kwonlyargs]
    for elem in ast.walk(method_def):
        if isinstance(elem, ast.Call) and getattr(elem.func, "attr", "") == "serialize":
            if not elem.func.value.id == "cls":
                skipped_recursive_calls += 1
                break
            kwargs = {y.arg: y.value for y in elem.keywords}
            for name in kwonly_args:
                if name not in kwargs or getattr(kwargs[name], "id", "") != name:
                    ref = f"{file}:{elem.lineno}"
                    message = (
                        f"Error at {ref}; recursive calls to `cls.serialize` "
                        f"must forward the `{name}` argument"
                    )
                    raise Exception(message)
                valid_recursive_call_count += 1
    print(f"validated calls: {valid_recursive_call_count}")
    assert valid_recursive_call_count > 0
    assert skipped_recursive_calls == 1


def test_strict_mode():
    """If strict=True, serialization should fail when object is not JSON serializable."""

    class Test:
        a = 1

    from airflow.serialization.serialized_objects import BaseSerialization

    obj = [[[Test()]]]  # nested to verify recursive behavior
    BaseSerialization.serialize(obj)  # does not raise
    with pytest.raises(SerializationError, match="Encountered unexpected type"):
        BaseSerialization.serialize(obj, strict=True)  # now raises


def test_prevent_re_serialization_of_serialized_operators():
    """SerializedBaseOperator should not be re-serializable."""
    from airflow.serialization.definitions.baseoperator import SerializedBaseOperator

    serialized_op = SerializedBaseOperator(task_id="test_task")

    with pytest.raises(SerializationError, match="Encountered unexpected type"):
        BaseSerialization.serialize(serialized_op, strict=True)


def test_validate_schema():
    from airflow.serialization.serialized_objects import BaseSerialization

    with pytest.raises(AirflowException, match="BaseSerialization is not set"):
        BaseSerialization.validate_schema({"any": "thing"})

    BaseSerialization._json_schema = object()
    with pytest.raises(TypeError, match="Invalid type: Only dict and str are supported"):
        BaseSerialization.validate_schema(123)


def test_serde_validate_schema_valid_json():
    from airflow.serialization.serialized_objects import BaseSerialization

    class Test:
        def validate(self, obj):
            self.obj = obj

    t = Test()
    BaseSerialization._json_schema = t
    BaseSerialization.validate_schema('{"foo": "bar"}')
    assert t.obj == {"foo": "bar"}


TI = TaskInstance(
    task=create_scheduler_operator(EmptyOperator(task_id="test-task")),
    run_id="fake_run",
    state=State.RUNNING,
    dag_version_id=uuid7(),
)

TI_WITH_START_DAY = TaskInstance(
    task=create_scheduler_operator(EmptyOperator(task_id="test-task")),
    run_id="fake_run",
    state=State.RUNNING,
    dag_version_id=uuid7(),
)
TI_WITH_START_DAY.start_date = timezone.datetime(2020, 1, 1, 0, 0, 0)

DAG_RUN = DagRun(
    dag_id="test_dag_id",
    run_id="test_dag_run_id",
    run_type=DagRunType.MANUAL,
    logical_date=timezone.utcnow(),
    start_date=timezone.utcnow(),
    state=DagRunState.SUCCESS,
)
DAG_RUN.id = 1


# we add the tasks out of order, to ensure they are deserialized in the correct order
DAG_WITH_TASKS = DAG(dag_id="test_dag", start_date=datetime.now())
EmptyOperator(task_id="task2", dag=DAG_WITH_TASKS)
EmptyOperator(task_id="task1", dag=DAG_WITH_TASKS)


def create_outlet_event_accessors(
    key: Asset | AssetAlias, extra: dict[str, JsonValue], asset_alias_events: list[AssetAliasEvent]
) -> OutletEventAccessors:
    o = OutletEventAccessors()
    o[key].extra = extra
    o[key].asset_alias_events = asset_alias_events

    return o


def equals(a, b) -> bool:
    return a == b


def equal_time(a: datetime, b: datetime) -> bool:
    return a.strftime("%s") == b.strftime("%s")


def equal_exception(a: AirflowException, b: AirflowException) -> bool:
    return a.__class__ == b.__class__ and str(a) == str(b)


def equal_outlet_event_accessors(a: OutletEventAccessors, b: OutletEventAccessors) -> bool:
    return a._dict.keys() == b._dict.keys() and all(
        equal_outlet_event_accessor(a._dict[key], b._dict[key]) for key in a._dict
    )


def equal_outlet_event_accessor(a: OutletEventAccessor, b: OutletEventAccessor) -> bool:
    return a.key == b.key and a.extra == b.extra and a.asset_alias_events == b.asset_alias_events


def equal_serialized_asset(a: SerializedAssetBase | BaseAsset, b: SerializedAssetBase | BaseAsset) -> bool:
    return ensure_serialized_asset(a) == ensure_serialized_asset(b)


def equal_serialized_deadline_alert(
    a: DeadlineAlert | SerializedDeadlineAlert,
    b: DeadlineAlert | SerializedDeadlineAlert,
) -> bool:
    """Compare deadline alerts for equality, normalizing to serialized form."""
    a_ser = ensure_serialized_deadline_alert(a)
    b_ser = ensure_serialized_deadline_alert(b)

    return a_ser == b_ser


class MockLazySelectSequence(LazySelectSequence):
    _data = ["a", "b", "c"]

    def __init__(self):
        super().__init__(None, None, session="MockSession")

    def __iter__(self) -> Iterator[str]:
        return iter(self._data)

    def __len__(self) -> int:
        return len(self._data)


@pytest.mark.parametrize(
    ("input", "encoded_type", "cmp_func"),
    [
        ("test_str", None, equals),
        (1, None, equals),
        (math.nan, None, lambda a, b: b == "nan"),
        (math.inf, None, lambda a, b: b == "inf"),
        (-math.inf, None, lambda a, b: b == "-inf"),
        (timezone.utcnow(), DAT.DATETIME, equal_time),
        (timedelta(minutes=2), DAT.TIMEDELTA, equals),
        (Timezone("UTC"), DAT.TIMEZONE, lambda a, b: a.name == b.name),
        (
            relativedelta.relativedelta(hours=+1),
            DAT.RELATIVEDELTA,
            lambda a, b: a.hours == b.hours,
        ),
        ({"test": "dict", "test-1": 1}, None, equals),
        (["array_item", 2], None, equals),
        (("tuple_item", 3), DAT.TUPLE, equals),
        (set(["set_item", 3]), DAT.SET, equals),
        (
            k8s.V1Pod(
                metadata=k8s.V1ObjectMeta(
                    name="test",
                    annotations={"test": "annotation"},
                    creation_timestamp=timezone.utcnow(),
                )
            ),
            DAT.POD,
            equals,
        ),
        (
            DAG(
                "fake-dag",
                schedule="*/10 * * * *",
                default_args={"depends_on_past": True},
                start_date=timezone.utcnow(),
                catchup=False,
            ),
            DAT.DAG,
            lambda a, b: a.dag_id == b.dag_id and equal_time(a.start_date, b.start_date),
        ),
        (Resources(cpus=0.1, ram=2048), None, None),
        (EmptyOperator(task_id="test-task"), None, None),
        (
            TaskGroup(
                group_id="test-group",
                dag=DAG(dag_id="test_dag", start_date=datetime.now()),
            ),
            None,
            None,
        ),
        (
            Param("test", "desc"),
            DAT.PARAM,
            lambda a, b: a.value == b.value and a.description == b.description,
        ),
        (
            XComArg(
                operator=PythonOperator(
                    python_callable=int,
                    task_id="test_xcom_op",
                    do_xcom_push=True,
                )
            ),
            DAT.XCOM_REF,
            None,
        ),
        (
            MockLazySelectSequence(),
            None,
            lambda a, b: len(a) == len(b) and isinstance(b, list),
        ),
        (Asset(uri="test://asset1", name="test"), DAT.ASSET, equal_serialized_asset),
        (
            Asset(
                uri="test://asset1",
                name="test",
                watchers=[AssetWatcher(name="test", trigger=FileDeleteTrigger(filepath="/tmp"))],
            ),
            DAT.ASSET,
            equal_serialized_asset,
        ),
        (
            Connection(conn_id="TEST_ID", uri="mysql://"),
            DAT.CONNECTION,
            lambda a, b: a.get_uri() == b.get_uri(),
        ),
        (
            TaskCallbackRequest(
                filepath="filepath",
                ti=TI,
                bundle_name="testing",
                bundle_version=None,
            ),
            DAT.TASK_CALLBACK_REQUEST,
            lambda a, b: a.ti == b.ti,
        ),
        (
            DagCallbackRequest(
                filepath="filepath",
                dag_id="fake_dag",
                run_id="fake_run",
                bundle_name="testing",
                bundle_version=None,
            ),
            DAT.DAG_CALLBACK_REQUEST,
            lambda a, b: a.dag_id == b.dag_id,
        ),
        (Asset.ref(name="test"), DAT.ASSET_REF, lambda a, b: a.name == b.name),
        (
            DeadlineAlert(
                reference=DeadlineReference.DAGRUN_LOGICAL_DATE,
                interval=timedelta(),
                callback=AsyncCallback("fake_callable"),
            ),
            None,
            None,
        ),
        (
            create_outlet_event_accessors(
                Asset(uri="test", name="test", group="test-group"), {"key": "value"}, []
            ),
            DAT.ASSET_EVENT_ACCESSORS,
            equal_outlet_event_accessors,
        ),
        (
            create_outlet_event_accessors(
                AssetAlias(name="test_alias", group="test-alias-group"),
                {"key": "value"},
                [
                    AssetAliasEvent(
                        source_alias_name="test_alias",
                        dest_asset_key=AssetUniqueKey(name="test_name", uri="test://asset-uri"),
                        dest_asset_extra={"extra": "from asset itself"},
                        extra={"extra": "from event"},
                    )
                ],
            ),
            DAT.ASSET_EVENT_ACCESSORS,
            equal_outlet_event_accessors,
        ),
        (
            AirflowException("test123 wohoo!"),
            DAT.AIRFLOW_EXC_SER,
            equal_exception,
        ),
        (
            AirflowFailException("uuups, failed :-("),
            DAT.AIRFLOW_EXC_SER,
            equal_exception,
        ),
        (
            DAG_WITH_TASKS,
            DAT.DAG,
            lambda _, b: list(b.task_group.children.keys()) == sorted(b.task_group.children.keys()),
        ),
        (
            DeadlineAlert(
                reference=DeadlineReference.DAGRUN_QUEUED_AT,
                interval=timedelta(hours=1),
                callback=AsyncCallback("valid.callback.path", kwargs={"arg1": "value1"}),
            ),
            DAT.DEADLINE_ALERT,
            equal_serialized_deadline_alert,
        ),
    ],
)
def test_serialize_deserialize(input, encoded_type, cmp_func):
    from airflow.serialization.serialized_objects import BaseSerialization

    serialized = BaseSerialization.serialize(input)  # does not raise
    json.dumps(serialized)  # does not raise
    if encoded_type is not None:
        assert serialized[Encoding.TYPE] == encoded_type
        assert serialized[Encoding.VAR] is not None
    if cmp_func is not None:
        deserialized = BaseSerialization.deserialize(serialized)
        assert cmp_func(input, deserialized)

    # Verify recursive behavior
    obj = [[input]]
    serialized = BaseSerialization.serialize(obj)  # does not raise
    # Verify the result is JSON-serializable
    json.dumps(serialized)  # does not raise


@pytest.mark.parametrize("reference", REFERENCE_TYPES)
def test_serialize_deserialize_deadline_alert(reference):
    public_deadline_alert_fields = {
        field.lower() for field in vars(DeadlineAlertFields) if not field.startswith("_")
    }
    original = DeadlineAlert(
        reference=reference,
        interval=timedelta(hours=1),
        callback=AsyncCallback(empty_callback_for_deadline, kwargs=TEST_CALLBACK_KWARGS),
    )

    # Use BaseSerialization like assets do
    serialized = BaseSerialization.serialize(original)
    assert serialized[Encoding.TYPE] == DAT.DEADLINE_ALERT
    assert set(serialized[Encoding.VAR].keys()) == public_deadline_alert_fields

    deserialized = BaseSerialization.deserialize(serialized)
    assert deserialized.reference.serialize_reference() == reference.serialize_reference()
    assert deserialized.interval == original.interval
    assert deserialized.callback == original.callback


@pytest.mark.parametrize(
    "conn_uri",
    [
        pytest.param("aws://", id="only-conn-type"),
        pytest.param(
            "postgres://username:password@ec2.compute.com:5432/the_database",
            id="all-non-extra",
        ),
        pytest.param(
            "///?__extra__=%7B%22foo%22%3A+%22bar%22%2C+%22answer%22%3A+42%2C+%22"
            "nullable%22%3A+null%2C+%22empty%22%3A+%22%22%2C+%22zero%22%3A+0%7D",
            id="extra",
        ),
    ],
)
def test_backcompat_deserialize_connection(conn_uri):
    """Test deserialize connection which serialised by previous serializer implementation."""
    from airflow.serialization.serialized_objects import BaseSerialization

    conn_obj = {
        Encoding.TYPE: DAT.CONNECTION,
        Encoding.VAR: {"conn_id": "TEST_ID", "uri": conn_uri},
    }
    deserialized = BaseSerialization.deserialize(conn_obj)
    assert deserialized.get_uri() == conn_uri


def test_ser_of_asset_event_accessor():
    # todo: (Airflow 3.0) we should force reserialization on upgrade
    d = OutletEventAccessors()
    d[
        Asset("hi")
    ].extra = "blah1"  # todo: this should maybe be forbidden?  i.e. can extra be any json or just dict?
    d[Asset(name="yo", uri="test://yo")].extra = {"this": "that", "the": "other"}
    ser = BaseSerialization.serialize(var=d)
    deser = BaseSerialization.deserialize(ser)
    assert deser[Asset(uri="hi", name="hi")].extra == "blah1"
    assert d[Asset(name="yo", uri="test://yo")].extra == {"this": "that", "the": "other"}


class MyTrigger(BaseTrigger):
    def __init__(self, hi):
        self.hi = hi

    def serialize(self):
        return "unit.serialization.test_serialized_objects.MyTrigger", {"hi": self.hi}

    async def run(self):
        yield


def test_roundtrip_exceptions():
    """
    This is for AIP-44 when we need to send certain non-error exceptions
    as part of an RPC call e.g. TaskDeferred or AirflowRescheduleException.
    """
    some_date = pendulum.now()
    resched_exc = AirflowRescheduleException(reschedule_date=some_date)
    ser = BaseSerialization.serialize(resched_exc)
    deser = BaseSerialization.deserialize(ser)
    assert isinstance(deser, AirflowRescheduleException)
    assert deser.reschedule_date == some_date
    del ser
    del deser
    exc = TaskDeferred(
        trigger=MyTrigger(hi="yo"),
        method_name="meth_name",
        kwargs={"have": "pie"},
        timeout=timedelta(seconds=30),
    )
    ser = BaseSerialization.serialize(exc)
    deser = BaseSerialization.deserialize(ser)
    assert deser.trigger.hi == "yo"
    assert deser.method_name == "meth_name"
    assert deser.kwargs == {"have": "pie"}
    assert deser.timeout == timedelta(seconds=30)


@pytest.mark.parametrize(
    "concurrency_parameter",
    [
        "max_active_tis_per_dag",
        "max_active_tis_per_dagrun",
    ],
)
@pytest.mark.db_test
def test_serialized_dag_has_task_concurrency_limits(dag_maker, concurrency_parameter):
    with dag_maker() as dag:
        BashOperator(task_id="task1", bash_command="echo 1", **{concurrency_parameter: 1})

    ser_dict = DagSerialization.to_dict(dag)
    lazy_serialized_dag = LazyDeserializedDAG(data=ser_dict)

    assert lazy_serialized_dag.has_task_concurrency_limits


@pytest.mark.parametrize(
    "concurrency_parameter",
    [
        "max_active_tis_per_dag",
        "max_active_tis_per_dagrun",
    ],
)
@pytest.mark.db_test
def test_serialized_dag_mapped_task_has_task_concurrency_limits(dag_maker, concurrency_parameter):
    with dag_maker() as dag:

        @task
        def my_task():
            return [1, 2, 3, 4, 5, 6, 7]

        @task(**{concurrency_parameter: 1})
        def map_me_but_slowly(a):
            pass

        map_me_but_slowly.expand(a=my_task())

    ser_dict = DagSerialization.to_dict(dag)
    lazy_serialized_dag = LazyDeserializedDAG(data=ser_dict)

    assert lazy_serialized_dag.has_task_concurrency_limits


def test_hash_property():
    from airflow.models.serialized_dag import SerializedDagModel

    data = {"dag": {"dag_id": "dag1"}}
    lazy_serialized_dag = LazyDeserializedDAG(data=data)
    assert lazy_serialized_dag.hash == SerializedDagModel.hash(data)


@pytest.mark.parametrize(
    ("payload", "expected_cls"),
    [
        pytest.param(
            {
                "__type": DAT.ASSET,
                "name": "test_asset",
                "uri": "test://asset-uri",
                "group": "test-group",
                "extra": {},
            },
            SerializedAsset,
            id="asset",
        ),
        pytest.param(
            {
                "__type": DAT.ASSET_ALL,
                "objects": [
                    {
                        "__type": DAT.ASSET,
                        "name": "x",
                        "uri": "test://x",
                        "group": "g",
                        "extra": {},
                    },
                    {
                        "__type": DAT.ASSET,
                        "name": "x",
                        "uri": "test://x",
                        "group": "g",
                        "extra": {},
                    },
                ],
            },
            SerializedAssetAll,
            id="asset_all",
        ),
        pytest.param(
            {
                "__type": DAT.ASSET_ANY,
                "objects": [
                    {
                        "__type": DAT.ASSET,
                        "name": "y",
                        "uri": "test://y",
                        "group": "g",
                        "extra": {},
                    }
                ],
            },
            SerializedAssetAny,
            id="asset_any",
        ),
        pytest.param(
            {"__type": DAT.ASSET_ALIAS, "name": "alias", "group": "g"},
            SerializedAssetAlias,
            id="asset_alias",
        ),
        pytest.param(
            {"__type": DAT.ASSET_REF, "name": "ref"},
            SerializedAssetRef,
            id="asset_ref",
        ),
    ],
)
def test_serde_decode_asset_condition_success(payload, expected_cls):
    from airflow.serialization.serialized_objects import decode_asset_like

    assert isinstance(decode_asset_like(payload), expected_cls)


def test_serde_decode_asset_condition_unknown_type():
    from airflow.serialization.serialized_objects import decode_asset_like

    with pytest.raises(
        ValueError,
        match="deserialization not implemented for DAT 'UNKNOWN_TYPE'",
    ):
        decode_asset_like({"__type": "UNKNOWN_TYPE"})


def test_encode_timezone():
    from airflow.serialization.serialized_objects import encode_timezone

    assert encode_timezone(FixedTimezone(0)) == "UTC"
    with pytest.raises(ValueError, match="DAG timezone should be a pendulum.tz.Timezone"):
        encode_timezone(object())


@pytest.mark.parametrize(
    ("cls", "args", "encode_type", "encode_var"),
    [
        (IdentityMapper, [], "airflow.partition_mappers.identity.IdentityMapper", {}),
        (
            HourlyMapper,
            [],
            "airflow.partition_mappers.temporal.HourlyMapper",
            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m-%dT%H"},
        ),
        (
            DailyMapper,
            [],
            "airflow.partition_mappers.temporal.DailyMapper",
            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m-%d"},
        ),
        (
            WeeklyMapper,
            [],
            "airflow.partition_mappers.temporal.WeeklyMapper",
            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m-%d (W%V)"},
        ),
        (
            MonthlyMapper,
            [],
            "airflow.partition_mappers.temporal.MonthlyMapper",
            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-%m"},
        ),
        (
            QuarterlyMapper,
            [],
            "airflow.partition_mappers.temporal.QuarterlyMapper",
            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y-Q{quarter}"},
        ),
        (
            YearlyMapper,
            [],
            "airflow.partition_mappers.temporal.YearlyMapper",
            {"input_format": "%Y-%m-%dT%H:%M:%S", "output_format": "%Y"},
        ),
    ],
)
def test_encode_partition_mapper(cls, args, encode_type, encode_var):
    from airflow.serialization.encoders import encode_partition_mapper

    partition_mapper = cls(*args)
    assert encode_partition_mapper(partition_mapper) == {
        Encoding.TYPE: encode_type,
        Encoding.VAR: encode_var,
    }


@pytest.mark.parametrize(
    ("sdk_cls", "core_cls"),
    [
        (IdentityMapper, CoreIdentityMapper),
        (HourlyMapper, CoureHourlyMapper),
        (DailyMapper, CoreDailyMapper),
        (WeeklyMapper, CoreWeeklyMapper),
        (MonthlyMapper, CoreMonthlyMapper),
        (QuarterlyMapper, CoreQuarterlyMapper),
        (YearlyMapper, CoreYearlyMapper),
    ],
)
def test_decode_partition_mapper(sdk_cls, core_cls):
    from airflow.serialization.decoders import decode_partition_mapper
    from airflow.serialization.encoders import encode_partition_mapper

    partition_mapper = sdk_cls()
    encoded_pm = encode_partition_mapper(partition_mapper)

    core_pm = decode_partition_mapper(encoded_pm)
    assert isinstance(core_pm, core_cls)


def test_decode_partition_mapper_not_exists():
    from airflow.serialization.decoders import decode_partition_mapper

    with pytest.raises(
        PartitionMapperNotFound,
        match=(
            "PartitionMapper class 'not_exists' could not be imported or "
            "you have a top level database access that disrupted the session. "
            "Please check the airflow best practices documentation."
        ),
    ):
        decode_partition_mapper({Encoding.TYPE: "not_exists", Encoding.VAR: {}})


def test_encode_product_mapper():
    from airflow.sdk import HourlyMapper, IdentityMapper, ProductMapper
    from airflow.serialization.encoders import encode_partition_mapper

    partition_mapper = ProductMapper(IdentityMapper(), HourlyMapper())
    assert encode_partition_mapper(partition_mapper) == {
        Encoding.TYPE: "airflow.partition_mappers.product.ProductMapper",
        Encoding.VAR: {
            "delimiter": "|",
            "mappers": [
                {
                    Encoding.TYPE: "airflow.partition_mappers.identity.IdentityMapper",
                    Encoding.VAR: {},
                },
                {
                    Encoding.TYPE: "airflow.partition_mappers.temporal.HourlyMapper",
                    Encoding.VAR: {
                        "input_format": "%Y-%m-%dT%H:%M:%S",
                        "output_format": "%Y-%m-%dT%H",
                    },
                },
            ],
        },
    }


def test_decode_product_mapper():
    from airflow.partition_mappers.product import ProductMapper as CoreProductMapper
    from airflow.sdk import DailyMapper, HourlyMapper, ProductMapper
    from airflow.serialization.decoders import decode_partition_mapper
    from airflow.serialization.encoders import encode_partition_mapper

    partition_mapper = ProductMapper(HourlyMapper(), DailyMapper())
    encoded_pm = encode_partition_mapper(partition_mapper)

    core_pm = decode_partition_mapper(encoded_pm)

    assert isinstance(core_pm, CoreProductMapper)
    assert len(core_pm.mappers) == 2
    assert core_pm.delimiter == "|"
    assert core_pm.to_downstream("2024-06-15T10:30:00|2024-06-15T10:30:00") == "2024-06-15T10|2024-06-15"


class TestSerializedBaseOperator:
    # ensure the default logging config is used for this test, no matter what ran before
    @pytest.mark.usefixtures("reset_logging_config")
    def test_logging_propogated_by_default(self, caplog):
        """Test that when set_context hasn't been called that log records are emitted"""
        BaseOperator(task_id="test").log.warning("test")
        # This looks like "how could it fail" but this actually checks that the handler called `emit`. Testing
        # the other case (that when we have set_context it goes to the file is harder to achieve without
        # leaking a lot of state)
        assert caplog.messages == ["test"]

    def test_resume_execution(self):
        from airflow.models.trigger import TriggerFailureReason
        from airflow.sdk.exceptions import TaskDeferralTimeout

        op = BaseOperator(task_id="hi")
        with pytest.raises(TaskDeferralTimeout):
            op.resume_execution(
                next_method="__fail__",
                next_kwargs={"error": TriggerFailureReason.TRIGGER_TIMEOUT},
                context={},
            )

    def test_deserialize_datetime_with_template_string(self):
        """Test that _deserialize_datetime handles template strings correctly."""
        from airflow.providers.standard.operators.trigger_dagrun import TriggerDagRunOperator
        from airflow.serialization.serialized_objects import DagSerialization

        # Create a DAG with a mapped operator that has a template string for logical_date
        with DAG(DAG_ID, start_date=DEFAULT_DATE) as dag:
            TriggerDagRunOperator.partial(
                task_id="test_trigger",
                logical_date="{{ ts_nodash_with_tz }}",  # Template string
                pool="default_pool",
            ).expand(trigger_dag_id=["dag1", "dag2"])

        # Serialize the DAG
        serialized_dag = DagSerialization.serialize_dag(dag)

        # Deserialize the DAG
        deserialized_dag = DagSerialization.deserialize_dag(serialized_dag)

        # Verify the operator was deserialized correctly
        task = deserialized_dag.task_dict["test_trigger"]
        assert task.partial_kwargs["logical_date"] == "{{ ts_nodash_with_tz }}"

    def test_deserialize_datetime_with_timestamp(self):
        """Test that _deserialize_datetime handles timestamp values correctly."""
        from airflow.serialization.serialized_objects import BaseSerialization

        # Test with a numeric timestamp
        timestamp = 1234567890.0
        result = BaseSerialization._deserialize_datetime(timestamp)

        # Should return a datetime object
        assert isinstance(result, datetime)
        assert result.timestamp() == timestamp


class TestKubernetesImportAvoidance:
    """Test that serialization doesn't import kubernetes unnecessarily."""

    def test_has_kubernetes_no_import_when_not_needed(self):
        """Ensure _has_kubernetes() doesn't import k8s when not already loaded."""
        # Remove kubernetes from sys.modules if present
        k8s_modules = [m for m in list(sys.modules.keys()) if m.startswith("kubernetes")]
        if k8s_modules:
            pytest.skip("Kubernetes already imported, cannot test import avoidance")

        # Call _has_kubernetes() - should check sys.modules and return False without importing
        _has_kubernetes.cache_clear()
        result = _has_kubernetes()

        assert result is False
        assert "kubernetes.client" not in sys.modules

    def test_has_kubernetes_uses_existing_import(self):
        """Ensure _has_kubernetes() uses k8s when it's already imported."""
        pytest.importorskip("kubernetes")

        # Now k8s is imported, should return True
        _has_kubernetes.cache_clear()
        result = _has_kubernetes()

        assert result is True