task-sdk/tests/task_sdk/execution_time/test_context.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
from datetime import UTC, datetime, timedelta
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 AssetScope, 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,
    AirflowSecretsBackendAccessDenied,
    ErrorType,
)
from airflow.sdk.execution_time.comms import (
    AssetEventDagRunReferenceResult,
    AssetEventResult,
    AssetEventSourceTaskInstance,
    AssetEventsResult,
    AssetResult,
    AssetsByAliasResult,
    AssetStateStoreResult,
    ClearAssetStateStoreByName,
    ClearAssetStateStoreByUri,
    ClearTaskStateStore,
    ConnectionResult,
    DagRunResult,
    DeleteAssetStateStoreByName,
    DeleteAssetStateStoreByUri,
    DeleteTaskStateStore,
    DeleteVariable,
    ErrorResponse,
    GetAssetByName,
    GetAssetByUri,
    GetAssetEventByAsset,
    GetAssetsByAlias,
    GetAssetStateStoreByName,
    GetAssetStateStoreByUri,
    GetDagRun,
    GetTaskStateStore,
    GetVariableKeys,
    GetXCom,
    OKResponse,
    PutVariable,
    SetAssetStateStoreByName,
    SetAssetStateStoreByUri,
    SetTaskStateStore,
    TaskStateStoreResult,
    VariableKeysResult,
    VariableResult,
    XComResult,
)
from airflow.sdk.execution_time.context import (
    _VARIABLE_KEYS_PAGE_SIZE,
    NEVER_EXPIRE,
    AssetStateStoreAccessor,
    AssetStateStoreAccessors,
    ConnectionAccessor,
    InletEventsAccessors,
    MacrosAccessor,
    OutletEventAccessor,
    OutletEventAccessors,
    TaskStateStoreAccessor,
    TriggeringAssetEventsAccessor,
    VariableAccessor,
    _AssetRefResolutionMixin,
    _async_delete_variable,
    _async_get_connection,
    _async_get_variable,
    _async_get_variable_keys,
    _async_set_variable,
    _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
from tests_common.test_utils.mock_plugins import mock_plugin_manager

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
        )

    def test__get_item__with_extra_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] * 3

        list(sample_inlet_evnets_accessor[TEST_ASSET].extra("region", "us"))
        list(sample_inlet_evnets_accessor[TEST_ASSET].extra("region", "us").extra("env", "prod"))
        list(sample_inlet_evnets_accessor[TEST_ASSET].extra("region", "us").extra("env", "prod").limit(5))

        assert mock_supervisor_comms.send.call_count == 3

        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,
            extra={"region": "us"},
        )
        assert calls[1][0][0] == GetAssetEventByAsset(
            name="test_uri",
            uri="test://test/",
            after=None,
            before=None,
            limit=None,
            ascending=True,
            extra={"region": "us", "env": "prod"},
        )
        assert calls[2][0][0] == GetAssetEventByAsset(
            name="test_uri",
            uri="test://test/",
            after=None,
            before=None,
            limit=5,
            ascending=True,
            extra={"region": "us", "env": "prod"},
        )

    @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()

    @pytest.mark.asyncio
    async def test_async_get_connection_uses_aget_uri_not_get_uri(self, mock_supervisor_comms):
        """_async_get_connection must call aget_uri() when caching, never the sync get_uri().

        get_uri() accesses extra_dejson which calls mask_secret() -> comms.send()
        from the event-loop thread, triggering DeadlockImminentError in Airflow 3.3.1.
        aget_uri() uses amask_secret() -> asend() and is safe in async contexts.
        """
        from airflow.sdk.execution_time.cache import SecretCache

        sample_connection = Connection(
            conn_id="test_conn",
            conn_type="postgres",
            host="localhost",
            port=5432,
            extra='{"sslmode": "require"}',
        )

        class MockSecretsBackend:
            def get_connection(self, conn_id: str) -> Connection | None:
                return sample_connection if conn_id == "test_conn" else None

        with (
            patch(
                "airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", autospec=True
            ) as mock_load,
            mock.patch.object(SecretCache, "save_connection_uri"),
        ):
            mock_load.return_value = [MockSecretsBackend()]

            await _async_get_connection("test_conn")

            # get_uri() would reach the sync mask_secret() -> comms.send(), which deadlocks
            # on the event-loop thread; aget_uri() must go through amask_secret() -> comms.asend().
            mock_supervisor_comms.send.assert_not_called()
            mock_supervisor_comms.asend.assert_awaited()


class TestAsyncVariableContext:
    """Test async variable context functions: _async_get_variable, _async_set_variable, etc."""

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        ("deserialize_json", "value", "expected_value"),
        [
            pytest.param(False, "my_value", "my_value", id="simple-value"),
            pytest.param(
                True,
                '{"key": "value", "number": 42, "flag": true}',
                {"key": "value", "number": 42, "flag": True},
                id="deser-object-value",
            ),
        ],
    )
    async def test_async_get_variable_from_api(
        self, deserialize_json, value, expected_value, mock_supervisor_comms
    ):
        """_async_get_variable fetches from the Execution API via ExecutionAPISecretsBackend (async asend)."""
        mock_supervisor_comms.asend.return_value = VariableResult(key="my_key", value=value)

        result = await _async_get_variable("my_key", deserialize_json=deserialize_json)

        assert result == expected_value
        mock_supervisor_comms.asend.assert_called()

    @pytest.mark.asyncio
    async def test_async_get_variable_from_secrets_backend(self, mock_supervisor_comms):
        """_async_get_variable returns the value from a secrets backend without calling asend."""

        class MockBackend:
            def get_variable(self, key: str):
                return "backend_value"

        with patch(
            "airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", autospec=True
        ) as mock_load:
            mock_load.return_value = [MockBackend()]
            result = await _async_get_variable("my_key", deserialize_json=False)

        assert result == "backend_value"
        # ExecutionAPISecretsBackend (which uses sync send) must not have been called
        mock_supervisor_comms.send.assert_not_called()

    @pytest.mark.asyncio
    async def test_async_get_variable_not_found_raises(self, mock_supervisor_comms):
        """_async_get_variable raises AirflowRuntimeError when the variable does not exist."""
        with patch(
            "airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", autospec=True
        ) as mock_load:
            mock_load.return_value = []

            with pytest.raises(AirflowRuntimeError) as exc_info:
                await _async_get_variable("missing_key", deserialize_json=False)

        assert exc_info.value.error.error == ErrorType.VARIABLE_NOT_FOUND

    @pytest.mark.asyncio
    async def test_async_get_variable_does_not_fall_through_after_deny(self, mock_supervisor_comms):
        """An authoritative deny from the Execution API raises; the next backend is never asked."""
        mock_supervisor_comms.asend.return_value = ErrorResponse(
            error=ErrorType.PERMISSION_DENIED,
            detail={"key": "denied_var", "status_code": 403},
        )

        later_backend = MagicMock(name="LaterBackend")
        # the dispatcher prefers aget_variable when present, so spy on both
        later_backend.aget_variable = mock.AsyncMock(return_value="leaked-value")
        later_backend.get_variable = MagicMock(return_value="leaked-value")

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

            with pytest.raises(AirflowSecretsBackendAccessDenied, match="variable 'denied_var'"):
                await _async_get_variable("denied_var", deserialize_json=False)

        later_backend.aget_variable.assert_not_awaited()
        later_backend.get_variable.assert_not_called()

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_json_masks_from_cache(self, mock_amask_secret, mock_supervisor_comms):
        """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):
            val = await _async_get_variable("db_config", deserialize_json=True)

        assert val == {"password": "s3cr3t", "host": "db.example.com"}
        mock_amask_secret.assert_any_await(raw_json, "db_config")
        mock_amask_secret.assert_any_await({"password": "s3cr3t", "host": "db.example.com"})
        # served from the cache, no backend was asked
        mock_supervisor_comms.asend.assert_not_awaited()

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_value_masks_secret(self, mock_amask_secret, mock_supervisor_comms):
        """A plain value is masked under the variable's key name."""
        mock_supervisor_comms.asend.return_value = VariableResult(key="my_password", value="s3cr3t")

        val = await _async_get_variable("my_password", deserialize_json=False)

        assert val == "s3cr3t"
        mock_amask_secret.assert_awaited_once_with("s3cr3t", "my_password")

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_json_masks_raw_string_and_dict_values(
        self, mock_amask_secret, mock_supervisor_comms
    ):
        """Both the raw JSON string and the deserialized dict's sensitive fields are masked."""
        raw_json = '{"password": "s3cr3t", "host": "db.example.com"}'
        mock_supervisor_comms.asend.return_value = VariableResult(key="db_config", value=raw_json)

        val = await _async_get_variable("db_config", deserialize_json=True)

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

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_json_masks_list_values(self, mock_amask_secret, mock_supervisor_comms):
        """A JSON list is handed to the masker whole, under the variable's key."""
        raw_json = '[{"password": "s3cr3t"}, {"password": "s3cr3t2"}]'
        mock_supervisor_comms.asend.return_value = VariableResult(key="db_configs", value=raw_json)

        val = await _async_get_variable("db_configs", deserialize_json=True)

        assert val == [{"password": "s3cr3t"}, {"password": "s3cr3t2"}]
        mock_amask_secret.assert_any_await(raw_json, "db_configs")
        mock_amask_secret.assert_any_await([{"password": "s3cr3t"}, {"password": "s3cr3t2"}], "db_configs")

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_json_sensitive_key_masks_raw_json(
        self, mock_amask_secret, mock_supervisor_comms
    ):
        """A sensitive variable key masks the entire raw JSON string."""
        raw_json = '{"endpoint": "https://api.example.com", "token": "abc123"}'
        mock_supervisor_comms.asend.return_value = VariableResult(key="my_secret", value=raw_json)

        val = await _async_get_variable("my_secret", deserialize_json=True)

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

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_json_string_value_masks_both_forms(
        self, mock_amask_secret, mock_supervisor_comms
    ):
        """A JSON string value masks both the quoted raw and the unquoted value."""
        mock_supervisor_comms.asend.return_value = VariableResult(key="my_token", value='"s3cr3t"')

        val = await _async_get_variable("my_token", deserialize_json=True)

        assert val == "s3cr3t"
        assert mock_amask_secret.await_count == 2
        mock_amask_secret.assert_any_await('"s3cr3t"', "my_token")
        mock_amask_secret.assert_any_await("s3cr3t", "my_token")

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_json_list_value_does_not_over_mask(
        self, mock_amask_secret, mock_supervisor_comms
    ):
        """A non-sensitive list variable is never masked anonymously."""
        raw_json = '["us-east-1", "eu-west-1"]'
        mock_supervisor_comms.asend.return_value = VariableResult(key="aws_regions", value=raw_json)

        val = await _async_get_variable("aws_regions", deserialize_json=True)

        assert val == ["us-east-1", "eu-west-1"]
        mock_amask_secret.assert_any_await(raw_json, "aws_regions")
        mock_amask_secret.assert_any_await(["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_amask_secret.await_args_list

    @pytest.mark.asyncio
    @mock.patch("airflow.sdk.execution_time.context.amask_secret")
    async def test_async_var_json_invalid_json_raises(self, mock_amask_secret, mock_supervisor_comms):
        """Invalid JSON raises JSONDecodeError; the raw value is still masked before the error."""
        from airflow.sdk.execution_time.cache import SecretCache

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

        mock_amask_secret.assert_awaited_once_with(raw, "bad_var")

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        ("key", "value", "description", "serialize_json"),
        [
            pytest.param("key", "value", "description", False, id="simple-value"),
            pytest.param(
                "key2",
                {"hi": "there", "hello": 42, "flag": True},
                "description2",
                True,
                id="serialize-json-value",
            ),
        ],
    )
    async def test_async_set_variable(self, key, value, description, serialize_json, mock_supervisor_comms):
        """_async_set_variable sends PutVariable via asend."""
        mock_supervisor_comms.asend.return_value = None

        with patch(
            "airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", autospec=True
        ) as mock_load:
            mock_load.return_value = []
            await _async_set_variable(key, value, description, serialize_json=serialize_json)

        expected_value = json.dumps(value, indent=2) if serialize_json else value
        mock_supervisor_comms.asend.assert_called_once_with(
            PutVariable(key=key, value=expected_value, description=description)
        )

    @pytest.mark.asyncio
    async def test_async_set_variable_warns_on_conflicting_backend(self, mock_supervisor_comms):
        """A worker-side backend that already holds the key gets a warning; the write still goes through."""
        mock_supervisor_comms.asend.return_value = None

        class ConflictingBackend:
            def get_variable(self, key: str):
                return "backend_value"

        class EmptyBackend:
            def get_variable(self, key: str):
                return None

        class FailingBackend:
            def get_variable(self, key: str):
                raise RuntimeError("backend down")

        execution_api_backend = mock.create_autospec(ExecutionAPISecretsBackend, instance=True)

        with (
            patch(
                "airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", autospec=True
            ) as mock_load,
            patch("airflow.sdk.execution_time.context.log") as mock_log,
        ):
            mock_load.return_value = [
                ConflictingBackend(),
                EmptyBackend(),
                FailingBackend(),
                execution_api_backend,
            ]
            await _async_set_variable("my_key", "new_value")

        mock_log.warning.assert_called_once()
        warning_args = mock_log.warning.call_args.args
        assert warning_args[1:] == ("my_key", "ConflictingBackend", "ConflictingBackend")
        mock_log.exception.assert_called_once()
        assert mock_log.exception.call_args.args[1] == "FailingBackend"
        # the Execution API is the write target, not a conflict
        execution_api_backend.get_variable.assert_not_called()
        mock_supervisor_comms.asend.assert_called_once_with(
            PutVariable(key="my_key", value="new_value", description=None)
        )

    @pytest.mark.asyncio
    async def test_async_delete_variable(self, mock_supervisor_comms):
        """_async_delete_variable sends DeleteVariable via asend."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)

        await _async_delete_variable("my_key")

        mock_supervisor_comms.asend.assert_called_once_with(DeleteVariable(key="my_key"))

    @pytest.mark.asyncio
    @pytest.mark.parametrize(
        ("prefix", "keys"),
        [
            pytest.param(None, ["prod_db", "prod_api", "dev_debug"], id="all"),
            pytest.param("prod_", ["prod_db", "prod_api"], id="with-prefix"),
            pytest.param("nonexistent_", [], id="empty-result"),
        ],
    )
    async def test_async_get_variable_keys(self, prefix, keys, mock_supervisor_comms):
        """_async_get_variable_keys fetches all keys matching the prefix in one page."""
        mock_supervisor_comms.asend.return_value = VariableKeysResult(keys=keys, total_entries=len(keys))

        result = await _async_get_variable_keys(prefix=prefix)

        assert result == keys
        mock_supervisor_comms.asend.assert_called_once_with(
            GetVariableKeys(prefix=prefix, limit=1000, offset=0)
        )

    @pytest.mark.asyncio
    async def test_async_get_variable_keys_paginates(self, mock_supervisor_comms):
        """_async_get_variable_keys accumulates results across multiple pages."""
        page1 = [f"k{i}" for i in range(_VARIABLE_KEYS_PAGE_SIZE)]
        page2 = ["last_key"]
        mock_supervisor_comms.asend.side_effect = [
            VariableKeysResult(keys=page1, total_entries=_VARIABLE_KEYS_PAGE_SIZE + 1),
            VariableKeysResult(keys=page2, total_entries=_VARIABLE_KEYS_PAGE_SIZE + 1),
        ]

        result = await _async_get_variable_keys(prefix=None)

        assert result == page1 + page2
        assert mock_supervisor_comms.asend.call_count == 2

    @pytest.mark.asyncio
    async def test_async_get_variable_keys_raises_on_error(self, mock_supervisor_comms):
        """_async_get_variable_keys raises AirflowRuntimeError on an ErrorResponse."""
        mock_supervisor_comms.asend.return_value = ErrorResponse(
            error=ErrorType.GENERIC_ERROR, detail={"message": "boom"}
        )

        with pytest.raises(AirflowRuntimeError):
            await _async_get_variable_keys(prefix="x_")

    @pytest.mark.asyncio
    async def test_async_get_variable_keys_raises_on_unexpected_response(self, mock_supervisor_comms):
        """_async_get_variable_keys raises TypeError for an unrecognised response type."""
        mock_supervisor_comms.asend.return_value = VariableResult(key="x", value="y")

        with pytest.raises(TypeError, match="Unexpected response type"):
            await _async_get_variable_keys(prefix="x_")


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=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=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=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=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=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")

    @pytest.mark.asyncio
    async def test_aget_returns_value(self, mock_supervisor_comms):
        """aget awaits asend and returns the stored value, without touching sync send."""
        mock_supervisor_comms.asend.return_value = TaskStateStoreResult(value="app_001")

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

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

    @pytest.mark.asyncio
    async def test_aget_returns_default_when_key_missing(self, mock_supervisor_comms):
        mock_supervisor_comms.asend.return_value = ErrorResponse(
            error=ErrorType.TASK_STORE_NOT_FOUND, detail={"key": "job_id"}
        )

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

        assert result == "default-id"

    @pytest.mark.asyncio
    async def test_aget_raises_on_error(self, mock_supervisor_comms):
        mock_supervisor_comms.asend.return_value = ErrorResponse(
            error=ErrorType.GENERIC_ERROR, detail={"message": "server error"}
        )

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

    @pytest.mark.asyncio
    async def test_aget_with_custom_backend_removes_decoration_marker(self, mock_supervisor_comms):
        """aget unwraps the external Store marker and resolves the ref via the backend."""
        mock_supervisor_comms.asend.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 = await TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).aget("job_id")

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

    @pytest.mark.asyncio
    async def test_aset_with_global_retention(self, mock_supervisor_comms, time_machine):
        """aset awaits asend with the message built from the global retention config."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
        now = datetime(2026, 5, 14, 12, 0, 0, tzinfo=UTC)
        time_machine.move_to(now, tick=False)

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

        mock_supervisor_comms.asend.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=UTC),
            )
        )
        mock_supervisor_comms.send.assert_not_called()

    @pytest.mark.asyncio
    async def test_aset_none_raises(self, mock_supervisor_comms):
        with pytest.raises(ValueError, match="Cannot set value as None"):
            await TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).aset("job_id", None)

        mock_supervisor_comms.asend.assert_not_called()

    @pytest.mark.asyncio
    async def test_aset_with_custom_backend_decorates_value_with_marker(self, mock_supervisor_comms):
        """aset wraps the custom backend ref in the external Store marker before sending."""
        mock_supervisor_comms.asend.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"}),
        ):
            await TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).aset("job_id", "spark_001")

        mock_supervisor_comms.asend.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,
            )
        )

    @pytest.mark.asyncio
    async def test_adelete_awaits_asend(self, mock_supervisor_comms):
        """adelete awaits asend without touching sync send."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)

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

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

    @pytest.mark.asyncio
    async def test_aclear_awaits_asend(self, mock_supervisor_comms):
        """aclear awaits asend without touching sync send."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)

        await TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).aclear()

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

    @pytest.mark.asyncio
    async def test_adelete_purges_via_async_backend(self, mock_supervisor_comms):
        """adelete awaits the async backend instead of blocking on the sync delete."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
        backend = MagicMock(spec=BaseStoreBackend)

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

        backend.adelete.assert_awaited_once_with(self.SCOPE, "job_id")
        backend.delete.assert_not_called()

    @pytest.mark.asyncio
    async def test_aclear_purges_via_async_backend(self, mock_supervisor_comms):
        """aclear awaits the async backend instead of blocking on the sync clear."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
        backend = MagicMock(spec=BaseStoreBackend)

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

        backend.aclear.assert_awaited_once_with(self.SCOPE)
        backend.clear.assert_not_called()


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()

    @pytest.mark.parametrize("lookup_by", ["name", "uri"])
    @pytest.mark.asyncio
    async def test_aget_returns_value(self, mock_supervisor_comms, lookup_by):
        """aget awaits asend and returns the stored value, without touching sync send."""
        mock_supervisor_comms.asend.return_value = AssetStateStoreResult(value="2026-04-30T00:00:00Z")
        if lookup_by == "name":
            accessor = AssetStateStoreAccessor(name=self.ASSET_NAME)
            expected_message = GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark")
        else:
            accessor = AssetStateStoreAccessor(uri=self.ASSET_URI)
            expected_message = GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark")

        result = await accessor.aget("watermark")

        assert result == "2026-04-30T00:00:00Z"
        mock_supervisor_comms.asend.assert_called_once_with(expected_message)
        mock_supervisor_comms.send.assert_not_called()

    @pytest.mark.asyncio
    async def test_aget_returns_default_when_key_missing(self, mock_supervisor_comms):
        mock_supervisor_comms.asend.return_value = ErrorResponse(
            error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key": "watermark"}
        )

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

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

    @pytest.mark.asyncio
    async def test_aget_raises_on_error(self, mock_supervisor_comms):
        mock_supervisor_comms.asend.return_value = ErrorResponse(
            error=ErrorType.GENERIC_ERROR, detail={"message": "server error"}
        )

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

    @pytest.mark.asyncio
    async def test_aget_with_custom_backend_removes_decoration_marker(self, mock_supervisor_comms):
        """aget unwraps the external Store marker and resolves the ref via the backend."""
        mock_supervisor_comms.asend.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 = await AssetStateStoreAccessor(name=self.ASSET_NAME).aget("watermark")

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

    @pytest.mark.parametrize("lookup_by", ["name", "uri"])
    @pytest.mark.asyncio
    async def test_aset_operation(self, mock_supervisor_comms, lookup_by):
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
        if lookup_by == "name":
            accessor = AssetStateStoreAccessor(name=self.ASSET_NAME)
            expected_message = SetAssetStateStoreByName(
                name=self.ASSET_NAME, key="watermark", value="2026-04-30T00:00:00Z"
            )
        else:
            accessor = AssetStateStoreAccessor(uri=self.ASSET_URI)
            expected_message = SetAssetStateStoreByUri(
                uri=self.ASSET_URI, key="watermark", value="2026-04-30T00:00:00Z"
            )

        await accessor.aset("watermark", "2026-04-30T00:00:00Z")

        mock_supervisor_comms.asend.assert_called_once_with(expected_message)
        mock_supervisor_comms.send.assert_not_called()

    @pytest.mark.asyncio
    async def test_aset_none_raises(self, mock_supervisor_comms):
        with pytest.raises(ValueError, match="Cannot set value as None"):
            await AssetStateStoreAccessor(name=self.ASSET_NAME).aset("watermark", None)

        mock_supervisor_comms.asend.assert_not_called()

    @pytest.mark.asyncio
    async def test_aset_with_custom_backend_decorates_value_with_marker(self, mock_supervisor_comms):
        """aset wraps the custom backend ref in the external Store marker before sending."""
        mock_supervisor_comms.asend.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
        ):
            await AssetStateStoreAccessor(name=self.ASSET_NAME).aset("watermark", "2026-05-01")

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

    @pytest.mark.asyncio
    async def test_adelete_awaits_asend(self, mock_supervisor_comms):
        """adelete awaits asend without touching sync send."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)

        await AssetStateStoreAccessor(name=self.ASSET_NAME).adelete("watermark")

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

    @pytest.mark.asyncio
    async def test_aclear_awaits_asend(self, mock_supervisor_comms):
        """aclear awaits asend without touching sync send."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)

        await AssetStateStoreAccessor(uri=self.ASSET_URI).aclear()

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

    @pytest.mark.asyncio
    async def test_adelete_purges_via_async_backend(self, mock_supervisor_comms):
        """adelete awaits the async backend instead of blocking on the sync delete."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
        backend = MagicMock(spec=BaseStoreBackend)

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

        backend.adelete.assert_awaited_once_with(AssetScope(name=self.ASSET_NAME, uri=None), "watermark")
        backend.delete.assert_not_called()

    @pytest.mark.asyncio
    async def test_aclear_purges_via_async_backend(self, mock_supervisor_comms):
        """aclear awaits the async backend instead of blocking on the sync clear."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)
        backend = MagicMock(spec=BaseStoreBackend)

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

        backend.aclear.assert_awaited_once_with(AssetScope(name=self.ASSET_NAME, uri=None))
        backend.clear.assert_not_called()


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()

    @pytest.mark.parametrize(
        ("method_name", "expected_message"),
        [
            ("aget", GetAssetStateStoreByName(name=ASSET_NAME, key="watermark")),
            ("aset", SetAssetStateStoreByName(name=ASSET_NAME, key="watermark", value="2026-05-01")),
            ("adelete", DeleteAssetStateStoreByName(name=ASSET_NAME, key="watermark")),
            ("aclear", ClearAssetStateStoreByName(name=ASSET_NAME)),
        ],
    )
    @pytest.mark.asyncio
    async def test_single_inlet_async_shorthand(self, mock_supervisor_comms, method_name, expected_message):
        asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}")
        mock_supervisor_comms.asend.return_value = (
            AssetStateStoreResult(value="v5") if method_name == "aget" else OKResponse(ok=True)
        )
        accessors = AssetStateStoreAccessors([asset])

        if method_name == "aget":
            result = await accessors.aget("watermark")
            assert result == "v5"
        elif method_name == "aset":
            await accessors.aset("watermark", "2026-05-01")
        elif method_name == "adelete":
            await accessors.adelete("watermark")
        else:
            await accessors.aclear()

        mock_supervisor_comms.asend.assert_called_once_with(expected_message)

    @pytest.mark.asyncio
    async def test_double_reference_raises_for_async_accessor(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"):
            await AssetStateStoreAccessors([a1, a2]).aget("watermark")


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=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))

    @pytest.mark.asyncio
    async def test_aset_returns_reference_to_storage(self, mock_supervisor_comms, backend, time_machine):
        """aset() stores actual value in backend and sends mem:// reference via comms."""
        mock_supervisor_comms.asend.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=UTC)
        time_machine.move_to(frozen_dt, tick=False)

        await TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).aset("job_id", "app_001")

        mock_supervisor_comms.asend.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),
            )
        )
        assert backend._actual_key_value_store["job_id"] == "app_001"
        assert backend.reference["job_id"] == expected_ref

    @pytest.mark.asyncio
    async def test_aget_resolves_reference_to_actual_value(self, mock_supervisor_comms, backend):
        """aget() 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.asend.return_value = TaskStateStoreResult(value=ref)

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

        assert result == "app_001"


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))

    @pytest.mark.asyncio
    async def test_aset_sends_reference_not_value(self, mock_supervisor_comms, backend):
        """aset() stores actual value in backend and sends mem:// reference via comms."""
        mock_supervisor_comms.asend.return_value = OKResponse(ok=True)

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

        expected_ref = f"mem://{self.ASSET_NAME}/watermark"
        mock_supervisor_comms.asend.assert_called_once_with(
            SetAssetStateStoreByName(
                name=self.ASSET_NAME,
                key="watermark",
                value=_wrap_external_ref(expected_ref),
            )
        )
        assert backend._actual_key_value_store["watermark"] == "2026-05-01"
        assert backend.reference["watermark"] == expected_ref

    @pytest.mark.asyncio
    async def test_aget_resolves_reference_to_actual_value(self, mock_supervisor_comms, backend):
        """aget() 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.asend.return_value = AssetStateStoreResult(value=ref)

        result = await AssetStateStoreAccessor(name=self.ASSET_NAME).aget("watermark")

        assert result == "2026-05-01"


class TestMacrosAccessorTeamScoping:
    """A task may use its own team's and the global plugins' macros, but not another team's."""

    @staticmethod
    def _plugins():
        from airflow.sdk.plugins_manager import AirflowPlugin

        def team_a_macro():
            return "team-a"

        def shared_macro():
            return "shared"

        class TeamAPlugin(AirflowPlugin):
            name = "team_a_macros"
            team_name = "team-a"
            macros = [team_a_macro]

        class GlobalPlugin(AirflowPlugin):
            name = "global_macros"
            macros = [shared_macro]

        return [TeamAPlugin, GlobalPlugin]

    @pytest.fixture
    def integrated_macros(self):
        from airflow.sdk.plugins_manager import integrate_macros_plugins

        with mock_plugin_manager(plugins=self._plugins()):
            integrate_macros_plugins()
            yield

    @pytest.mark.parametrize(
        ("team_name", "reachable"),
        [
            pytest.param("team-a", True, id="owning-team"),
            pytest.param("team-b", False, id="other-team"),
            pytest.param(None, False, id="teamless-task"),
        ],
    )
    def test_team_macros_reachable_only_by_their_team(self, integrated_macros, team_name, reachable):
        accessor = MacrosAccessor(team_name=team_name, multi_team=True)

        if reachable:
            assert accessor.team_a_macros.team_a_macro() == "team-a"
        else:
            with pytest.raises(AttributeError, match="belong to team 'team-a'"):
                accessor.team_a_macros

    @pytest.mark.parametrize(
        "team_name",
        [pytest.param("team-a", id="team-task"), pytest.param(None, id="teamless-task")],
    )
    def test_global_plugin_macros_stay_reachable(self, integrated_macros, team_name):
        accessor = MacrosAccessor(team_name=team_name, multi_team=True)

        assert accessor.global_macros.shared_macro() == "shared"

    def test_builtin_macros_stay_reachable(self, integrated_macros):
        """Only plugin submodules are scoped; the macros module's own contents are not."""
        accessor = MacrosAccessor(team_name="team-b", multi_team=True)

        assert accessor.ds_add("2026-01-01", 1) == "2026-01-02"

    def test_nothing_is_hidden_when_multi_team_is_off(self, integrated_macros):
        """A plugin declaring a team in a single-team deployment keeps working as before."""
        accessor = MacrosAccessor()

        assert accessor.team_a_macros.team_a_macro() == "team-a"
        assert accessor.global_macros.shared_macro() == "shared"