kaxil commented on code in PR #68926:
URL: https://github.com/apache/airflow/pull/68926#discussion_r3542445390
##########
providers/common/ai/tests/unit/common/ai/operators/test_agent.py:
##########
@@ -557,17 +557,44 @@ def test_durable_default_false(self):
op = AgentOperator(task_id="test", prompt="test", llm_conn_id="my_llm")
assert op.durable is False
+ @patch("airflow.providers.common.ai.operators.agent.AIRFLOW_V_3_3_PLUS",
True)
+ def test_build_durable_storage_uses_task_state_store_on_3_3(self):
+ """On Airflow >= 3.3 the cache lives in the task state store -- no
durable_cache_path needed."""
+ from airflow.providers.common.ai.durable.task_state_store import
TaskStateStoreDurableStorage
+
+ accessor = MagicMock()
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c",
durable=True)
+
+ storage = op._build_durable_storage({"task_state_store": accessor})
+
+ assert isinstance(storage, TaskStateStoreDurableStorage)
+ assert storage._store is accessor
+
+ @patch("airflow.providers.common.ai.operators.agent.AIRFLOW_V_3_3_PLUS",
False)
+ def
test_build_durable_storage_falls_back_to_object_storage_below_3_3(self):
+ """On Airflow < 3.3 the cache falls back to the ObjectStorage
backend."""
+ from airflow.providers.common.ai.durable.storage import DurableStorage
+
+ ti = MagicMock(dag_id="d", task_id="t", run_id="r", map_index=-1)
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="c",
durable=True)
+
+ storage = op._build_durable_storage({"task_instance": ti})
+
+ assert isinstance(storage, DurableStorage)
+ assert storage._cache_id == "d_t_r"
Review Comment:
Good call, switched to `@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, ...)`
and dropped the flag patch. This was load-bearing too: once the collection
error was fixed, the compat runners would reach this test and fail on its
3.3-only inline import. The `< 3.3` fallback test keeps the patch since it
deliberately exercises the ObjectStorage branch on any core.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
