# 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 datetime import datetime, timedelta, timezone as dt_timezone
from typing import TYPE_CHECKING
from unittest import mock
from unittest.mock import MagicMock, patch
from uuid import UUID

import pytest
from pydantic import ValidationError

from airflow.sdk import BaseOperator, get_current_context, timezone
from airflow.sdk._shared.state import TaskScope
from airflow.sdk.api.datamodels._generated import (
    AssetEventResponse,
    AssetResponse,
    DagRun,
)
from airflow.sdk.bases.xcom import BaseXCom
from airflow.sdk.definitions.asset import (
    Asset,
    AssetAlias,
    AssetAliasEvent,
    AssetAliasUniqueKey,
    AssetNameRef,
    AssetUniqueKey,
    AssetUriRef,
)
from airflow.sdk.definitions.connection import Connection
from airflow.sdk.definitions.variable import Variable
from airflow.sdk.exceptions import AirflowNotFoundException, AirflowRuntimeError, ErrorType
from airflow.sdk.execution_time.comms import (
    AssetEventDagRunReferenceResult,
    AssetEventResult,
    AssetEventSourceTaskInstance,
    AssetEventsResult,
    AssetResult,
    AssetsByAliasResult,
    AssetStateStoreResult,
    ClearAssetStateStoreByName,
    ClearAssetStateStoreByUri,
    ClearTaskStateStore,
    ConnectionResult,
    DagRunResult,
    DeleteAssetStateStoreByName,
    DeleteAssetStateStoreByUri,
    DeleteTaskStateStore,
    ErrorResponse,
    GetAssetByName,
    GetAssetByUri,
    GetAssetEventByAsset,
    GetAssetsByAlias,
    GetAssetStateStoreByName,
    GetAssetStateStoreByUri,
    GetDagRun,
    GetTaskStateStore,
    GetXCom,
    OKResponse,
    SetAssetStateStoreByName,
    SetAssetStateStoreByUri,
    SetTaskStateStore,
    TaskStateStoreResult,
    VariableResult,
    XComResult,
)
from airflow.sdk.execution_time.context import (
    NEVER_EXPIRE,
    AssetStateStoreAccessor,
    AssetStateStoreAccessors,
    ConnectionAccessor,
    InletEventsAccessors,
    OutletEventAccessor,
    OutletEventAccessors,
    TaskStateStoreAccessor,
    TriggeringAssetEventsAccessor,
    VariableAccessor,
    _AssetRefResolutionMixin,
    _async_get_connection,
    _convert_variable_result_to_variable,
    _get_connection,
    _process_connection_result_conn,
    _wrap_external_ref,
    context_to_airflow_vars,
    set_current_context,
)
from airflow.sdk.execution_time.secrets import ExecutionAPISecretsBackend
from airflow.sdk.state import BaseStoreBackend

from tests_common.test_utils.config import conf_vars

if TYPE_CHECKING:
    from pydantic import JsonValue


def test_convert_connection_result_conn():
    """Test that the ConnectionResult is converted to a Connection object."""
    conn = ConnectionResult(
        conn_id="test_conn",
        conn_type="mysql",
        host="mysql",
        schema="airflow",
        login="root",
        password="password",
        port=1234,
        extra='{"extra_key": "extra_value"}',
    )
    conn = _process_connection_result_conn(conn)
    assert conn == Connection(
        conn_id="test_conn",
        conn_type="mysql",
        host="mysql",
        schema="airflow",
        login="root",
        password="password",
        port=1234,
        extra='{"extra_key": "extra_value"}',
    )


def test_convert_variable_result_to_variable():
    """Test that the VariableResult is converted to a Variable object."""
    var = VariableResult(
        key="test_key",
        value="test_value",
    )
    var = _convert_variable_result_to_variable(var, deserialize_json=False)
    assert var == Variable(
        key="test_key",
        value="test_value",
    )


def test_convert_variable_result_to_variable_with_deserialize_json():
    """Test that the VariableResult is converted to a Variable object with deserialize_json set to True."""
    var = VariableResult(
        key="test_key",
        value='{\r\n  "key1": "value1",\r\n  "key2": "value2",\r\n  "enabled": true,\r\n  "threshold": 42\r\n}',
    )
    var = _convert_variable_result_to_variable(var, deserialize_json=True)
    assert var == Variable(
        key="test_key", value={"key1": "value1", "key2": "value2", "enabled": True, "threshold": 42}
    )


class TestAirflowContextHelpers:
    def test_context_to_airflow_vars_empty_context(self):
        assert context_to_airflow_vars({}) == {}

    def test_context_to_airflow_vars_all_context(self, create_runtime_ti):
        task = BaseOperator(
            task_id="test_context_vars",
            owner=["owner1", "owner2"],
            email="email1@test.com",
        )

        rti = create_runtime_ti(
            task=task,
            dag_id="dag_id",
            run_id="dag_run_id",
            logical_date="2017-05-21T00:00:00Z",
            try_number=1,
        )
        context = rti.get_template_context()
        assert context_to_airflow_vars(context) == {
            "airflow.ctx.dag_id": "dag_id",
            "airflow.ctx.logical_date": "2017-05-21T00:00:00+00:00",
            "airflow.ctx.task_id": "test_context_vars",
            "airflow.ctx.dag_run_id": "dag_run_id",
            "airflow.ctx.try_number": "1",
            "airflow.ctx.dag_owner": "owner1,owner2",
            "airflow.ctx.dag_email": "email1@test.com",
        }

        assert context_to_airflow_vars(context, in_env_var_format=True) == {
            "AIRFLOW_CTX_DAG_ID": "dag_id",
            "AIRFLOW_CTX_LOGICAL_DATE": "2017-05-21T00:00:00+00:00",
            "AIRFLOW_CTX_TASK_ID": "test_context_vars",
            "AIRFLOW_CTX_TRY_NUMBER": "1",
            "AIRFLOW_CTX_DAG_RUN_ID": "dag_run_id",
            "AIRFLOW_CTX_DAG_OWNER": "owner1,owner2",
            "AIRFLOW_CTX_DAG_EMAIL": "email1@test.com",
        }

    def test_context_to_airflow_vars_team_name(self, create_runtime_ti):
        """``team_name`` on dag_run surfaces as AIRFLOW_CTX_TEAM_NAME when set; omitted when None."""
        task = BaseOperator(task_id="task")
        rti = create_runtime_ti(task=task)
        context = rti.get_template_context()

        # Default (team_name is None) -> key not present
        assert "AIRFLOW_CTX_TEAM_NAME" not in context_to_airflow_vars(context, in_env_var_format=True)

        context["dag_run"].team_name = "team-a"
        env_vars = context_to_airflow_vars(context, in_env_var_format=True)
        assert env_vars["AIRFLOW_CTX_TEAM_NAME"] == "team-a"
        assert context_to_airflow_vars(context)["airflow.ctx.team_name"] == "team-a"

    def test_context_to_airflow_vars_from_policy(self):
        with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method:
            airflow_cluster = "cluster-a"
            mock_method.return_value = {"airflow_cluster": airflow_cluster}

            context_vars = context_to_airflow_vars({})
            assert context_vars["airflow.ctx.airflow_cluster"] == airflow_cluster

            context_vars = context_to_airflow_vars({}, in_env_var_format=True)
            assert context_vars["AIRFLOW_CTX_AIRFLOW_CLUSTER"] == airflow_cluster

        with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method:
            mock_method.return_value = {"airflow_cluster": [1, 2]}
            with pytest.raises(TypeError) as error:
                context_to_airflow_vars({})
            assert str(error.value) == "value of key <airflow_cluster> must be string, not <class 'list'>"

        with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method:
            mock_method.return_value = {1: "value"}
            with pytest.raises(TypeError) as error:
                context_to_airflow_vars({})
            assert str(error.value) == "key <1> must be string"


class TestConnectionAccessor:
    def test_getattr_connection(self, mock_supervisor_comms):
        """
        Test that the connection is fetched when accessed via __getattr__.

        The __getattr__ method is used for template rendering. Example: ``{{ conn.mysql_conn.host }}``.
        """
        accessor = ConnectionAccessor()

        # Conn from the supervisor / API Server
        conn_result = ConnectionResult(
            conn_id="mysql_conn",
            conn_type="mysql",
            host="mysql",
            port=3306,
            schema=None,
            login=None,
            password=None,
            extra=None,
        )

        mock_supervisor_comms.send.return_value = conn_result

        # Fetch the connection; triggers __getattr__
        conn = accessor.mysql_conn

        expected_conn = Connection(conn_id="mysql_conn", conn_type="mysql", host="mysql", port=3306)
        assert conn == expected_conn

    def test_get_method_valid_connection(self, mock_supervisor_comms):
        """Test that the get method returns the requested connection using `conn.get`."""
        accessor = ConnectionAccessor()
        conn_result = ConnectionResult(
            conn_id="mysql_conn",
            conn_type="mysql",
            host="mysql",
            port=3306,
            schema=None,
            login=None,
            password=None,
            extra=None,
        )

        mock_supervisor_comms.send.return_value = conn_result

        conn = accessor.get("mysql_conn")
        assert conn == Connection(conn_id="mysql_conn", conn_type="mysql", host="mysql", port=3306)

    def test_get_method_with_default(self, mock_supervisor_comms):
        """Test that the get method returns the default connection when the requested connection is not found."""
        accessor = ConnectionAccessor()
        default_conn = {"conn_id": "default_conn", "conn_type": "sqlite"}
        error_response = ErrorResponse(
            error=ErrorType.CONNECTION_NOT_FOUND, detail={"conn_id": "nonexistent_conn"}
        )

        mock_supervisor_comms.send.return_value = error_response

        conn = accessor.get("nonexistent_conn", default_conn=default_conn)
        assert conn == default_conn

    def test_getattr_connection_for_extra_dejson(self, mock_supervisor_comms):
        accessor = ConnectionAccessor()

        # Conn from the supervisor / API Server
        conn_result = ConnectionResult(
            conn_id="mysql_conn",
            conn_type="mysql",
            host="mysql",
            port=3306,
            extra='{"extra_key": "extra_value"}',
            schema=None,
            login=None,
            password=None,
        )

        mock_supervisor_comms.send.return_value = conn_result

        # Fetch the connection's dejson; triggers __getattr__
        dejson = accessor.mysql_conn.extra_dejson

        assert dejson == {"extra_key": "extra_value"}

    @patch("airflow.sdk.definitions.connection.log", create=True)
    def test_getattr_connection_for_extra_dejson_decode_error(self, mock_log, mock_supervisor_comms):
        mock_log.return_value = MagicMock()

        accessor = ConnectionAccessor()

        # Conn from the supervisor / API Server
        conn_result = ConnectionResult(
            conn_id="mysql_conn",
            conn_type="mysql",
            host="mysql",
            port=3306,
            extra="This is not JSON!",
            schema=None,
            login=None,
            password=None,
        )

        mock_supervisor_comms.send.return_value = conn_result

        # Fetch the connection's dejson; triggers __getattr__
        dejson = accessor.mysql_conn.extra_dejson

        # empty in case of failed deserialising
        assert dejson == {}

        mock_log.exception.assert_any_call(
            "Failed to deserialize extra property `extra`, returning empty dictionary"
        )


class TestVariableAccessor:
    def test_getattr_variable(self, mock_supervisor_comms):
        """
        Test that the variable is fetched when accessed via __getattr__.
        """
        accessor = VariableAccessor(deserialize_json=False)

        # Variable from the supervisor / API Server
        var_result = VariableResult(key="test_key", value="test_value")

        mock_supervisor_comms.send.return_value = var_result

        # Fetch the variable; triggers __getattr__
        value = accessor.test_key

        assert value == var_result.value

    def test_get_method_valid_variable(self, mock_supervisor_comms):
        """Test that the get method returns the requested variable using `var.get`."""
        accessor = VariableAccessor(deserialize_json=False)
        var_result = VariableResult(key="test_key", value="test_value")

        mock_supervisor_comms.send.return_value = var_result

        val = accessor.get("test_key")
        assert val == var_result.value

    def test_get_method_with_default(self, mock_supervisor_comms):
        """Test that the get method returns the default variable when the requested variable is not found."""

        accessor = VariableAccessor(deserialize_json=False)
        error_response = ErrorResponse(error=ErrorType.VARIABLE_NOT_FOUND, detail={"test_key": "test_value"})

        mock_supervisor_comms.send.return_value = error_response

        val = accessor.get("nonexistent_var_key", default="default_value")
        assert val == "default_value"

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_json_masks_from_cache(self, mock_mask_secret):
        """SecretCache hit path applies the same masking as the backends path."""
        from airflow.sdk.execution_time.cache import SecretCache

        raw_json = '{"password": "s3cr3t", "host": "db.example.com"}'
        with mock.patch.object(SecretCache, "get_variable", return_value=raw_json):
            accessor = VariableAccessor(deserialize_json=True)
            val = accessor.db_config

        assert val == {"password": "s3cr3t", "host": "db.example.com"}
        mock_mask_secret.assert_any_call(raw_json, "db_config")
        mock_mask_secret.assert_any_call({"password": "s3cr3t", "host": "db.example.com"})

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_value_masks_secret(self, mock_mask_secret, mock_supervisor_comms):
        """var.value.<key> masks the string value when the key name is sensitive."""
        accessor = VariableAccessor(deserialize_json=False)
        mock_supervisor_comms.send.return_value = VariableResult(key="my_password", value="s3cr3t")

        val = accessor.my_password

        assert val == "s3cr3t"
        mock_mask_secret.assert_called_once_with("s3cr3t", "my_password")

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_json_masks_raw_string_and_dict_values(self, mock_mask_secret, mock_supervisor_comms):
        """var.json.<key> masks both the raw JSON string and the deserialized dict's sensitive fields."""
        accessor = VariableAccessor(deserialize_json=True)
        raw_json = '{"password": "s3cr3t", "host": "db.example.com"}'
        mock_supervisor_comms.send.return_value = VariableResult(key="db_config", value=raw_json)

        val = accessor.db_config

        assert val == {"password": "s3cr3t", "host": "db.example.com"}
        # First call: raw JSON string with variable key (masks if key name is sensitive)
        mock_mask_secret.assert_any_call(raw_json, "db_config")
        # Second call: deserialized dict so internal sensitive fields like "password" get masked
        mock_mask_secret.assert_any_call({"password": "s3cr3t", "host": "db.example.com"})

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_json_masks_list_values(self, mock_mask_secret, mock_supervisor_comms):
        """A JSON list is handed to the masker whole, exactly as a dict is."""
        accessor = VariableAccessor(deserialize_json=True)
        raw_json = '[{"password": "s3cr3t"}, {"password": "s3cr3t2"}]'
        mock_supervisor_comms.send.return_value = VariableResult(key="db_configs", value=raw_json)

        val = accessor.db_configs

        assert val == [{"password": "s3cr3t"}, {"password": "s3cr3t2"}]
        mock_mask_secret.assert_any_call(raw_json, "db_configs")
        # under the variable's key; dicts inside are still masked by their own key names
        mock_mask_secret.assert_any_call([{"password": "s3cr3t"}, {"password": "s3cr3t2"}], "db_configs")

    @pytest.mark.parametrize(
        ("raw", "expected"),
        [
            pytest.param("12345", 12345, id="int"),
            pytest.param("true", True, id="bool"),
            pytest.param("null", None, id="null"),
            pytest.param("1.5", 1.5, id="float"),
        ],
    )
    def test_var_json_scalar_values_pass_through(self, raw, expected, mock_supervisor_comms):
        """Handing a scalar to the masker is a no-op and must not change the value returned."""
        accessor = VariableAccessor(deserialize_json=True)
        mock_supervisor_comms.send.return_value = VariableResult(key="some_number", value=raw)

        assert accessor.some_number == expected

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_json_sensitive_key_masks_raw_json(self, mock_mask_secret, mock_supervisor_comms):
        """var.json.<sensitive_key> masks the entire raw JSON string because the variable key is sensitive."""
        accessor = VariableAccessor(deserialize_json=True)
        raw_json = '{"endpoint": "https://api.example.com", "token": "abc123"}'
        mock_supervisor_comms.send.return_value = VariableResult(key="my_secret", value=raw_json)

        val = accessor.my_secret

        assert val == {"endpoint": "https://api.example.com", "token": "abc123"}
        mock_mask_secret.assert_any_call(raw_json, "my_secret")
        mock_mask_secret.assert_any_call({"endpoint": "https://api.example.com", "token": "abc123"})

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_json_string_value_masks_both_forms(self, mock_mask_secret, mock_supervisor_comms):
        """var.json with a JSON string value masks both the quoted raw and the unquoted value."""
        accessor = VariableAccessor(deserialize_json=True)
        mock_supervisor_comms.send.return_value = VariableResult(key="my_token", value='"s3cr3t"')

        val = accessor.my_token

        assert val == "s3cr3t"
        assert mock_mask_secret.call_count == 2
        mock_mask_secret.assert_any_call('"s3cr3t"', "my_token")
        mock_mask_secret.assert_any_call("s3cr3t", "my_token")

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_json_list_value_does_not_over_mask(self, mock_mask_secret, mock_supervisor_comms):
        """var.json with a non-sensitive list variable does not mask individual list elements."""
        accessor = VariableAccessor(deserialize_json=True)
        raw_json = '["us-east-1", "eu-west-1"]'
        mock_supervisor_comms.send.return_value = VariableResult(key="aws_regions", value=raw_json)

        val = accessor.aws_regions

        assert val == ["us-east-1", "eu-west-1"]
        mock_mask_secret.assert_any_call(raw_json, "aws_regions")
        mock_mask_secret.assert_any_call(["us-east-1", "eu-west-1"], "aws_regions")
        # never anonymously -- that is what would mask the elements globally
        assert mock.call(["us-east-1", "eu-west-1"]) not in mock_mask_secret.call_args_list

    @mock.patch("airflow.sdk.execution_time.context.mask_secret")
    def test_var_json_invalid_json_raises(self, mock_mask_secret):
        """Invalid JSON raises JSONDecodeError; the raw value is still masked before the error."""
        import json

        from airflow.sdk.execution_time.cache import SecretCache

        raw = "not-valid-json"
        with mock.patch.object(SecretCache, "get_variable", return_value=raw):
            accessor = VariableAccessor(deserialize_json=True)
            with pytest.raises(json.JSONDecodeError):
                accessor.bad_var

        mock_mask_secret.assert_called_once_with(raw, "bad_var")


class TestCurrentContext:
    def test_current_context_roundtrip(self):
        example_context = {"Hello": "World"}

        with set_current_context(example_context):
            assert get_current_context() == example_context

    def test_context_removed_after_exit(self):
        example_context = {"Hello": "World"}

        with set_current_context(example_context):
            pass
        with pytest.raises(RuntimeError):
            get_current_context()

    def test_nested_context(self):
        """
        Nested execution context should be supported in case the user uses multiple context managers.
        Each time the execute method of an operator is called, we set a new 'current' context.
        This test verifies that no matter how many contexts are entered - order is preserved
        """
        max_stack_depth = 15
        ctx_list = []
        for i in range(max_stack_depth):
            # Create all contexts in ascending order
            new_context = {"ContextId": i}
            # Like 15 nested with statements
            ctx_obj = set_current_context(new_context)
            ctx_obj.__enter__()
            ctx_list.append(ctx_obj)
        for i in reversed(range(max_stack_depth)):
            # Iterate over contexts in reverse order - stack is LIFO
            ctx = get_current_context()
            assert ctx["ContextId"] == i
            # End of with statement
            ctx_list[i].__exit__(None, None, None)


class TestOutletEventAccessor:
    @pytest.mark.parametrize(
        "add_args",
        [
            (Asset("name", "uri", extra={"extra": "from asset itself"}), {"extra": "from event"}),
            (Asset.ref(name="name"), {"extra": "from event"}),
            (Asset.ref(uri="uri"), {"extra": "from event"}),
        ],
        ids=["asset", "asset name ref", "asset uri ref"],
    )
    @pytest.mark.parametrize(
        ("key", "asset_alias_events"),
        (
            (AssetUniqueKey.from_asset(Asset("test_uri")), []),
            (
                AssetAliasUniqueKey.from_asset_alias(AssetAlias("test_alias")),
                [
                    AssetAliasEvent(
                        source_alias_name="test_alias",
                        dest_asset_key=AssetUniqueKey(name="name", uri="uri"),
                        dest_asset_extra={"extra": "from asset itself"},
                        extra={"extra": "from event"},
                    )
                ],
            ),
        ),
        ids=["inactive asset", "active asset"],
    )
    def test_add(self, add_args, key, asset_alias_events, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = AssetResponse(
            name="name", uri="uri", group="", extra={"extra": "from asset itself"}
        )

        outlet_event_accessor = OutletEventAccessor(key=key, extra={})
        outlet_event_accessor.add(*add_args)
        assert outlet_event_accessor.asset_alias_events == asset_alias_events

    @pytest.mark.parametrize(
        "add_args",
        [
            (Asset(name="name", uri="uri", extra={"extra": "from asset itself"}), {"extra": "from event"}),
            (Asset.ref(name="name"), {"extra": "from event"}),
            (Asset.ref(uri="uri"), {"extra": "from event"}),
        ],
        ids=["asset", "asset name ref", "asset uri ref"],
    )
    @pytest.mark.parametrize(
        ("key", "asset_alias_events"),
        (
            (AssetUniqueKey.from_asset(Asset("test_uri")), []),
            (
                AssetAliasUniqueKey.from_asset_alias(AssetAlias("test_alias")),
                [
                    AssetAliasEvent(
                        source_alias_name="test_alias",
                        dest_asset_key=AssetUniqueKey(name="name", uri="uri"),
                        dest_asset_extra={"extra": "from asset itself"},
                        extra={"extra": "from event"},
                    )
                ],
            ),
        ),
        ids=["inactive asset", "active asset"],
    )
    def test_add_with_db(self, add_args, key, asset_alias_events, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = AssetResponse(
            name="name", uri="uri", group="", extra={"extra": "from asset itself"}
        )

        outlet_event_accessor = OutletEventAccessor(key=key)
        outlet_event_accessor.add(*add_args)
        assert outlet_event_accessor.asset_alias_events == asset_alias_events


class TestOutletEventAccessorPartitionKeys:
    @pytest.fixture
    def accessor(self) -> OutletEventAccessor:
        return OutletEventAccessor(key=AssetUniqueKey.from_asset(Asset("a")))

    def test_default_is_empty(self, accessor):
        assert accessor.partition_keys == set()

    def test_direct_assignment(self, accessor):
        accessor.partition_keys = {"us", "eu"}
        assert accessor.partition_keys == {"us", "eu"}

    def test_add_partitions(self, accessor):
        accessor.add_partitions("us")
        assert accessor.partition_keys == {"us"}

    def test_add_partitions_appends(self, accessor):
        accessor.add_partitions("us")
        accessor.add_partitions("eu")
        accessor.add_partitions("apac")
        assert accessor.partition_keys == {"us", "eu", "apac"}

    def test_add_partitions_dedupes(self, accessor):
        accessor.add_partitions("us")
        accessor.add_partitions("us")
        accessor.add_partitions(["us", "eu"])
        assert accessor.partition_keys == {"us", "eu"}

    @pytest.mark.parametrize(
        "key",
        [
            "",
            "   ",
            "\t",
        ],
        ids=["empty", "spaces", "tab"],
    )
    def test_add_partitions_rejects_empty_key(self, accessor, key):
        with pytest.raises(ValueError, match="must not be empty or whitespace-only"):
            accessor.add_partitions(key)

    @pytest.mark.parametrize(
        "key",
        [
            "a" * 250,
            "a" * 251,
        ],
        ids=["at_limit_accepted", "over_limit_rejected"],
    )
    def test_add_partitions_length_boundary(self, accessor, key):
        if len(key) <= accessor._PARTITION_KEY_MAX_LENGTH:
            accessor.add_partitions(key)
            assert key in accessor.partition_keys
        else:
            with pytest.raises(ValueError, match="at most 250 characters"):
                accessor.add_partitions(key)

    def test_add_partitions_rejects_any_invalid_in_list(self, accessor):
        """A list with a mix of valid and invalid keys fails before any are added."""
        with pytest.raises(ValueError, match="must not be empty or whitespace-only"):
            accessor.add_partitions(["us", ""])
        assert accessor.partition_keys == set()

    def test_add_partitions_rejects_asset_alias_accessor(self):
        alias_accessor = OutletEventAccessor(
            key=AssetAliasUniqueKey.from_asset_alias(AssetAlias("test_alias"))
        )
        with pytest.raises(TypeError, match="not supported on asset alias"):
            alias_accessor.add_partitions("us")
        assert alias_accessor.partition_keys == set()


class TestTriggeringAssetEventsAccessor:
    @pytest.fixture(autouse=True)
    def clear_cache(self):
        _AssetRefResolutionMixin._asset_ref_cache = {}
        yield
        _AssetRefResolutionMixin._asset_ref_cache = {}

    @pytest.fixture
    def event_data(self):
        return [
            {
                "asset": {
                    "name": "1",
                    "uri": "1",
                    "extra": {},
                },
                "extra": {},
                "source_task_id": "t1",
                "source_dag_id": "d1",
                "source_run_id": "r1",
                "source_map_index": -1,
                "source_aliases": [],
                "timestamp": "2025-01-01T00:00:12Z",
            },
            {
                "asset": {
                    "name": "1",
                    "uri": "1",
                    "extra": {},
                },
                "extra": {},
                "source_task_id": "t2",
                "source_dag_id": "d1",
                "source_run_id": "r1",
                "source_map_index": -1,
                "source_aliases": [
                    {"name": "a"},
                    {"name": "b"},
                ],
                "timestamp": "2025-01-01T00:05:43Z",
            },
            {
                "asset": {
                    "name": "2",
                    "uri": "2",
                    "extra": {},
                },
                "extra": {},
                "source_task_id": "t2",
                "source_dag_id": "d1",
                "source_run_id": "r1",
                "source_map_index": -1,
                "source_aliases": [],
                "timestamp": "2025-01-01T00:06:07Z",
            },
        ]

    @pytest.fixture
    def accessor(self, event_data):
        return TriggeringAssetEventsAccessor.build(
            AssetEventDagRunReferenceResult.model_validate(d) for d in event_data
        )

    @pytest.mark.parametrize(
        ("key", "result_indexes"),
        [
            (Asset("1"), [0, 1]),
            (Asset("2"), [2]),
            (AssetAlias("a"), [1]),
            (AssetAlias("b"), [1]),
        ],
    )
    def test_getitem(self, event_data, accessor, key, result_indexes):
        expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes]
        assert accessor[key] == expected

    @pytest.mark.parametrize(
        ("name", "resolved_asset", "result_indexes"),
        [
            ("1", AssetResult(name="1", uri="1", group="whatever"), [0, 1]),
            ("2", AssetResult(name="2", uri="2", group="whatever"), [2]),
        ],
    )
    def test_getitem_name_ref(
        self,
        mock_supervisor_comms,
        event_data,
        accessor,
        name,
        resolved_asset,
        result_indexes,
    ):
        mock_supervisor_comms.send.return_value = resolved_asset
        expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes]
        assert accessor[Asset.ref(name=name)] == expected

        assert mock_supervisor_comms.send.mock_calls == [
            mock.call(GetAssetByName(name=name, type="GetAssetByName"))
        ]
        assert _AssetRefResolutionMixin._asset_ref_cache

    @pytest.mark.parametrize(
        ("uri", "resolved_asset", "result_indexes"),
        [
            ("1", AssetResult(name="1", uri="1", group="whatever"), [0, 1]),
            ("2", AssetResult(name="2", uri="2", group="whatever"), [2]),
        ],
    )
    def test_getitem_uri_ref(
        self,
        mock_supervisor_comms,
        event_data,
        accessor,
        uri,
        resolved_asset,
        result_indexes,
    ):
        mock_supervisor_comms.send.return_value = resolved_asset
        expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes]
        assert accessor[Asset.ref(uri=uri)] == expected
        assert mock_supervisor_comms.send.mock_calls == [mock.call(GetAssetByUri(uri=uri))]
        assert _AssetRefResolutionMixin._asset_ref_cache

    def test_partition_key_exposed(self):
        """A consumed asset event's partition key is reachable via triggering_asset_events."""
        event = {
            "asset": {"name": "1", "uri": "1", "extra": {}},
            "extra": {},
            "source_task_id": "t1",
            "source_dag_id": "d1",
            "source_run_id": "r1",
            "source_map_index": -1,
            "source_aliases": [],
            "timestamp": "2025-01-01T00:00:12Z",
            "partition_key": "2024-01-15",
        }
        accessor = TriggeringAssetEventsAccessor.build(
            [AssetEventDagRunReferenceResult.model_validate(event)]
        )
        assert [e.partition_key for e in accessor[Asset("1")]] == ["2024-01-15"]

    def test_source_task_instance_xcom_pull(self, mock_supervisor_comms, accessor):
        events = accessor[Asset("2")]
        assert len(events) == 1

        mock_dag_run = mock.Mock(dag_id="d1", run_id="r1")
        mock_supervisor_comms.send.side_effect = [mock_dag_run]
        source = events[0].source_task_instance
        assert source == AssetEventSourceTaskInstance(dag_run=mock_dag_run, task_id="t2", map_index=-1)
        assert mock_supervisor_comms.send.mock_calls == [mock.call(GetDagRun(dag_id="d1", run_id="r1"))]

        mock_supervisor_comms.reset_mock()
        mock_supervisor_comms.send.side_effect = [
            XComResult(key=BaseXCom.XCOM_RETURN_KEY, value="__example_xcom_value__"),
        ]
        assert source.xcom_pull() == "__example_xcom_value__"
        assert mock_supervisor_comms.send.mock_calls == [
            mock.call(
                GetXCom(
                    key=BaseXCom.XCOM_RETURN_KEY,
                    dag_id="d1",
                    run_id="r1",
                    task_id="t2",
                    map_index=-1,
                ),
            )
        ]


TEST_ASSET = Asset(name="test_uri", uri="test://test")
TEST_ASSET_ALIAS = AssetAlias(name="name")
TEST_ASSET_REFS = [Asset.ref(name="test_uri"), Asset.ref(uri="test://test/")]
TEST_INLETS = [TEST_ASSET, TEST_ASSET_ALIAS] + TEST_ASSET_REFS


class TestOutletEventAccessors:
    @pytest.mark.parametrize(
        ("access_key", "internal_key"),
        (
            (Asset("test"), AssetUniqueKey.from_asset(Asset("test"))),
            (
                Asset(name="test", uri="test://asset"),
                AssetUniqueKey.from_asset(Asset(name="test", uri="test://asset")),
            ),
            (AssetAlias("test_alias"), AssetAliasUniqueKey.from_asset_alias(AssetAlias("test_alias"))),
        ),
    )
    def test__get_item__dict_key_not_exists(self, access_key, internal_key):
        outlet_event_accessors = OutletEventAccessors()
        assert len(outlet_event_accessors) == 0
        outlet_event_accessor = outlet_event_accessors[access_key]
        assert len(outlet_event_accessors) == 1
        assert outlet_event_accessor.key == internal_key
        assert outlet_event_accessor.extra == {}

    @pytest.mark.parametrize(
        ("access_key", "asset"),
        (
            (Asset.ref(name="test"), Asset(name="test")),
            (Asset.ref(name="test1"), Asset(name="test1", uri="test://asset-uri")),
            (Asset.ref(uri="test://asset-uri"), Asset(uri="test://asset-uri")),
        ),
    )
    def test__get_item__asset_ref(self, access_key, asset, mock_supervisor_comms):
        """Test accessing OutletEventAccessors with AssetRef resolves to correct Asset."""
        internal_key = AssetUniqueKey.from_asset(asset)
        outlet_event_accessors = OutletEventAccessors()
        assert len(outlet_event_accessors) == 0

        # Asset from the API Server via the supervisor
        mock_supervisor_comms.send.return_value = AssetResult(
            name=asset.name,
            uri=asset.uri,
            group=asset.group,
        )

        outlet_event_accessor = outlet_event_accessors[access_key]
        assert len(outlet_event_accessors) == 1
        assert outlet_event_accessor.key == internal_key
        assert outlet_event_accessor.extra == {}

    @pytest.mark.parametrize(
        ("name", "uri", "expected_key"),
        (
            ("test_uri", "test://test/", TEST_ASSET),
            ("test_uri", None, TEST_ASSET_REFS[0]),
            (None, "test://test/", TEST_ASSET_REFS[1]),
        ),
    )
    @mock.patch("airflow.sdk.execution_time.context.OutletEventAccessors.__getitem__")
    def test_for_asset(self, mocked__getitem__, name, uri, expected_key):
        outlet_event_accessors = OutletEventAccessors()
        outlet_event_accessors.for_asset(name=name, uri=uri)
        assert mocked__getitem__.call_args[0][0] == expected_key

    @mock.patch("airflow.sdk.execution_time.context.OutletEventAccessors.__getitem__")
    def test_for_asset_alias(self, mocked__getitem__):
        outlet_event_accessors = OutletEventAccessors()
        outlet_event_accessors.for_asset_alias(name="name")
        assert mocked__getitem__.call_args[0][0] == TEST_ASSET_ALIAS


class TestInletEventAccessor:
    @pytest.fixture
    def sample_inlet_evnets_accessor(self, mock_supervisor_comms):
        mock_supervisor_comms.send.side_effect = [
            AssetResult(name="test_uri", uri="test://test", group="asset"),
            AssetResult(name="test_uri", uri="test://test", group="asset"),
        ]
        obj = InletEventsAccessors(inlets=TEST_INLETS)
        mock_supervisor_comms.reset_mock()
        return obj

    @pytest.mark.usefixtures("mock_supervisor_comms")
    def test__iter__(self, sample_inlet_evnets_accessor):
        for actual, expected in zip(sample_inlet_evnets_accessor, TEST_INLETS):
            assert actual == expected

    @pytest.mark.usefixtures("mock_supervisor_comms")
    def test__len__(self, sample_inlet_evnets_accessor):
        assert len(sample_inlet_evnets_accessor) == 4

    @pytest.mark.parametrize("key", TEST_INLETS + [0, 1, 2, 3])
    def test__get_item__(self, key, sample_inlet_evnets_accessor, mock_supervisor_comms):
        # This test only verifies a valid key can be used to access inlet events,
        # but not access asset events are fetched. That is verified in test_asset_events in execution_api
        asset_event_resp = AssetEventResult(
            id=1,
            created_dagruns=[],
            timestamp=timezone.utcnow(),
            asset=AssetResponse(name="test", uri="test", group="asset"),
        )
        events_result = AssetEventsResult(asset_events=[asset_event_resp])
        mock_supervisor_comms.send.side_effect = [events_result] * 4

        assert list(sample_inlet_evnets_accessor[key]) == [asset_event_resp]

    @pytest.mark.usefixtures("mock_supervisor_comms")
    def test__get_item__out_of_index(self, sample_inlet_evnets_accessor):
        with pytest.raises(IndexError):
            sample_inlet_evnets_accessor[5]

    def test__get_item__with_filters(self, sample_inlet_evnets_accessor, mock_supervisor_comms):
        asset_event_resp = AssetEventResult(
            id=1,
            created_dagruns=[],
            timestamp=timezone.utcnow(),
            asset=AssetResponse(name="test_uri", uri="test_uri", group="asset"),
        )
        events_result = AssetEventsResult(asset_events=[asset_event_resp])
        mock_supervisor_comms.send.side_effect = [events_result] * 10

        list(sample_inlet_evnets_accessor[TEST_ASSET])
        list(sample_inlet_evnets_accessor[TEST_ASSET].after("2024-01-01T00:00:00Z"))
        list(sample_inlet_evnets_accessor[TEST_ASSET].before("2024-01-01T00:00:00Z"))
        list(sample_inlet_evnets_accessor[TEST_ASSET].limit(10))
        list(
            sample_inlet_evnets_accessor[TEST_ASSET]
            .after("2024-01-01T00:00:00Z")
            .before("2024-01-02T00:00:00Z")
            .limit(10)
        )
        list(sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10))

        assert mock_supervisor_comms.send.call_count == 6

        # test accessing the accessor without list() or []
        sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10)

        assert mock_supervisor_comms.send.call_count == 6

        # test accessing one of the elements
        res = sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10)[0]
        assert res == asset_event_resp
        assert mock_supervisor_comms.send.call_count == 7

        # test evaluating the accessor multiple times with the same filters
        res = sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10)
        assert res[0] == asset_event_resp
        assert res[0] == asset_event_resp

        assert mock_supervisor_comms.send.call_count == 8

        # test changing one of the filters
        assert res.after("2024-01-01T00:00:00Z")[0] == asset_event_resp

        assert mock_supervisor_comms.send.call_count == 9

        # test len()
        assert len(sample_inlet_evnets_accessor[TEST_ASSET].ascending(True).limit(10)) == 1
        assert mock_supervisor_comms.send.call_count == 10

        calls = mock_supervisor_comms.send.call_args_list
        assert calls[0][0][0] == GetAssetEventByAsset(
            name="test_uri", uri="test://test/", after=None, before=None, limit=None, ascending=True
        )
        assert calls[1][0][0] == GetAssetEventByAsset(
            name="test_uri",
            uri="test://test/",
            after="2024-01-01T00:00:00Z",
            before=None,
            limit=None,
            ascending=True,
        )
        assert calls[2][0][0] == GetAssetEventByAsset(
            name="test_uri",
            uri="test://test/",
            after=None,
            before="2024-01-01T00:00:00Z",
            limit=None,
            ascending=True,
        )
        assert calls[3][0][0] == GetAssetEventByAsset(
            name="test_uri", uri="test://test/", after=None, before=None, limit=10, ascending=True
        )
        assert calls[4][0][0] == GetAssetEventByAsset(
            name="test_uri",
            uri="test://test/",
            after="2024-01-01T00:00:00Z",
            before="2024-01-02T00:00:00Z",
            limit=10,
            ascending=True,
        )
        assert calls[5][0][0] == GetAssetEventByAsset(
            name="test_uri", uri="test://test/", after=None, before=None, limit=10, ascending=False
        )

    @pytest.mark.parametrize(
        ("name", "uri", "expected_key"),
        (
            ("test_uri", "test://test/", TEST_ASSET),
            ("test_uri", None, TEST_ASSET_REFS[0]),
            (None, "test://test/", TEST_ASSET_REFS[1]),
        ),
    )
    @mock.patch("airflow.sdk.execution_time.context.InletEventsAccessors.__getitem__")
    def test_for_asset(self, mocked__getitem__, sample_inlet_evnets_accessor, name, uri, expected_key):
        sample_inlet_evnets_accessor.for_asset(name=name, uri=uri)
        assert mocked__getitem__.call_args[0][0] == expected_key

    @mock.patch("airflow.sdk.execution_time.context.InletEventsAccessors.__getitem__")
    def test_for_asset_alias(self, mocked__getitem__, sample_inlet_evnets_accessor):
        sample_inlet_evnets_accessor.for_asset_alias(name="name")
        assert mocked__getitem__.call_args[0][0] == TEST_ASSET_ALIAS

    def test_source_task_instance_xcom_pull(self, sample_inlet_evnets_accessor, mock_supervisor_comms):
        mock_supervisor_comms.send.side_effect = [
            AssetEventsResult(
                asset_events=[
                    AssetEventResponse(
                        id=1,
                        timestamp=timezone.utcnow(),
                        asset=AssetResponse(name="test_uri", uri="test://test", group="asset"),
                        created_dagruns=[],
                        source_dag_id="__dag__",
                        source_run_id="__run__",
                        source_task_id="__task__",
                        source_map_index=0,
                    ),
                    AssetEventResponse(
                        id=1,
                        timestamp=timezone.utcnow(),
                        asset=AssetResponse(name="test_uri", uri="test://test", group="asset"),
                        created_dagruns=[],
                    ),
                ],
            )
        ]
        events = list(sample_inlet_evnets_accessor[Asset.ref(name="test_uri")])
        assert mock_supervisor_comms.send.mock_calls == [
            mock.call(
                GetAssetEventByAsset(
                    name="test_uri",
                    uri=None,
                    after=None,
                    before=None,
                    limit=None,
                    ascending=True,
                )
            )
        ]

        assert len(events) == 2

        dag_run_result = DagRunResult(
            dag_id="__dag__",
            run_id="__run__",
            run_after=timezone.utcnow(),
            start_date=timezone.utcnow(),
            run_type="scheduled",
            state="success",
            consumed_asset_events=[],
            logical_date=None,
            data_interval_start=None,
            data_interval_end=None,
            end_date=None,
            partition_key=None,
        )
        mock_supervisor_comms.reset_mock()
        mock_supervisor_comms.send.side_effect = [dag_run_result]
        assert events[1].source_task_instance is None
        source = events[0].source_task_instance
        assert source == AssetEventSourceTaskInstance(dag_run=dag_run_result, task_id="__task__", map_index=0)
        assert mock_supervisor_comms.send.mock_calls == [
            mock.call(GetDagRun(dag_id="__dag__", run_id="__run__"))
        ]

        mock_supervisor_comms.reset_mock()
        mock_supervisor_comms.send.side_effect = [
            XComResult(key=BaseXCom.XCOM_RETURN_KEY, value="__example_xcom_value__"),
        ]
        assert source.xcom_pull() == "__example_xcom_value__"
        assert mock_supervisor_comms.send.mock_calls == [
            mock.call(
                GetXCom(
                    key=BaseXCom.XCOM_RETURN_KEY,
                    dag_id="__dag__",
                    run_id="__run__",
                    task_id="__task__",
                    map_index=0,
                ),
            )
        ]


class TestDagRunStartDateNullable:
    """Test that DagRun and TIRunContext accept start_date=None (queued runs that haven't started)."""

    def test_dag_run_model_accepts_null_start_date(self):
        """DagRun datamodel should accept start_date=None for runs that haven't started yet."""
        dag_run = DagRun(
            dag_id="test_dag",
            run_id="test_run",
            logical_date="2024-12-01T01:00:00Z",
            data_interval_start="2024-12-01T00:00:00Z",
            data_interval_end="2024-12-01T01:00:00Z",
            start_date=None,
            run_after="2024-12-01T01:00:00Z",
            run_type="manual",
            state="queued",
            conf=None,
            consumed_asset_events=[],
            end_date=None,
            partition_key=None,
        )

        assert dag_run.start_date is None

    def test_ti_run_context_with_null_start_date(self, make_ti_context):
        """TIRunContext should be constructable when the DagRun has start_date=None."""
        ti_context = make_ti_context(start_date=None)

        assert ti_context.dag_run.start_date is None


class TestAsyncGetConnection:
    """Test async connection retrieval with secrets backends."""

    @pytest.mark.asyncio
    async def test_async_get_connection_from_secrets_backend(self, mock_supervisor_comms):
        """Test that _async_get_connection successfully retrieves from secrets backend using sync_to_async."""
        sample_connection = Connection(
            conn_id="test_conn", conn_type="postgres", host="localhost", port=5432, login="user"
        )

        class MockSecretsBackend:
            """Simple mock secrets backend for testing."""

            def __init__(self, connections: dict[str, Connection | None] | None = None):
                self.connections = connections or {}

            def get_connection(self, conn_id: str) -> Connection | None:
                return self.connections.get(conn_id)

        backend = MockSecretsBackend({"test_conn": sample_connection})

        with patch(
            "airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", autospec=True
        ) as mock_load:
            mock_load.return_value = [backend]

            result = await _async_get_connection("test_conn")

            assert result == sample_connection
            # Should not have tried SUPERVISOR_COMMS since secrets backend had the connection
            mock_supervisor_comms.send.assert_not_called()
            mock_supervisor_comms.asend.assert_not_called()


class TestSecretsBackend:
    """Test that connection resolution uses the backend chain correctly."""

    def test_execution_api_backend_in_worker_chain(self):
        """Test that ExecutionAPISecretsBackend is in the worker search path."""
        from airflow.sdk.execution_time.secrets import DEFAULT_SECRETS_SEARCH_PATH_WORKERS

        assert (
            "airflow.sdk.execution_time.secrets.execution_api.ExecutionAPISecretsBackend"
            in DEFAULT_SECRETS_SEARCH_PATH_WORKERS
        )

    def test_metastore_backend_in_server_chain(self):
        """Test that MetastoreBackend is in the API server search path."""
        from airflow.sdk.execution_time.secrets import _SERVER_DEFAULT_SECRETS_SEARCH_PATH

        assert "airflow.secrets.metastore.MetastoreBackend" in _SERVER_DEFAULT_SECRETS_SEARCH_PATH
        assert (
            "airflow.sdk.execution_time.secrets.execution_api.ExecutionAPISecretsBackend"
            not in _SERVER_DEFAULT_SECRETS_SEARCH_PATH
        )

    def test_get_connection_uses_backend_chain(self, mock_supervisor_comms):
        """Test that _get_connection properly iterates through backends."""
        from airflow.sdk.api.datamodels._generated import ConnectionResponse
        from airflow.sdk.execution_time.comms import ConnectionResult

        # Mock connection response
        conn_response = ConnectionResponse(
            conn_id="test_conn",
            conn_type="http",
            host="example.com",
            port=443,
            schema=None,
            login=None,
            password=None,
            extra=None,
        )
        conn_result = ConnectionResult.from_conn_response(conn_response)
        mock_supervisor_comms.send.return_value = conn_result

        # Mock the backend loading to include our SupervisorComms backend
        supervisor_backend = ExecutionAPISecretsBackend()

        with patch("airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded") as mock_load:
            mock_load.return_value = [supervisor_backend]

            conn = _get_connection("test_conn")

            assert conn is not None
            assert conn.conn_id == "test_conn"
            assert conn.host == "example.com"
            mock_supervisor_comms.send.assert_called_once()

    def test_get_connection_backend_fallback(self, mock_supervisor_comms):
        """Test that _get_connection falls through backends correctly."""
        from airflow.sdk.api.datamodels._generated import ConnectionResponse
        from airflow.sdk.execution_time.comms import ConnectionResult

        # First backend returns nothing (simulating env var backend with no env var)
        class EmptyBackend:
            def get_connection(self, conn_id):
                return None

        # Second backend returns the connection
        conn_response = ConnectionResponse(
            conn_id="test_conn",
            conn_type="postgres",
            host="db.example.com",
            schema=None,
            login=None,
            password=None,
            port=None,
            extra=None,
        )
        conn_result = ConnectionResult.from_conn_response(conn_response)
        mock_supervisor_comms.send.return_value = conn_result

        supervisor_backend = ExecutionAPISecretsBackend()

        with patch("airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded") as mock_load:
            mock_load.return_value = [EmptyBackend(), supervisor_backend]

            conn = _get_connection("test_conn")

            assert conn is not None
            assert conn.conn_id == "test_conn"
            # SupervisorComms backend was called (first backend returned None)
            mock_supervisor_comms.send.assert_called_once()

    def test_get_connection_not_found_raises_error(self, mock_supervisor_comms):
        """Test that _get_connection raises error when no backend finds connection."""

        # Backend returns None (not found)
        class EmptyBackend:
            def get_connection(self, conn_id):
                return None

        with patch("airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded") as mock_load:
            mock_load.return_value = [EmptyBackend()]

            with pytest.raises(AirflowNotFoundException, match="isn't defined"):
                _get_connection("nonexistent_conn")


class TestTaskStateStoreAccessor:
    TI_ID = UUID("01900000-0000-0000-0000-000000000001")
    SCOPE = TaskScope(dag_id="dag", run_id="run", task_id="task")

    def test_get_returns_value(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = TaskStateStoreResult(value="app_001")

        result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("job_id")

        assert result == "app_001"
        mock_supervisor_comms.send.assert_called_once_with(GetTaskStateStore(ti_id=self.TI_ID, key="job_id"))

    def test_get_returns_none_on_404(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = ErrorResponse(
            error=ErrorType.TASK_STORE_NOT_FOUND, detail={"key": "missing_key"}
        )

        result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("missing_key")

        assert result is None

    def test_get_returns_default_when_key_missing(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = ErrorResponse(
            error=ErrorType.TASK_STORE_NOT_FOUND, detail={"key": "job_id"}
        )

        result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get(
            "job_id", default="default-id"
        )

        assert result == "default-id"

    def test_get_ignores_default_when_key_exists(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = TaskStateStoreResult(value="job-001")

        result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get(
            "job_id", default="do-not-start-here"
        )

        assert result == "job-001"

    def test_get_raises_on_error(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = ErrorResponse(
            error=ErrorType.GENERIC_ERROR, detail={"message": "server error"}
        )

        with pytest.raises(AirflowRuntimeError):
            TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("some_key")

    def test_set_none_raises(self, mock_supervisor_comms):
        with pytest.raises(ValueError, match="Cannot set value as None"):
            TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", None)

    def test_set_operation_with_global_retention(self, mock_supervisor_comms, time_machine):
        """set() with no retention uses global default_retention_days config."""

        mock_supervisor_comms.send.return_value = OKResponse(ok=True)
        now = datetime(2026, 5, 14, 12, 0, 0, tzinfo=dt_timezone.utc)
        time_machine.move_to(now, tick=False)

        with conf_vars({("state_store", "default_retention_days"): "30"}):
            TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001")

        mock_supervisor_comms.send.assert_called_once_with(
            SetTaskStateStore(
                ti_id=self.TI_ID,
                key="job_id",
                value="app_001",
                expires_at=datetime(2026, 6, 13, 12, 0, 0, tzinfo=dt_timezone.utc),
            )
        )

    def test_set_with_retention_computes_expires_at(self, mock_supervisor_comms, time_machine):
        """set(retention=timedelta(...)) computes expires_at on the worker and sends it."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)
        now = datetime(2026, 5, 14, 12, 0, 0, tzinfo=dt_timezone.utc)
        time_machine.move_to(now, tick=False)

        TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set(
            "job_id", "app_001", retention=timedelta(days=7)
        )

        mock_supervisor_comms.send.assert_called_once_with(
            SetTaskStateStore(
                ti_id=self.TI_ID,
                key="job_id",
                value="app_001",
                expires_at=datetime(2026, 5, 21, 12, 0, 0, tzinfo=dt_timezone.utc),
            )
        )

    def test_set_with_never_expire_sends_null_expires_at(self, mock_supervisor_comms):
        """set(retention=NEVER_EXPIRE) sends expires_at=None"""

        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set(
            "job_id", "app_001", retention=NEVER_EXPIRE
        )

        mock_supervisor_comms.send.assert_called_once_with(
            SetTaskStateStore(ti_id=self.TI_ID, key="job_id", value="app_001", expires_at=None)
        )

    def test_set_global_default_zero_sends_null_expires_at(self, mock_supervisor_comms):
        """When default_retention_days=0 (never expire globally), expires_at=None (stored as NULL)."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        with conf_vars({("state_store", "default_retention_days"): "0"}):
            TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001")

        mock_supervisor_comms.send.assert_called_once_with(
            SetTaskStateStore(ti_id=self.TI_ID, key="job_id", value="app_001", expires_at=None)
        )

    def test_set_raises_on_negative_retention_days(self, mock_supervisor_comms):
        """set() raises ValueError when default_retention_days is negative."""
        with conf_vars({("state_store", "default_retention_days"): "-1"}):
            with pytest.raises(ValueError, match="default_retention_days must be >= 0"):
                TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001")

    def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms):
        """set() logs a warning when the serialized value exceeds max_value_storage_bytes."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)
        big = "x" * 110
        with conf_vars({("state_store", "max_value_storage_bytes"): "100"}):
            with patch("airflow.sdk.execution_time.context.log") as mock_log:
                TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", big)
                mock_log.warning.assert_called_once()
                assert "max_value_storage_bytes" in mock_log.warning.call_args[0][0]
        mock_supervisor_comms.send.assert_called_once()

    def test_delete_operation(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).delete("job_id")

        mock_supervisor_comms.send.assert_called_once_with(
            DeleteTaskStateStore(ti_id=self.TI_ID, key="job_id")
        )

    def test_clear_sends_comms_message(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).clear()

        mock_supervisor_comms.send.assert_called_once_with(ClearTaskStateStore(ti_id=self.TI_ID))

    def test_set_datetime_raises_validation_error(self, mock_supervisor_comms):
        """datetime is not JSON-serializable; callers must use .isoformat() first."""
        with pytest.raises(ValidationError):
            TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set(
                "watermark",
                datetime(2026, 5, 15, tzinfo=dt_timezone.utc),
            )

        mock_supervisor_comms.send.assert_not_called()

    def test_set_with_custom_backend_decorates_value_with_marker(self, mock_supervisor_comms):
        """Custom backend ref is wrapped in external Store marker before going to DB."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        backend = MagicMock(spec=BaseStoreBackend)
        backend.serialize_task_state_store_to_ref.return_value = "s3://bucket/ti_123/job_id"

        with (
            patch("airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend),
            conf_vars({("state_store", "default_retention_days"): "0"}),
        ):
            TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "spark_001")

        mock_supervisor_comms.send.assert_called_once_with(
            SetTaskStateStore(
                ti_id=self.TI_ID,
                key="job_id",
                value=_wrap_external_ref("s3://bucket/ti_123/job_id"),
                expires_at=None,
            )
        )

    def test_get_with_custom_backend_removes_decoration_marker(self, mock_supervisor_comms):
        """External Store marker is detected and the ref is passed to deserialize."""
        mock_supervisor_comms.send.return_value = TaskStateStoreResult(
            value=_wrap_external_ref("s3://bucket/ti_123/job_id")
        )

        backend = MagicMock(spec=BaseStoreBackend)
        backend.deserialize_task_state_store_from_ref.return_value = {"rows": 123}

        with patch(
            "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend
        ):
            result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("job_id")

        assert result == {"rows": 123}
        backend.deserialize_task_state_store_from_ref.assert_called_once_with("s3://bucket/ti_123/job_id")


class TestAssetStateStoreAccessor:
    ASSET_NAME = "debug_watcher_asset"
    ASSET_URI = "s3://bucket/key"

    def test_get_returns_value(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="2026-04-30T00:00:00Z")

        result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("watermark")

        assert result == "2026-04-30T00:00:00Z"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_get_returns_none_on_404(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = ErrorResponse(
            error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key": "missing_key"}
        )

        result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("missing_key")

        assert result is None

    def test_get_returns_default_when_key_missing(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = ErrorResponse(
            error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key": "watermark"}
        )

        result = AssetStateStoreAccessor(name=self.ASSET_NAME).get(
            "watermark", default="2026-01-01T00:00:00+00:00"
        )

        assert result == "2026-01-01T00:00:00+00:00"

    def test_get_ignores_default_when_key_exists(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="2026-06-01T00:00:00+00:00")

        result = AssetStateStoreAccessor(name=self.ASSET_NAME).get(
            "watermark", default="2026-01-01T00:00:00+00:00"
        )

        assert result == "2026-06-01T00:00:00+00:00"

    def test_get_raises_on_error(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = ErrorResponse(
            error=ErrorType.GENERIC_ERROR, detail={"message": "server error"}
        )

        with pytest.raises(AirflowRuntimeError):
            AssetStateStoreAccessor(name=self.ASSET_NAME).get("some_key")

    def test_set_operation(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", "2026-04-30T00:00:00Z")

        mock_supervisor_comms.send.assert_called_once_with(
            SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="2026-04-30T00:00:00Z")
        )

    def test_set_none_raises(self, mock_supervisor_comms):
        with pytest.raises(ValueError, match="Cannot set value as None"):
            AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", None)

    def test_delete_operation(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(name=self.ASSET_NAME).delete("watermark")

        mock_supervisor_comms.send.assert_called_once_with(
            DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_clear_operation(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(name=self.ASSET_NAME).clear()

        mock_supervisor_comms.send.assert_called_once_with(ClearAssetStateStoreByName(name=self.ASSET_NAME))

    def test_get_by_uri(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="2026-04-30T00:00:00Z")

        result = AssetStateStoreAccessor(uri=self.ASSET_URI).get("watermark")

        assert result == "2026-04-30T00:00:00Z"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark")
        )

    def test_set_by_uri(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(uri=self.ASSET_URI).set("watermark", "2026-04-30T00:00:00Z")

        mock_supervisor_comms.send.assert_called_once_with(
            SetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark", value="2026-04-30T00:00:00Z")
        )

    def test_delete_by_uri(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(uri=self.ASSET_URI).delete("watermark")

        mock_supervisor_comms.send.assert_called_once_with(
            DeleteAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark")
        )

    def test_clear_by_uri(self, mock_supervisor_comms):
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(uri=self.ASSET_URI).clear()

        mock_supervisor_comms.send.assert_called_once_with(ClearAssetStateStoreByUri(uri=self.ASSET_URI))

    def test_set_with_custom_backend_decorates_value_with_marker(self, mock_supervisor_comms):
        """Custom backend ref is wrapped in external Store marker before going to DB."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        backend = MagicMock(spec=BaseStoreBackend)
        backend.serialize_asset_state_store_to_ref.return_value = "s3://bucket/assets/orders/watermark"

        with patch(
            "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend
        ):
            AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", "2026-05-01")

        mock_supervisor_comms.send.assert_called_once_with(
            SetAssetStateStoreByName(
                name=self.ASSET_NAME,
                key="watermark",
                value=_wrap_external_ref("s3://bucket/assets/orders/watermark"),
            )
        )

    def test_get_with_custom_backend_removes_decoration_marker(self, mock_supervisor_comms):
        """External Store marker is detected and the ref is passed to deserialize."""
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(
            value=_wrap_external_ref("s3://bucket/assets/orders/watermark")
        )

        backend = MagicMock(spec=BaseStoreBackend)
        backend.deserialize_asset_state_store_from_ref.return_value = "2026-05-01"

        with patch(
            "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend
        ):
            result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("watermark")

        assert result == "2026-05-01"
        backend.deserialize_asset_state_store_from_ref.assert_called_once_with(
            "s3://bucket/assets/orders/watermark"
        )

    def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms):
        """set() logs a warning when the serialized value exceeds max_value_storage_bytes."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)
        big = "x" * 110
        with conf_vars({("state_store", "max_value_storage_bytes"): "100"}):
            with patch("airflow.sdk.execution_time.context.log") as mock_log:
                AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", big)
                mock_log.warning.assert_called_once()
                assert "max_value_storage_bytes" in mock_log.warning.call_args[0][0]
        mock_supervisor_comms.send.assert_called_once()


class TestAssetStateStoreAccessors:
    ASSET_NAME = "my_asset"
    ASSET_URI = "s3://bucket/key"

    def test_subscript_by_asset_routes_by_name(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v1")

        result = AssetStateStoreAccessors([asset])[asset].get("watermark")

        assert result == "v1"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_subscript_by_asset_name_ref(self, mock_supervisor_comms):
        ref = AssetNameRef(name=self.ASSET_NAME)
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v2")

        result = AssetStateStoreAccessors([ref])[ref].get("watermark")

        assert result == "v2"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_subscript_by_uri_ref(self, mock_supervisor_comms):
        ref = AssetUriRef(uri=self.ASSET_URI)
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v3")

        result = AssetStateStoreAccessors([ref])[ref].get("watermark")

        assert result == "v3"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark")
        )

    def test_get_single_inlet_simplified(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v4")

        result = AssetStateStoreAccessors([asset]).get("watermark")

        assert result == "v4"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_set_single_inlet_simplified(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessors([asset]).set("watermark", "2026-05-01")

        mock_supervisor_comms.send.assert_called_once_with(
            SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="2026-05-01")
        )

    def test_delete_single_inlet_simplified(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessors([asset]).delete("watermark")

        mock_supervisor_comms.send.assert_called_once_with(
            DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_clear_single_inlet_simplified(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessors([asset]).clear()

        mock_supervisor_comms.send.assert_called_once_with(ClearAssetStateStoreByName(name=self.ASSET_NAME))

    def test_double_reference_raises(self):
        a1 = Asset(name="asset_one", uri="s3://one")
        a2 = Asset(name="asset_two", uri="s3://two")

        with pytest.raises(ValueError, match="2 concrete inlets and outlets"):
            AssetStateStoreAccessors([a1, a2]).get("watermark")

    def test_alias_inlet_resolves_to_concrete_assets(self, mock_supervisor_comms):
        alias = AssetAlias(name="my_alias")
        mock_supervisor_comms.send.return_value = AssetsByAliasResult(
            assets=[AssetResult(name="resolved_asset", uri="s3://bucket/resolved", group="asset")]
        )
        mock_supervisor_comms.send.return_value = AssetsByAliasResult(
            assets=[AssetResult(name="resolved_asset", uri="s3://bucket/resolved", group="asset")]
        )

        accessors = AssetStateStoreAccessors([alias])

        mock_supervisor_comms.send.assert_called_once_with(GetAssetsByAlias(alias_name="my_alias"))
        resolved = Asset(name="resolved_asset", uri="s3://bucket/resolved")
        assert resolved.name in accessors._by_name

    def test_alias_inlet_no_resolved_assets_contributes_nothing(self, mock_supervisor_comms):
        alias = AssetAlias(name="empty_alias")
        mock_supervisor_comms.send.return_value = AssetsByAliasResult(assets=[])

        accessors = AssetStateStoreAccessors([alias])

        assert accessors._total == 0

    def test_outlet_only_asset_is_accessible(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v1")

        result = AssetStateStoreAccessors([], [asset])[asset].get("watermark")

        assert result == "v1"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_outlet_only_name_ref_is_accessible(self, mock_supervisor_comms):
        ref = AssetNameRef(name=self.ASSET_NAME)
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v2")

        result = AssetStateStoreAccessors([], [ref])[ref].get("watermark")

        assert result == "v2"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_outlet_only_uri_ref_is_accessible(self, mock_supervisor_comms):
        ref = AssetUriRef(uri=self.ASSET_URI)
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v2")

        result = AssetStateStoreAccessors([], [ref])[ref].get("watermark")

        assert result == "v2"
        mock_supervisor_comms.send.assert_called_once_with(
            GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark")
        )

    def test_outlet_only_single_shorthand_works(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v3")

        result = AssetStateStoreAccessors([], [asset]).get("watermark")

        assert result == "v3"

    def test_asset_in_both_inlets_and_outlets_not_duplicated(self, mock_supervisor_comms):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")

        accessors = AssetStateStoreAccessors([asset], [asset])

        assert accessors._total == 1

    def test_outlet_alias_is_ignored(self, mock_supervisor_comms):
        alias = AssetAlias(name="my_alias")

        accessors = AssetStateStoreAccessors([], [alias])

        assert accessors._total == 0
        mock_supervisor_comms.send.assert_not_called()


class InMemoryStoreBackend(BaseStoreBackend):
    """Simple in-memory test backend."""

    def __init__(self):
        self._actual_key_value_store: dict[str, str] = {}  # key -> actual value
        self.reference: dict[str, str] = {}  # key -> stored ref (mem:// URI)

    def serialize_task_state_store_to_ref(self, *, value, key: str, scope) -> str:
        ref = f"mem://{scope.dag_id}/{scope.run_id}/{scope.task_id}/{scope.map_index}/{key}"
        self._actual_key_value_store[key] = value
        self.reference[key] = ref
        return ref

    def deserialize_task_state_store_from_ref(self, stored: str) -> JsonValue:
        key = stored.rsplit("/", 1)[-1]
        return self._actual_key_value_store.get(key, stored)

    def serialize_asset_state_store_to_ref(self, *, value, key: str, scope) -> str:
        ref = f"mem://{scope.name or scope.uri}/{key}"
        self._actual_key_value_store[key] = value
        self.reference[key] = ref
        return ref

    def deserialize_asset_state_store_from_ref(self, stored: str) -> JsonValue:
        key = stored.rsplit("/", 1)[-1]
        return self._actual_key_value_store.get(key, stored)

    def get(self, scope, key, *, session=None): ...
    def set(self, scope, key, value, *, session=None): ...

    def delete(self, scope, key, *, session=None) -> None:
        self._actual_key_value_store.pop(key, None)
        self.reference.pop(key, None)

    def clear(self, scope, *, all_map_indices=False, session=None) -> None:
        self._actual_key_value_store.clear()
        self.reference.clear()

    async def aget(self, scope, key): ...
    async def aset(self, scope, key, value): ...
    async def adelete(self, scope, key): ...
    async def aclear(self, scope, *, all_map_indices=False): ...


class TestTaskStateStoreAccessorWithCustomBackend:
    TI_ID = UUID("01900000-0000-0000-0000-000000000002")
    SCOPE = TaskScope(dag_id="dag", run_id="run", task_id="task")

    @pytest.fixture(autouse=True)
    def backend(self):
        b = InMemoryStoreBackend()
        with mock.patch(
            "airflow.sdk.execution_time.context._get_worker_state_store_backend",
            return_value=b,
        ):
            yield b

    def test_set_returns_reference_to_storage(self, mock_supervisor_comms, backend, time_machine):
        """set() stores actual value in backend and sends mem:// reference via comms."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)
        expected_ref = f"mem://{self.SCOPE.dag_id}/{self.SCOPE.run_id}/{self.SCOPE.task_id}/{self.SCOPE.map_index}/job_id"

        frozen_dt = datetime(2026, 1, 1, 12, 0, 0, tzinfo=dt_timezone.utc)
        time_machine.move_to(frozen_dt, tick=False)

        TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001")
        # comms message has the mem:// reference, not the actual value
        mock_supervisor_comms.send.assert_called_once_with(
            SetTaskStateStore(
                ti_id=self.TI_ID,
                key="job_id",
                value=_wrap_external_ref(expected_ref),
                expires_at=frozen_dt + timedelta(days=30),
            )
        )
        # actual value is stored on the backend, reference is stored for DB
        assert backend._actual_key_value_store["job_id"] == "app_001"
        assert backend.reference["job_id"] == expected_ref

    def test_get_resolves_reference_to_actual_value(self, mock_supervisor_comms, backend):
        """get() fetches mem:// reference from DB, resolves it to actual value via backend."""
        ref = _wrap_external_ref(
            f"mem://{self.SCOPE.dag_id}/{self.SCOPE.run_id}/{self.SCOPE.task_id}/{self.SCOPE.map_index}/job_id"
        )
        backend._actual_key_value_store["job_id"] = "app_001"
        mock_supervisor_comms.send.return_value = TaskStateStoreResult(value=ref)

        result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("job_id")
        # actual value is resolved from mem:// reference via backend
        assert result == "app_001"

    def test_deletes_from_backend_and_removes_db_ref(self, mock_supervisor_comms, backend):
        """delete() purges from backend storage and removes the DB reference."""
        backend._actual_key_value_store["job_id"] = "app_001"
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).delete("job_id")

        # backend does not have the value anymore
        assert "job_id" not in backend._actual_key_value_store
        # request to delete reference in DB was made
        mock_supervisor_comms.send.assert_any_call(DeleteTaskStateStore(ti_id=self.TI_ID, key="job_id"))

    def test_clears_all_from_backend_and_clears_db(self, mock_supervisor_comms, backend):
        """clear() purges all backend objects for the TI and removes all DB references."""
        backend._actual_key_value_store["job_id"] = "app_001"
        backend._actual_key_value_store["checkpoint"] = "step_3"
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).clear()

        assert "job_id" not in backend._actual_key_value_store
        assert "checkpoint" not in backend._actual_key_value_store
        mock_supervisor_comms.send.assert_any_call(ClearTaskStateStore(ti_id=self.TI_ID))


class TestAssetStateStoreAccessorWithCustomBackend:
    ASSET_NAME = "my_asset"

    @pytest.fixture(autouse=True)
    def backend(self):
        b = InMemoryStoreBackend()
        with mock.patch(
            "airflow.sdk.execution_time.context._get_worker_state_store_backend",
            return_value=b,
        ):
            yield b

    def test_set_sends_reference_not_value(self, mock_supervisor_comms, backend):
        """set() stores actual value in backend and sends mem:// reference via comms."""
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", "2026-05-01")

        expected_ref = f"mem://{self.ASSET_NAME}/watermark"
        # comms message has the mem:// reference, not the actual value
        mock_supervisor_comms.send.assert_called_once_with(
            SetAssetStateStoreByName(
                name=self.ASSET_NAME,
                key="watermark",
                value=_wrap_external_ref(expected_ref),
            )
        )
        # actual value is stored on the backend, reference is stored for DB
        assert backend._actual_key_value_store["watermark"] == "2026-05-01"
        assert backend.reference["watermark"] == expected_ref

    def test_get_resolves_reference_to_actual_value(self, mock_supervisor_comms, backend):
        """get() fetches mem:// reference from DB, resolves it to actual value via backend."""
        ref = _wrap_external_ref(f"mem://{self.ASSET_NAME}/watermark")
        backend._actual_key_value_store["watermark"] = "2026-05-01"
        mock_supervisor_comms.send.return_value = AssetStateStoreResult(value=ref)

        result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("watermark")

        # actual value is resolved from mem:// reference via backend
        assert result == "2026-05-01"

    def test_delete_purges_from_backend_and_removes_db_ref(self, mock_supervisor_comms, backend):
        """delete() purges from backend storage and removes the DB reference."""
        backend._actual_key_value_store["watermark"] = "2026-05-01"
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(name=self.ASSET_NAME).delete("watermark")

        # backend doesn't have the value anymore
        assert "watermark" not in backend._actual_key_value_store
        # request to delete reference in DB was made
        mock_supervisor_comms.send.assert_any_call(
            DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        )

    def test_clear_purges_all_from_backend_and_clears_db(self, mock_supervisor_comms, backend):
        """clear() purges all backend objects and removes all DB references."""
        backend._actual_key_value_store["watermark"] = "2026-05-01"
        backend._actual_key_value_store["file_count"] = "42"
        mock_supervisor_comms.send.return_value = OKResponse(ok=True)

        AssetStateStoreAccessor(name=self.ASSET_NAME).clear()

        assert "watermark" not in backend._actual_key_value_store
        assert "file_count" not in backend._actual_key_value_store
        mock_supervisor_comms.send.assert_any_call(ClearAssetStateStoreByName(name=self.ASSET_NAME))
