| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| from __future__ import annotations |
| |
| from datetime import datetime, timedelta, timezone as dt_timezone |
| from typing import TYPE_CHECKING |
| from unittest import mock |
| from unittest.mock import MagicMock, patch |
| from uuid import UUID |
| |
| import pytest |
| from pydantic import ValidationError |
| |
| from airflow.sdk import BaseOperator, get_current_context, timezone |
| from airflow.sdk._shared.state import TaskScope |
| from airflow.sdk.api.datamodels._generated import ( |
| AssetEventResponse, |
| AssetResponse, |
| DagRun, |
| ) |
| from airflow.sdk.bases.xcom import BaseXCom |
| from airflow.sdk.definitions.asset import ( |
| Asset, |
| AssetAlias, |
| AssetAliasEvent, |
| AssetAliasUniqueKey, |
| AssetNameRef, |
| AssetUniqueKey, |
| AssetUriRef, |
| ) |
| from airflow.sdk.definitions.connection import Connection |
| from airflow.sdk.definitions.variable import Variable |
| from airflow.sdk.exceptions import AirflowNotFoundException, AirflowRuntimeError, ErrorType |
| from airflow.sdk.execution_time.comms import ( |
| AssetEventDagRunReferenceResult, |
| AssetEventResult, |
| AssetEventSourceTaskInstance, |
| AssetEventsResult, |
| AssetResult, |
| AssetsByAliasResult, |
| AssetStateStoreResult, |
| ClearAssetStateStoreByName, |
| ClearAssetStateStoreByUri, |
| ClearTaskStateStore, |
| ConnectionResult, |
| DagRunResult, |
| DeleteAssetStateStoreByName, |
| DeleteAssetStateStoreByUri, |
| DeleteTaskStateStore, |
| ErrorResponse, |
| GetAssetByName, |
| GetAssetByUri, |
| GetAssetEventByAsset, |
| GetAssetsByAlias, |
| GetAssetStateStoreByName, |
| GetAssetStateStoreByUri, |
| GetDagRun, |
| GetTaskStateStore, |
| GetXCom, |
| OKResponse, |
| SetAssetStateStoreByName, |
| SetAssetStateStoreByUri, |
| SetTaskStateStore, |
| TaskStateStoreResult, |
| VariableResult, |
| XComResult, |
| ) |
| from airflow.sdk.execution_time.context import ( |
| NEVER_EXPIRE, |
| AssetStateStoreAccessor, |
| AssetStateStoreAccessors, |
| ConnectionAccessor, |
| InletEventsAccessors, |
| OutletEventAccessor, |
| OutletEventAccessors, |
| TaskStateStoreAccessor, |
| TriggeringAssetEventsAccessor, |
| VariableAccessor, |
| _AssetRefResolutionMixin, |
| _async_get_connection, |
| _convert_variable_result_to_variable, |
| _get_connection, |
| _process_connection_result_conn, |
| _wrap_external_ref, |
| context_to_airflow_vars, |
| set_current_context, |
| ) |
| from airflow.sdk.execution_time.secrets import ExecutionAPISecretsBackend |
| from airflow.sdk.state import BaseStoreBackend |
| |
| from tests_common.test_utils.config import conf_vars |
| |
| if TYPE_CHECKING: |
| from pydantic import JsonValue |
| |
| |
| def test_convert_connection_result_conn(): |
| """Test that the ConnectionResult is converted to a Connection object.""" |
| conn = ConnectionResult( |
| conn_id="test_conn", |
| conn_type="mysql", |
| host="mysql", |
| schema="airflow", |
| login="root", |
| password="password", |
| port=1234, |
| extra='{"extra_key": "extra_value"}', |
| ) |
| conn = _process_connection_result_conn(conn) |
| assert conn == Connection( |
| conn_id="test_conn", |
| conn_type="mysql", |
| host="mysql", |
| schema="airflow", |
| login="root", |
| password="password", |
| port=1234, |
| extra='{"extra_key": "extra_value"}', |
| ) |
| |
| |
| def test_convert_variable_result_to_variable(): |
| """Test that the VariableResult is converted to a Variable object.""" |
| var = VariableResult( |
| key="test_key", |
| value="test_value", |
| ) |
| var = _convert_variable_result_to_variable(var, deserialize_json=False) |
| assert var == Variable( |
| key="test_key", |
| value="test_value", |
| ) |
| |
| |
| def test_convert_variable_result_to_variable_with_deserialize_json(): |
| """Test that the VariableResult is converted to a Variable object with deserialize_json set to True.""" |
| var = VariableResult( |
| key="test_key", |
| value='{\r\n "key1": "value1",\r\n "key2": "value2",\r\n "enabled": true,\r\n "threshold": 42\r\n}', |
| ) |
| var = _convert_variable_result_to_variable(var, deserialize_json=True) |
| assert var == Variable( |
| key="test_key", value={"key1": "value1", "key2": "value2", "enabled": True, "threshold": 42} |
| ) |
| |
| |
| class TestAirflowContextHelpers: |
| def test_context_to_airflow_vars_empty_context(self): |
| assert context_to_airflow_vars({}) == {} |
| |
| def test_context_to_airflow_vars_all_context(self, create_runtime_ti): |
| task = BaseOperator( |
| task_id="test_context_vars", |
| owner=["owner1", "owner2"], |
| email="email1@test.com", |
| ) |
| |
| rti = create_runtime_ti( |
| task=task, |
| dag_id="dag_id", |
| run_id="dag_run_id", |
| logical_date="2017-05-21T00:00:00Z", |
| try_number=1, |
| ) |
| context = rti.get_template_context() |
| assert context_to_airflow_vars(context) == { |
| "airflow.ctx.dag_id": "dag_id", |
| "airflow.ctx.logical_date": "2017-05-21T00:00:00+00:00", |
| "airflow.ctx.task_id": "test_context_vars", |
| "airflow.ctx.dag_run_id": "dag_run_id", |
| "airflow.ctx.try_number": "1", |
| "airflow.ctx.dag_owner": "owner1,owner2", |
| "airflow.ctx.dag_email": "email1@test.com", |
| } |
| |
| assert context_to_airflow_vars(context, in_env_var_format=True) == { |
| "AIRFLOW_CTX_DAG_ID": "dag_id", |
| "AIRFLOW_CTX_LOGICAL_DATE": "2017-05-21T00:00:00+00:00", |
| "AIRFLOW_CTX_TASK_ID": "test_context_vars", |
| "AIRFLOW_CTX_TRY_NUMBER": "1", |
| "AIRFLOW_CTX_DAG_RUN_ID": "dag_run_id", |
| "AIRFLOW_CTX_DAG_OWNER": "owner1,owner2", |
| "AIRFLOW_CTX_DAG_EMAIL": "email1@test.com", |
| } |
| |
| def test_context_to_airflow_vars_team_name(self, create_runtime_ti): |
| """``team_name`` on dag_run surfaces as AIRFLOW_CTX_TEAM_NAME when set; omitted when None.""" |
| task = BaseOperator(task_id="task") |
| rti = create_runtime_ti(task=task) |
| context = rti.get_template_context() |
| |
| # Default (team_name is None) -> key not present |
| assert "AIRFLOW_CTX_TEAM_NAME" not in context_to_airflow_vars(context, in_env_var_format=True) |
| |
| context["dag_run"].team_name = "team-a" |
| env_vars = context_to_airflow_vars(context, in_env_var_format=True) |
| assert env_vars["AIRFLOW_CTX_TEAM_NAME"] == "team-a" |
| assert context_to_airflow_vars(context)["airflow.ctx.team_name"] == "team-a" |
| |
| def test_context_to_airflow_vars_from_policy(self): |
| with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: |
| airflow_cluster = "cluster-a" |
| mock_method.return_value = {"airflow_cluster": airflow_cluster} |
| |
| context_vars = context_to_airflow_vars({}) |
| assert context_vars["airflow.ctx.airflow_cluster"] == airflow_cluster |
| |
| context_vars = context_to_airflow_vars({}, in_env_var_format=True) |
| assert context_vars["AIRFLOW_CTX_AIRFLOW_CLUSTER"] == airflow_cluster |
| |
| with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: |
| mock_method.return_value = {"airflow_cluster": [1, 2]} |
| with pytest.raises(TypeError) as error: |
| context_to_airflow_vars({}) |
| assert str(error.value) == "value of key <airflow_cluster> must be string, not <class 'list'>" |
| |
| with mock.patch("airflow.settings.get_airflow_context_vars") as mock_method: |
| mock_method.return_value = {1: "value"} |
| with pytest.raises(TypeError) as error: |
| context_to_airflow_vars({}) |
| assert str(error.value) == "key <1> must be string" |
| |
| |
| class TestConnectionAccessor: |
| def test_getattr_connection(self, mock_supervisor_comms): |
| """ |
| Test that the connection is fetched when accessed via __getattr__. |
| |
| The __getattr__ method is used for template rendering. Example: ``{{ conn.mysql_conn.host }}``. |
| """ |
| accessor = ConnectionAccessor() |
| |
| # Conn from the supervisor / API Server |
| conn_result = ConnectionResult(conn_id="mysql_conn", conn_type="mysql", host="mysql", port=3306) |
| |
| 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) |
| |
| 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"}', |
| ) |
| |
| 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!" |
| ) |
| |
| 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" |
| |
| |
| 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() |
| |
| |
| class TestTriggeringAssetEventsAccessor: |
| @pytest.fixture(autouse=True) |
| def clear_cache(self): |
| _AssetRefResolutionMixin._asset_ref_cache = {} |
| yield |
| _AssetRefResolutionMixin._asset_ref_cache = {} |
| |
| @pytest.fixture |
| def event_data(self): |
| return [ |
| { |
| "asset": { |
| "name": "1", |
| "uri": "1", |
| "extra": {}, |
| }, |
| "extra": {}, |
| "source_task_id": "t1", |
| "source_dag_id": "d1", |
| "source_run_id": "r1", |
| "source_map_index": -1, |
| "source_aliases": [], |
| "timestamp": "2025-01-01T00:00:12Z", |
| }, |
| { |
| "asset": { |
| "name": "1", |
| "uri": "1", |
| "extra": {}, |
| }, |
| "extra": {}, |
| "source_task_id": "t2", |
| "source_dag_id": "d1", |
| "source_run_id": "r1", |
| "source_map_index": -1, |
| "source_aliases": [ |
| {"name": "a"}, |
| {"name": "b"}, |
| ], |
| "timestamp": "2025-01-01T00:05:43Z", |
| }, |
| { |
| "asset": { |
| "name": "2", |
| "uri": "2", |
| "extra": {}, |
| }, |
| "extra": {}, |
| "source_task_id": "t2", |
| "source_dag_id": "d1", |
| "source_run_id": "r1", |
| "source_map_index": -1, |
| "source_aliases": [], |
| "timestamp": "2025-01-01T00:06:07Z", |
| }, |
| ] |
| |
| @pytest.fixture |
| def accessor(self, event_data): |
| return TriggeringAssetEventsAccessor.build( |
| AssetEventDagRunReferenceResult.model_validate(d) for d in event_data |
| ) |
| |
| @pytest.mark.parametrize( |
| ("key", "result_indexes"), |
| [ |
| (Asset("1"), [0, 1]), |
| (Asset("2"), [2]), |
| (AssetAlias("a"), [1]), |
| (AssetAlias("b"), [1]), |
| ], |
| ) |
| def test_getitem(self, event_data, accessor, key, result_indexes): |
| expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes] |
| assert accessor[key] == expected |
| |
| @pytest.mark.parametrize( |
| ("name", "resolved_asset", "result_indexes"), |
| [ |
| ("1", AssetResult(name="1", uri="1", group="whatever"), [0, 1]), |
| ("2", AssetResult(name="2", uri="2", group="whatever"), [2]), |
| ], |
| ) |
| def test_getitem_name_ref( |
| self, |
| mock_supervisor_comms, |
| event_data, |
| accessor, |
| name, |
| resolved_asset, |
| result_indexes, |
| ): |
| mock_supervisor_comms.send.return_value = resolved_asset |
| expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes] |
| assert accessor[Asset.ref(name=name)] == expected |
| |
| assert mock_supervisor_comms.send.mock_calls == [ |
| mock.call(GetAssetByName(name=name, type="GetAssetByName")) |
| ] |
| assert _AssetRefResolutionMixin._asset_ref_cache |
| |
| @pytest.mark.parametrize( |
| ("uri", "resolved_asset", "result_indexes"), |
| [ |
| ("1", AssetResult(name="1", uri="1", group="whatever"), [0, 1]), |
| ("2", AssetResult(name="2", uri="2", group="whatever"), [2]), |
| ], |
| ) |
| def test_getitem_uri_ref( |
| self, |
| mock_supervisor_comms, |
| event_data, |
| accessor, |
| uri, |
| resolved_asset, |
| result_indexes, |
| ): |
| mock_supervisor_comms.send.return_value = resolved_asset |
| expected = [AssetEventDagRunReferenceResult.model_validate(event_data[i]) for i in result_indexes] |
| assert accessor[Asset.ref(uri=uri)] == expected |
| assert mock_supervisor_comms.send.mock_calls == [mock.call(GetAssetByUri(uri=uri))] |
| assert _AssetRefResolutionMixin._asset_ref_cache |
| |
| def test_partition_key_exposed(self): |
| """A consumed asset event's partition key is reachable via triggering_asset_events.""" |
| event = { |
| "asset": {"name": "1", "uri": "1", "extra": {}}, |
| "extra": {}, |
| "source_task_id": "t1", |
| "source_dag_id": "d1", |
| "source_run_id": "r1", |
| "source_map_index": -1, |
| "source_aliases": [], |
| "timestamp": "2025-01-01T00:00:12Z", |
| "partition_key": "2024-01-15", |
| } |
| accessor = TriggeringAssetEventsAccessor.build( |
| [AssetEventDagRunReferenceResult.model_validate(event)] |
| ) |
| assert [e.partition_key for e in accessor[Asset("1")]] == ["2024-01-15"] |
| |
| def test_source_task_instance_xcom_pull(self, mock_supervisor_comms, accessor): |
| events = accessor[Asset("2")] |
| assert len(events) == 1 |
| |
| mock_dag_run = mock.Mock(dag_id="d1", run_id="r1") |
| mock_supervisor_comms.send.side_effect = [mock_dag_run] |
| source = events[0].source_task_instance |
| assert source == AssetEventSourceTaskInstance(dag_run=mock_dag_run, task_id="t2", map_index=-1) |
| assert mock_supervisor_comms.send.mock_calls == [mock.call(GetDagRun(dag_id="d1", run_id="r1"))] |
| |
| mock_supervisor_comms.reset_mock() |
| mock_supervisor_comms.send.side_effect = [ |
| XComResult(key=BaseXCom.XCOM_RETURN_KEY, value="__example_xcom_value__"), |
| ] |
| assert source.xcom_pull() == "__example_xcom_value__" |
| assert mock_supervisor_comms.send.mock_calls == [ |
| mock.call( |
| GetXCom( |
| key=BaseXCom.XCOM_RETURN_KEY, |
| dag_id="d1", |
| run_id="r1", |
| task_id="t2", |
| map_index=-1, |
| ), |
| ) |
| ] |
| |
| |
| TEST_ASSET = Asset(name="test_uri", uri="test://test") |
| TEST_ASSET_ALIAS = AssetAlias(name="name") |
| TEST_ASSET_REFS = [Asset.ref(name="test_uri"), Asset.ref(uri="test://test/")] |
| TEST_INLETS = [TEST_ASSET, TEST_ASSET_ALIAS] + TEST_ASSET_REFS |
| |
| |
| class TestOutletEventAccessors: |
| @pytest.mark.parametrize( |
| ("access_key", "internal_key"), |
| ( |
| (Asset("test"), AssetUniqueKey.from_asset(Asset("test"))), |
| ( |
| Asset(name="test", uri="test://asset"), |
| AssetUniqueKey.from_asset(Asset(name="test", uri="test://asset")), |
| ), |
| (AssetAlias("test_alias"), AssetAliasUniqueKey.from_asset_alias(AssetAlias("test_alias"))), |
| ), |
| ) |
| def test__get_item__dict_key_not_exists(self, access_key, internal_key): |
| outlet_event_accessors = OutletEventAccessors() |
| assert len(outlet_event_accessors) == 0 |
| outlet_event_accessor = outlet_event_accessors[access_key] |
| assert len(outlet_event_accessors) == 1 |
| assert outlet_event_accessor.key == internal_key |
| assert outlet_event_accessor.extra == {} |
| |
| @pytest.mark.parametrize( |
| ("access_key", "asset"), |
| ( |
| (Asset.ref(name="test"), Asset(name="test")), |
| (Asset.ref(name="test1"), Asset(name="test1", uri="test://asset-uri")), |
| (Asset.ref(uri="test://asset-uri"), Asset(uri="test://asset-uri")), |
| ), |
| ) |
| def test__get_item__asset_ref(self, access_key, asset, mock_supervisor_comms): |
| """Test accessing OutletEventAccessors with AssetRef resolves to correct Asset.""" |
| internal_key = AssetUniqueKey.from_asset(asset) |
| outlet_event_accessors = OutletEventAccessors() |
| assert len(outlet_event_accessors) == 0 |
| |
| # Asset from the API Server via the supervisor |
| mock_supervisor_comms.send.return_value = AssetResult( |
| name=asset.name, |
| uri=asset.uri, |
| group=asset.group, |
| ) |
| |
| outlet_event_accessor = outlet_event_accessors[access_key] |
| assert len(outlet_event_accessors) == 1 |
| assert outlet_event_accessor.key == internal_key |
| assert outlet_event_accessor.extra == {} |
| |
| @pytest.mark.parametrize( |
| ("name", "uri", "expected_key"), |
| ( |
| ("test_uri", "test://test/", TEST_ASSET), |
| ("test_uri", None, TEST_ASSET_REFS[0]), |
| (None, "test://test/", TEST_ASSET_REFS[1]), |
| ), |
| ) |
| @mock.patch("airflow.sdk.execution_time.context.OutletEventAccessors.__getitem__") |
| def test_for_asset(self, mocked__getitem__, name, uri, expected_key): |
| outlet_event_accessors = OutletEventAccessors() |
| outlet_event_accessors.for_asset(name=name, uri=uri) |
| assert mocked__getitem__.call_args[0][0] == expected_key |
| |
| @mock.patch("airflow.sdk.execution_time.context.OutletEventAccessors.__getitem__") |
| def test_for_asset_alias(self, mocked__getitem__): |
| outlet_event_accessors = OutletEventAccessors() |
| outlet_event_accessors.for_asset_alias(name="name") |
| assert mocked__getitem__.call_args[0][0] == TEST_ASSET_ALIAS |
| |
| |
| class TestInletEventAccessor: |
| @pytest.fixture |
| def sample_inlet_evnets_accessor(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.side_effect = [ |
| AssetResult(name="test_uri", uri="test://test", group="asset"), |
| AssetResult(name="test_uri", uri="test://test", group="asset"), |
| ] |
| obj = InletEventsAccessors(inlets=TEST_INLETS) |
| mock_supervisor_comms.reset_mock() |
| return obj |
| |
| @pytest.mark.usefixtures("mock_supervisor_comms") |
| def test__iter__(self, sample_inlet_evnets_accessor): |
| for actual, expected in zip(sample_inlet_evnets_accessor, TEST_INLETS): |
| assert actual == expected |
| |
| @pytest.mark.usefixtures("mock_supervisor_comms") |
| def test__len__(self, sample_inlet_evnets_accessor): |
| assert len(sample_inlet_evnets_accessor) == 4 |
| |
| @pytest.mark.parametrize("key", TEST_INLETS + [0, 1, 2, 3]) |
| def test__get_item__(self, key, sample_inlet_evnets_accessor, mock_supervisor_comms): |
| # This test only verifies a valid key can be used to access inlet events, |
| # but not access asset events are fetched. That is verified in test_asset_events in execution_api |
| asset_event_resp = AssetEventResult( |
| id=1, |
| created_dagruns=[], |
| timestamp=timezone.utcnow(), |
| asset=AssetResponse(name="test", uri="test", group="asset"), |
| ) |
| events_result = AssetEventsResult(asset_events=[asset_event_resp]) |
| mock_supervisor_comms.send.side_effect = [events_result] * 4 |
| |
| assert list(sample_inlet_evnets_accessor[key]) == [asset_event_resp] |
| |
| @pytest.mark.usefixtures("mock_supervisor_comms") |
| def test__get_item__out_of_index(self, sample_inlet_evnets_accessor): |
| with pytest.raises(IndexError): |
| sample_inlet_evnets_accessor[5] |
| |
| def test__get_item__with_filters(self, sample_inlet_evnets_accessor, mock_supervisor_comms): |
| asset_event_resp = AssetEventResult( |
| id=1, |
| created_dagruns=[], |
| timestamp=timezone.utcnow(), |
| asset=AssetResponse(name="test_uri", uri="test_uri", group="asset"), |
| ) |
| events_result = AssetEventsResult(asset_events=[asset_event_resp]) |
| mock_supervisor_comms.send.side_effect = [events_result] * 10 |
| |
| list(sample_inlet_evnets_accessor[TEST_ASSET]) |
| list(sample_inlet_evnets_accessor[TEST_ASSET].after("2024-01-01T00:00:00Z")) |
| list(sample_inlet_evnets_accessor[TEST_ASSET].before("2024-01-01T00:00:00Z")) |
| list(sample_inlet_evnets_accessor[TEST_ASSET].limit(10)) |
| list( |
| sample_inlet_evnets_accessor[TEST_ASSET] |
| .after("2024-01-01T00:00:00Z") |
| .before("2024-01-02T00:00:00Z") |
| .limit(10) |
| ) |
| list(sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10)) |
| |
| assert mock_supervisor_comms.send.call_count == 6 |
| |
| # test accessing the accessor without list() or [] |
| sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10) |
| |
| assert mock_supervisor_comms.send.call_count == 6 |
| |
| # test accessing one of the elements |
| res = sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10)[0] |
| assert res == asset_event_resp |
| assert mock_supervisor_comms.send.call_count == 7 |
| |
| # test evaluating the accessor multiple times with the same filters |
| res = sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10) |
| assert res[0] == asset_event_resp |
| assert res[0] == asset_event_resp |
| |
| assert mock_supervisor_comms.send.call_count == 8 |
| |
| # test changing one of the filters |
| assert res.after("2024-01-01T00:00:00Z")[0] == asset_event_resp |
| |
| assert mock_supervisor_comms.send.call_count == 9 |
| |
| # test len() |
| assert len(sample_inlet_evnets_accessor[TEST_ASSET].ascending(True).limit(10)) == 1 |
| assert mock_supervisor_comms.send.call_count == 10 |
| |
| calls = mock_supervisor_comms.send.call_args_list |
| assert calls[0][0][0] == GetAssetEventByAsset( |
| name="test_uri", uri="test://test/", after=None, before=None, limit=None, ascending=True |
| ) |
| assert calls[1][0][0] == GetAssetEventByAsset( |
| name="test_uri", |
| uri="test://test/", |
| after="2024-01-01T00:00:00Z", |
| before=None, |
| limit=None, |
| ascending=True, |
| ) |
| assert calls[2][0][0] == GetAssetEventByAsset( |
| name="test_uri", |
| uri="test://test/", |
| after=None, |
| before="2024-01-01T00:00:00Z", |
| limit=None, |
| ascending=True, |
| ) |
| assert calls[3][0][0] == GetAssetEventByAsset( |
| name="test_uri", uri="test://test/", after=None, before=None, limit=10, ascending=True |
| ) |
| assert calls[4][0][0] == GetAssetEventByAsset( |
| name="test_uri", |
| uri="test://test/", |
| after="2024-01-01T00:00:00Z", |
| before="2024-01-02T00:00:00Z", |
| limit=10, |
| ascending=True, |
| ) |
| assert calls[5][0][0] == GetAssetEventByAsset( |
| name="test_uri", uri="test://test/", after=None, before=None, limit=10, ascending=False |
| ) |
| |
| @pytest.mark.parametrize( |
| ("name", "uri", "expected_key"), |
| ( |
| ("test_uri", "test://test/", TEST_ASSET), |
| ("test_uri", None, TEST_ASSET_REFS[0]), |
| (None, "test://test/", TEST_ASSET_REFS[1]), |
| ), |
| ) |
| @mock.patch("airflow.sdk.execution_time.context.InletEventsAccessors.__getitem__") |
| def test_for_asset(self, mocked__getitem__, sample_inlet_evnets_accessor, name, uri, expected_key): |
| sample_inlet_evnets_accessor.for_asset(name=name, uri=uri) |
| assert mocked__getitem__.call_args[0][0] == expected_key |
| |
| @mock.patch("airflow.sdk.execution_time.context.InletEventsAccessors.__getitem__") |
| def test_for_asset_alias(self, mocked__getitem__, sample_inlet_evnets_accessor): |
| sample_inlet_evnets_accessor.for_asset_alias(name="name") |
| assert mocked__getitem__.call_args[0][0] == TEST_ASSET_ALIAS |
| |
| def test_source_task_instance_xcom_pull(self, sample_inlet_evnets_accessor, mock_supervisor_comms): |
| mock_supervisor_comms.send.side_effect = [ |
| AssetEventsResult( |
| asset_events=[ |
| AssetEventResponse( |
| id=1, |
| timestamp=timezone.utcnow(), |
| asset=AssetResponse(name="test_uri", uri="test://test", group="asset"), |
| created_dagruns=[], |
| source_dag_id="__dag__", |
| source_run_id="__run__", |
| source_task_id="__task__", |
| source_map_index=0, |
| ), |
| AssetEventResponse( |
| id=1, |
| timestamp=timezone.utcnow(), |
| asset=AssetResponse(name="test_uri", uri="test://test", group="asset"), |
| created_dagruns=[], |
| ), |
| ], |
| ) |
| ] |
| events = list(sample_inlet_evnets_accessor[Asset.ref(name="test_uri")]) |
| assert mock_supervisor_comms.send.mock_calls == [ |
| mock.call( |
| GetAssetEventByAsset( |
| name="test_uri", |
| uri=None, |
| after=None, |
| before=None, |
| limit=None, |
| ascending=True, |
| ) |
| ) |
| ] |
| |
| assert len(events) == 2 |
| |
| dag_run_result = DagRunResult( |
| dag_id="__dag__", |
| run_id="__run__", |
| run_after=timezone.utcnow(), |
| start_date=timezone.utcnow(), |
| run_type="scheduled", |
| state="success", |
| consumed_asset_events=[], |
| ) |
| 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=[], |
| ) |
| |
| assert dag_run.start_date is None |
| |
| def test_ti_run_context_with_null_start_date(self, make_ti_context): |
| """TIRunContext should be constructable when the DagRun has start_date=None.""" |
| ti_context = make_ti_context(start_date=None) |
| |
| assert ti_context.dag_run.start_date is None |
| |
| |
| class TestAsyncGetConnection: |
| """Test async connection retrieval with secrets backends.""" |
| |
| @pytest.mark.asyncio |
| async def test_async_get_connection_from_secrets_backend(self, mock_supervisor_comms): |
| """Test that _async_get_connection successfully retrieves from secrets backend using sync_to_async.""" |
| sample_connection = Connection( |
| conn_id="test_conn", conn_type="postgres", host="localhost", port=5432, login="user" |
| ) |
| |
| class MockSecretsBackend: |
| """Simple mock secrets backend for testing.""" |
| |
| def __init__(self, connections: dict[str, Connection | None] | None = None): |
| self.connections = connections or {} |
| |
| def get_connection(self, conn_id: str) -> Connection | None: |
| return self.connections.get(conn_id) |
| |
| backend = MockSecretsBackend({"test_conn": sample_connection}) |
| |
| with patch( |
| "airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", autospec=True |
| ) as mock_load: |
| mock_load.return_value = [backend] |
| |
| result = await _async_get_connection("test_conn") |
| |
| assert result == sample_connection |
| # Should not have tried SUPERVISOR_COMMS since secrets backend had the connection |
| mock_supervisor_comms.send.assert_not_called() |
| mock_supervisor_comms.asend.assert_not_called() |
| |
| |
| class TestSecretsBackend: |
| """Test that connection resolution uses the backend chain correctly.""" |
| |
| def test_execution_api_backend_in_worker_chain(self): |
| """Test that ExecutionAPISecretsBackend is in the worker search path.""" |
| from airflow.sdk.execution_time.secrets import DEFAULT_SECRETS_SEARCH_PATH_WORKERS |
| |
| assert ( |
| "airflow.sdk.execution_time.secrets.execution_api.ExecutionAPISecretsBackend" |
| in DEFAULT_SECRETS_SEARCH_PATH_WORKERS |
| ) |
| |
| def test_metastore_backend_in_server_chain(self): |
| """Test that MetastoreBackend is in the API server search path.""" |
| from airflow.sdk.execution_time.secrets import _SERVER_DEFAULT_SECRETS_SEARCH_PATH |
| |
| assert "airflow.secrets.metastore.MetastoreBackend" in _SERVER_DEFAULT_SECRETS_SEARCH_PATH |
| assert ( |
| "airflow.sdk.execution_time.secrets.execution_api.ExecutionAPISecretsBackend" |
| not in _SERVER_DEFAULT_SECRETS_SEARCH_PATH |
| ) |
| |
| def test_get_connection_uses_backend_chain(self, mock_supervisor_comms): |
| """Test that _get_connection properly iterates through backends.""" |
| from airflow.sdk.api.datamodels._generated import ConnectionResponse |
| from airflow.sdk.execution_time.comms import ConnectionResult |
| |
| # Mock connection response |
| conn_response = ConnectionResponse( |
| conn_id="test_conn", |
| conn_type="http", |
| host="example.com", |
| port=443, |
| ) |
| 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", |
| ) |
| conn_result = ConnectionResult.from_conn_response(conn_response) |
| mock_supervisor_comms.send.return_value = conn_result |
| |
| supervisor_backend = ExecutionAPISecretsBackend() |
| |
| with patch("airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded") as mock_load: |
| mock_load.return_value = [EmptyBackend(), supervisor_backend] |
| |
| conn = _get_connection("test_conn") |
| |
| assert conn is not None |
| assert conn.conn_id == "test_conn" |
| # SupervisorComms backend was called (first backend returned None) |
| mock_supervisor_comms.send.assert_called_once() |
| |
| def test_get_connection_not_found_raises_error(self, mock_supervisor_comms): |
| """Test that _get_connection raises error when no backend finds connection.""" |
| |
| # Backend returns None (not found) |
| class EmptyBackend: |
| def get_connection(self, conn_id): |
| return None |
| |
| with patch("airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded") as mock_load: |
| mock_load.return_value = [EmptyBackend()] |
| |
| with pytest.raises(AirflowNotFoundException, match="isn't defined"): |
| _get_connection("nonexistent_conn") |
| |
| |
| class TestTaskStateStoreAccessor: |
| TI_ID = UUID("01900000-0000-0000-0000-000000000001") |
| SCOPE = TaskScope(dag_id="dag", run_id="run", task_id="task") |
| |
| def test_get_returns_value(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = TaskStateStoreResult(value="app_001") |
| |
| result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("job_id") |
| |
| assert result == "app_001" |
| mock_supervisor_comms.send.assert_called_once_with(GetTaskStateStore(ti_id=self.TI_ID, key="job_id")) |
| |
| def test_get_returns_none_on_404(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = ErrorResponse( |
| error=ErrorType.TASK_STORE_NOT_FOUND, detail={"key": "missing_key"} |
| ) |
| |
| result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("missing_key") |
| |
| assert result is None |
| |
| def test_get_returns_default_when_key_missing(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = ErrorResponse( |
| error=ErrorType.TASK_STORE_NOT_FOUND, detail={"key": "job_id"} |
| ) |
| |
| result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get( |
| "job_id", default="default-id" |
| ) |
| |
| assert result == "default-id" |
| |
| def test_get_ignores_default_when_key_exists(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = TaskStateStoreResult(value="job-001") |
| |
| result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get( |
| "job_id", default="do-not-start-here" |
| ) |
| |
| assert result == "job-001" |
| |
| def test_get_raises_on_error(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = ErrorResponse( |
| error=ErrorType.GENERIC_ERROR, detail={"message": "server error"} |
| ) |
| |
| with pytest.raises(AirflowRuntimeError): |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("some_key") |
| |
| def test_set_none_raises(self, mock_supervisor_comms): |
| with pytest.raises(ValueError, match="Cannot set value as None"): |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", None) |
| |
| def test_set_operation_with_global_retention(self, mock_supervisor_comms, time_machine): |
| """set() with no retention uses global default_retention_days config.""" |
| |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| now = datetime(2026, 5, 14, 12, 0, 0, tzinfo=dt_timezone.utc) |
| time_machine.move_to(now, tick=False) |
| |
| with conf_vars({("state_store", "default_retention_days"): "30"}): |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetTaskStateStore( |
| ti_id=self.TI_ID, |
| key="job_id", |
| value="app_001", |
| expires_at=datetime(2026, 6, 13, 12, 0, 0, tzinfo=dt_timezone.utc), |
| ) |
| ) |
| |
| def test_set_with_retention_computes_expires_at(self, mock_supervisor_comms, time_machine): |
| """set(retention=timedelta(...)) computes expires_at on the worker and sends it.""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| now = datetime(2026, 5, 14, 12, 0, 0, tzinfo=dt_timezone.utc) |
| time_machine.move_to(now, tick=False) |
| |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set( |
| "job_id", "app_001", retention=timedelta(days=7) |
| ) |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetTaskStateStore( |
| ti_id=self.TI_ID, |
| key="job_id", |
| value="app_001", |
| expires_at=datetime(2026, 5, 21, 12, 0, 0, tzinfo=dt_timezone.utc), |
| ) |
| ) |
| |
| def test_set_with_never_expire_sends_null_expires_at(self, mock_supervisor_comms): |
| """set(retention=NEVER_EXPIRE) sends expires_at=None""" |
| |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set( |
| "job_id", "app_001", retention=NEVER_EXPIRE |
| ) |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetTaskStateStore(ti_id=self.TI_ID, key="job_id", value="app_001", expires_at=None) |
| ) |
| |
| def test_set_global_default_zero_sends_null_expires_at(self, mock_supervisor_comms): |
| """When default_retention_days=0 (never expire globally), expires_at=None (stored as NULL).""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| with conf_vars({("state_store", "default_retention_days"): "0"}): |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetTaskStateStore(ti_id=self.TI_ID, key="job_id", value="app_001", expires_at=None) |
| ) |
| |
| def test_set_raises_on_negative_retention_days(self, mock_supervisor_comms): |
| """set() raises ValueError when default_retention_days is negative.""" |
| with conf_vars({("state_store", "default_retention_days"): "-1"}): |
| with pytest.raises(ValueError, match="default_retention_days must be >= 0"): |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001") |
| |
| def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms): |
| """set() logs a warning when the serialized value exceeds max_value_storage_bytes.""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| big = "x" * 110 |
| with conf_vars({("state_store", "max_value_storage_bytes"): "100"}): |
| with patch("airflow.sdk.execution_time.context.log") as mock_log: |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", big) |
| mock_log.warning.assert_called_once() |
| assert "max_value_storage_bytes" in mock_log.warning.call_args[0][0] |
| mock_supervisor_comms.send.assert_called_once() |
| |
| def test_delete_operation(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).delete("job_id") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| DeleteTaskStateStore(ti_id=self.TI_ID, key="job_id") |
| ) |
| |
| def test_clear_sends_comms_message(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).clear() |
| |
| mock_supervisor_comms.send.assert_called_once_with(ClearTaskStateStore(ti_id=self.TI_ID)) |
| |
| def test_set_datetime_raises_validation_error(self, mock_supervisor_comms): |
| """datetime is not JSON-serializable; callers must use .isoformat() first.""" |
| with pytest.raises(ValidationError): |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set( |
| "watermark", |
| datetime(2026, 5, 15, tzinfo=dt_timezone.utc), |
| ) |
| |
| mock_supervisor_comms.send.assert_not_called() |
| |
| def test_set_with_custom_backend_decorates_value_with_marker(self, mock_supervisor_comms): |
| """Custom backend ref is wrapped in external Store marker before going to DB.""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| backend = MagicMock(spec=BaseStoreBackend) |
| backend.serialize_task_state_store_to_ref.return_value = "s3://bucket/ti_123/job_id" |
| |
| with ( |
| patch("airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend), |
| conf_vars({("state_store", "default_retention_days"): "0"}), |
| ): |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "spark_001") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetTaskStateStore( |
| ti_id=self.TI_ID, |
| key="job_id", |
| value=_wrap_external_ref("s3://bucket/ti_123/job_id"), |
| expires_at=None, |
| ) |
| ) |
| |
| def test_get_with_custom_backend_removes_decoration_marker(self, mock_supervisor_comms): |
| """External Store marker is detected and the ref is passed to deserialize.""" |
| mock_supervisor_comms.send.return_value = TaskStateStoreResult( |
| value=_wrap_external_ref("s3://bucket/ti_123/job_id") |
| ) |
| |
| backend = MagicMock(spec=BaseStoreBackend) |
| backend.deserialize_task_state_store_from_ref.return_value = {"rows": 123} |
| |
| with patch( |
| "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend |
| ): |
| result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("job_id") |
| |
| assert result == {"rows": 123} |
| backend.deserialize_task_state_store_from_ref.assert_called_once_with("s3://bucket/ti_123/job_id") |
| |
| |
| class TestAssetStateStoreAccessor: |
| ASSET_NAME = "debug_watcher_asset" |
| ASSET_URI = "s3://bucket/key" |
| |
| def test_get_returns_value(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="2026-04-30T00:00:00Z") |
| |
| result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("watermark") |
| |
| assert result == "2026-04-30T00:00:00Z" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_get_returns_none_on_404(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = ErrorResponse( |
| error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key": "missing_key"} |
| ) |
| |
| result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("missing_key") |
| |
| assert result is None |
| |
| def test_get_returns_default_when_key_missing(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = ErrorResponse( |
| error=ErrorType.ASSET_STORE_NOT_FOUND, detail={"key": "watermark"} |
| ) |
| |
| result = AssetStateStoreAccessor(name=self.ASSET_NAME).get( |
| "watermark", default="2026-01-01T00:00:00+00:00" |
| ) |
| |
| assert result == "2026-01-01T00:00:00+00:00" |
| |
| def test_get_ignores_default_when_key_exists(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="2026-06-01T00:00:00+00:00") |
| |
| result = AssetStateStoreAccessor(name=self.ASSET_NAME).get( |
| "watermark", default="2026-01-01T00:00:00+00:00" |
| ) |
| |
| assert result == "2026-06-01T00:00:00+00:00" |
| |
| def test_get_raises_on_error(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = ErrorResponse( |
| error=ErrorType.GENERIC_ERROR, detail={"message": "server error"} |
| ) |
| |
| with pytest.raises(AirflowRuntimeError): |
| AssetStateStoreAccessor(name=self.ASSET_NAME).get("some_key") |
| |
| def test_set_operation(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", "2026-04-30T00:00:00Z") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="2026-04-30T00:00:00Z") |
| ) |
| |
| def test_set_none_raises(self, mock_supervisor_comms): |
| with pytest.raises(ValueError, match="Cannot set value as None"): |
| AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", None) |
| |
| def test_delete_operation(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(name=self.ASSET_NAME).delete("watermark") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_clear_operation(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(name=self.ASSET_NAME).clear() |
| |
| mock_supervisor_comms.send.assert_called_once_with(ClearAssetStateStoreByName(name=self.ASSET_NAME)) |
| |
| def test_get_by_uri(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="2026-04-30T00:00:00Z") |
| |
| result = AssetStateStoreAccessor(uri=self.ASSET_URI).get("watermark") |
| |
| assert result == "2026-04-30T00:00:00Z" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark") |
| ) |
| |
| def test_set_by_uri(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(uri=self.ASSET_URI).set("watermark", "2026-04-30T00:00:00Z") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark", value="2026-04-30T00:00:00Z") |
| ) |
| |
| def test_delete_by_uri(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(uri=self.ASSET_URI).delete("watermark") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| DeleteAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark") |
| ) |
| |
| def test_clear_by_uri(self, mock_supervisor_comms): |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(uri=self.ASSET_URI).clear() |
| |
| mock_supervisor_comms.send.assert_called_once_with(ClearAssetStateStoreByUri(uri=self.ASSET_URI)) |
| |
| def test_set_with_custom_backend_decorates_value_with_marker(self, mock_supervisor_comms): |
| """Custom backend ref is wrapped in external Store marker before going to DB.""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| backend = MagicMock(spec=BaseStoreBackend) |
| backend.serialize_asset_state_store_to_ref.return_value = "s3://bucket/assets/orders/watermark" |
| |
| with patch( |
| "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend |
| ): |
| AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", "2026-05-01") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetAssetStateStoreByName( |
| name=self.ASSET_NAME, |
| key="watermark", |
| value=_wrap_external_ref("s3://bucket/assets/orders/watermark"), |
| ) |
| ) |
| |
| def test_get_with_custom_backend_removes_decoration_marker(self, mock_supervisor_comms): |
| """External Store marker is detected and the ref is passed to deserialize.""" |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult( |
| value=_wrap_external_ref("s3://bucket/assets/orders/watermark") |
| ) |
| |
| backend = MagicMock(spec=BaseStoreBackend) |
| backend.deserialize_asset_state_store_from_ref.return_value = "2026-05-01" |
| |
| with patch( |
| "airflow.sdk.execution_time.context._get_worker_state_store_backend", return_value=backend |
| ): |
| result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("watermark") |
| |
| assert result == "2026-05-01" |
| backend.deserialize_asset_state_store_from_ref.assert_called_once_with( |
| "s3://bucket/assets/orders/watermark" |
| ) |
| |
| def test_set_warns_when_value_exceeds_limit(self, mock_supervisor_comms): |
| """set() logs a warning when the serialized value exceeds max_value_storage_bytes.""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| big = "x" * 110 |
| with conf_vars({("state_store", "max_value_storage_bytes"): "100"}): |
| with patch("airflow.sdk.execution_time.context.log") as mock_log: |
| AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", big) |
| mock_log.warning.assert_called_once() |
| assert "max_value_storage_bytes" in mock_log.warning.call_args[0][0] |
| mock_supervisor_comms.send.assert_called_once() |
| |
| |
| class TestAssetStateStoreAccessors: |
| ASSET_NAME = "my_asset" |
| ASSET_URI = "s3://bucket/key" |
| |
| def test_subscript_by_asset_routes_by_name(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v1") |
| |
| result = AssetStateStoreAccessors([asset])[asset].get("watermark") |
| |
| assert result == "v1" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_subscript_by_asset_name_ref(self, mock_supervisor_comms): |
| ref = AssetNameRef(name=self.ASSET_NAME) |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v2") |
| |
| result = AssetStateStoreAccessors([ref])[ref].get("watermark") |
| |
| assert result == "v2" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_subscript_by_uri_ref(self, mock_supervisor_comms): |
| ref = AssetUriRef(uri=self.ASSET_URI) |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v3") |
| |
| result = AssetStateStoreAccessors([ref])[ref].get("watermark") |
| |
| assert result == "v3" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark") |
| ) |
| |
| def test_get_single_inlet_simplified(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v4") |
| |
| result = AssetStateStoreAccessors([asset]).get("watermark") |
| |
| assert result == "v4" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_set_single_inlet_simplified(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessors([asset]).set("watermark", "2026-05-01") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark", value="2026-05-01") |
| ) |
| |
| def test_delete_single_inlet_simplified(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessors([asset]).delete("watermark") |
| |
| mock_supervisor_comms.send.assert_called_once_with( |
| DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_clear_single_inlet_simplified(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessors([asset]).clear() |
| |
| mock_supervisor_comms.send.assert_called_once_with(ClearAssetStateStoreByName(name=self.ASSET_NAME)) |
| |
| def test_double_reference_raises(self): |
| a1 = Asset(name="asset_one", uri="s3://one") |
| a2 = Asset(name="asset_two", uri="s3://two") |
| |
| with pytest.raises(ValueError, match="2 concrete inlets and outlets"): |
| AssetStateStoreAccessors([a1, a2]).get("watermark") |
| |
| def test_alias_inlet_resolves_to_concrete_assets(self, mock_supervisor_comms): |
| alias = AssetAlias(name="my_alias") |
| mock_supervisor_comms.send.return_value = AssetsByAliasResult( |
| assets=[AssetResult(name="resolved_asset", uri="s3://bucket/resolved", group="asset")] |
| ) |
| mock_supervisor_comms.send.return_value = AssetsByAliasResult( |
| assets=[AssetResult(name="resolved_asset", uri="s3://bucket/resolved", group="asset")] |
| ) |
| |
| accessors = AssetStateStoreAccessors([alias]) |
| |
| mock_supervisor_comms.send.assert_called_once_with(GetAssetsByAlias(alias_name="my_alias")) |
| resolved = Asset(name="resolved_asset", uri="s3://bucket/resolved") |
| assert resolved.name in accessors._by_name |
| |
| def test_alias_inlet_no_resolved_assets_contributes_nothing(self, mock_supervisor_comms): |
| alias = AssetAlias(name="empty_alias") |
| mock_supervisor_comms.send.return_value = AssetsByAliasResult(assets=[]) |
| |
| accessors = AssetStateStoreAccessors([alias]) |
| |
| assert accessors._total == 0 |
| |
| def test_outlet_only_asset_is_accessible(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v1") |
| |
| result = AssetStateStoreAccessors([], [asset])[asset].get("watermark") |
| |
| assert result == "v1" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_outlet_only_name_ref_is_accessible(self, mock_supervisor_comms): |
| ref = AssetNameRef(name=self.ASSET_NAME) |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v2") |
| |
| result = AssetStateStoreAccessors([], [ref])[ref].get("watermark") |
| |
| assert result == "v2" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_outlet_only_uri_ref_is_accessible(self, mock_supervisor_comms): |
| ref = AssetUriRef(uri=self.ASSET_URI) |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v2") |
| |
| result = AssetStateStoreAccessors([], [ref])[ref].get("watermark") |
| |
| assert result == "v2" |
| mock_supervisor_comms.send.assert_called_once_with( |
| GetAssetStateStoreByUri(uri=self.ASSET_URI, key="watermark") |
| ) |
| |
| def test_outlet_only_single_shorthand_works(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value="v3") |
| |
| result = AssetStateStoreAccessors([], [asset]).get("watermark") |
| |
| assert result == "v3" |
| |
| def test_asset_in_both_inlets_and_outlets_not_duplicated(self, mock_supervisor_comms): |
| asset = Asset(name=self.ASSET_NAME, uri=f"s3://{self.ASSET_NAME}") |
| |
| accessors = AssetStateStoreAccessors([asset], [asset]) |
| |
| assert accessors._total == 1 |
| |
| def test_outlet_alias_is_ignored(self, mock_supervisor_comms): |
| alias = AssetAlias(name="my_alias") |
| |
| accessors = AssetStateStoreAccessors([], [alias]) |
| |
| assert accessors._total == 0 |
| mock_supervisor_comms.send.assert_not_called() |
| |
| |
| class InMemoryStoreBackend(BaseStoreBackend): |
| """Simple in-memory test backend.""" |
| |
| def __init__(self): |
| self._actual_key_value_store: dict[str, str] = {} # key -> actual value |
| self.reference: dict[str, str] = {} # key -> stored ref (mem:// URI) |
| |
| def serialize_task_state_store_to_ref(self, *, value, key: str, scope) -> str: |
| ref = f"mem://{scope.dag_id}/{scope.run_id}/{scope.task_id}/{scope.map_index}/{key}" |
| self._actual_key_value_store[key] = value |
| self.reference[key] = ref |
| return ref |
| |
| def deserialize_task_state_store_from_ref(self, stored: str) -> JsonValue: |
| key = stored.rsplit("/", 1)[-1] |
| return self._actual_key_value_store.get(key, stored) |
| |
| def serialize_asset_state_store_to_ref(self, *, value, key: str, scope) -> str: |
| ref = f"mem://{scope.name or scope.uri}/{key}" |
| self._actual_key_value_store[key] = value |
| self.reference[key] = ref |
| return ref |
| |
| def deserialize_asset_state_store_from_ref(self, stored: str) -> JsonValue: |
| key = stored.rsplit("/", 1)[-1] |
| return self._actual_key_value_store.get(key, stored) |
| |
| def get(self, scope, key, *, session=None): ... |
| def set(self, scope, key, value, *, session=None): ... |
| |
| def delete(self, scope, key, *, session=None) -> None: |
| self._actual_key_value_store.pop(key, None) |
| self.reference.pop(key, None) |
| |
| def clear(self, scope, *, all_map_indices=False, session=None) -> None: |
| self._actual_key_value_store.clear() |
| self.reference.clear() |
| |
| async def aget(self, scope, key): ... |
| async def aset(self, scope, key, value): ... |
| async def adelete(self, scope, key): ... |
| async def aclear(self, scope, *, all_map_indices=False): ... |
| |
| |
| class TestTaskStateStoreAccessorWithCustomBackend: |
| TI_ID = UUID("01900000-0000-0000-0000-000000000002") |
| SCOPE = TaskScope(dag_id="dag", run_id="run", task_id="task") |
| |
| @pytest.fixture(autouse=True) |
| def backend(self): |
| b = InMemoryStoreBackend() |
| with mock.patch( |
| "airflow.sdk.execution_time.context._get_worker_state_store_backend", |
| return_value=b, |
| ): |
| yield b |
| |
| def test_set_returns_reference_to_storage(self, mock_supervisor_comms, backend, time_machine): |
| """set() stores actual value in backend and sends mem:// reference via comms.""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| expected_ref = f"mem://{self.SCOPE.dag_id}/{self.SCOPE.run_id}/{self.SCOPE.task_id}/{self.SCOPE.map_index}/job_id" |
| |
| frozen_dt = datetime(2026, 1, 1, 12, 0, 0, tzinfo=dt_timezone.utc) |
| time_machine.move_to(frozen_dt, tick=False) |
| |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).set("job_id", "app_001") |
| # comms message has the mem:// reference, not the actual value |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetTaskStateStore( |
| ti_id=self.TI_ID, |
| key="job_id", |
| value=_wrap_external_ref(expected_ref), |
| expires_at=frozen_dt + timedelta(days=30), |
| ) |
| ) |
| # actual value is stored on the backend, reference is stored for DB |
| assert backend._actual_key_value_store["job_id"] == "app_001" |
| assert backend.reference["job_id"] == expected_ref |
| |
| def test_get_resolves_reference_to_actual_value(self, mock_supervisor_comms, backend): |
| """get() fetches mem:// reference from DB, resolves it to actual value via backend.""" |
| ref = _wrap_external_ref( |
| f"mem://{self.SCOPE.dag_id}/{self.SCOPE.run_id}/{self.SCOPE.task_id}/{self.SCOPE.map_index}/job_id" |
| ) |
| backend._actual_key_value_store["job_id"] = "app_001" |
| mock_supervisor_comms.send.return_value = TaskStateStoreResult(value=ref) |
| |
| result = TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).get("job_id") |
| # actual value is resolved from mem:// reference via backend |
| assert result == "app_001" |
| |
| def test_deletes_from_backend_and_removes_db_ref(self, mock_supervisor_comms, backend): |
| """delete() purges from backend storage and removes the DB reference.""" |
| backend._actual_key_value_store["job_id"] = "app_001" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).delete("job_id") |
| |
| # backend does not have the value anymore |
| assert "job_id" not in backend._actual_key_value_store |
| # request to delete reference in DB was made |
| mock_supervisor_comms.send.assert_any_call(DeleteTaskStateStore(ti_id=self.TI_ID, key="job_id")) |
| |
| def test_clears_all_from_backend_and_clears_db(self, mock_supervisor_comms, backend): |
| """clear() purges all backend objects for the TI and removes all DB references.""" |
| backend._actual_key_value_store["job_id"] = "app_001" |
| backend._actual_key_value_store["checkpoint"] = "step_3" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| TaskStateStoreAccessor(ti_id=self.TI_ID, scope=self.SCOPE).clear() |
| |
| assert "job_id" not in backend._actual_key_value_store |
| assert "checkpoint" not in backend._actual_key_value_store |
| mock_supervisor_comms.send.assert_any_call(ClearTaskStateStore(ti_id=self.TI_ID)) |
| |
| |
| class TestAssetStateStoreAccessorWithCustomBackend: |
| ASSET_NAME = "my_asset" |
| |
| @pytest.fixture(autouse=True) |
| def backend(self): |
| b = InMemoryStoreBackend() |
| with mock.patch( |
| "airflow.sdk.execution_time.context._get_worker_state_store_backend", |
| return_value=b, |
| ): |
| yield b |
| |
| def test_set_sends_reference_not_value(self, mock_supervisor_comms, backend): |
| """set() stores actual value in backend and sends mem:// reference via comms.""" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(name=self.ASSET_NAME).set("watermark", "2026-05-01") |
| |
| expected_ref = f"mem://{self.ASSET_NAME}/watermark" |
| # comms message has the mem:// reference, not the actual value |
| mock_supervisor_comms.send.assert_called_once_with( |
| SetAssetStateStoreByName( |
| name=self.ASSET_NAME, |
| key="watermark", |
| value=_wrap_external_ref(expected_ref), |
| ) |
| ) |
| # actual value is stored on the backend, reference is stored for DB |
| assert backend._actual_key_value_store["watermark"] == "2026-05-01" |
| assert backend.reference["watermark"] == expected_ref |
| |
| def test_get_resolves_reference_to_actual_value(self, mock_supervisor_comms, backend): |
| """get() fetches mem:// reference from DB, resolves it to actual value via backend.""" |
| ref = _wrap_external_ref(f"mem://{self.ASSET_NAME}/watermark") |
| backend._actual_key_value_store["watermark"] = "2026-05-01" |
| mock_supervisor_comms.send.return_value = AssetStateStoreResult(value=ref) |
| |
| result = AssetStateStoreAccessor(name=self.ASSET_NAME).get("watermark") |
| |
| # actual value is resolved from mem:// reference via backend |
| assert result == "2026-05-01" |
| |
| def test_delete_purges_from_backend_and_removes_db_ref(self, mock_supervisor_comms, backend): |
| """delete() purges from backend storage and removes the DB reference.""" |
| backend._actual_key_value_store["watermark"] = "2026-05-01" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(name=self.ASSET_NAME).delete("watermark") |
| |
| # backend doesn't have the value anymore |
| assert "watermark" not in backend._actual_key_value_store |
| # request to delete reference in DB was made |
| mock_supervisor_comms.send.assert_any_call( |
| DeleteAssetStateStoreByName(name=self.ASSET_NAME, key="watermark") |
| ) |
| |
| def test_clear_purges_all_from_backend_and_clears_db(self, mock_supervisor_comms, backend): |
| """clear() purges all backend objects and removes all DB references.""" |
| backend._actual_key_value_store["watermark"] = "2026-05-01" |
| backend._actual_key_value_store["file_count"] = "42" |
| mock_supervisor_comms.send.return_value = OKResponse(ok=True) |
| |
| AssetStateStoreAccessor(name=self.ASSET_NAME).clear() |
| |
| assert "watermark" not in backend._actual_key_value_store |
| assert "file_count" not in backend._actual_key_value_store |
| mock_supervisor_comms.send.assert_any_call(ClearAssetStateStoreByName(name=self.ASSET_NAME)) |
