This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch commonai-durable-code-mode
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/commonai-durable-code-mode by
this push:
new 9583d595a94 Import the harness directly in the CodeMode durable tests
9583d595a94 is described below
commit 9583d595a9422d92daad0b935f1b1f64af6e8b62
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Oct 15 13:56:10 2026 +0100
Import the harness directly in the CodeMode durable tests
The dev group installs pydantic-ai-harness, so a missing or broken install
now
fails collection instead of skipping the tests. The capability-function
check
runs on every Airflow version again, with the tool-approval assertion split
into its own 3.3+ test, and the durable eager case waits 0.5 s instead of 1
s.
---
.../common/ai/tests/unit/common/ai/durable/test_code_mode.py | 7 +++----
.../common/ai/tests/unit/common/ai/operators/test_agent.py | 11 +++++++++--
2 files changed, 12 insertions(+), 6 deletions(-)
diff --git a/providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py
b/providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py
index a69e3b7b400..3a0138817df 100644
--- a/providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py
+++ b/providers/common/ai/tests/unit/common/ai/durable/test_code_mode.py
@@ -35,15 +35,13 @@ from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart,
ToolCallPart
from pydantic_ai.models.function import AgentInfo, DeltaToolCall,
DeltaToolCalls, FunctionModel
from pydantic_ai.toolsets import FunctionToolset
+from pydantic_ai_harness.code_mode import CodeMode
from airflow.providers.common.ai.durable import AirflowDurability
from airflow.providers.common.ai.durable.journal import DurableJournal,
journal_scope
from unit.common.ai.durable.test_capability import Calls, attempt,
responses_so_far, tool_returns
-code_mode = pytest.importorskip("pydantic_ai_harness.code_mode")
-CodeMode = code_mode.CodeMode
-
GATHER = "import asyncio\na, b = await asyncio.gather(fetch(name='a'),
fetch(name='b'))\n[a, b]"
@@ -205,8 +203,9 @@ class TestEarlyExecution:
return
yield {0: DeltaToolCall(name="run_code",
json_args=args[:first_statement], tool_call_id="rc")}
# Give an eager pump time to run the first statement before the
call completes.
+ # Durable runs never start it, so they wait only long enough to
show that.
with contextlib.suppress(TimeoutError):
- await asyncio.wait_for(a_started.wait(), timeout=1)
+ await asyncio.wait_for(a_started.wait(), timeout=0.5 if
durable else 5)
seen_mid_stream.append(calls["fetch:a"])
yield {0: DeltaToolCall(json_args=args[first_statement:])}
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
index fb22b27fd9b..238fa47e093 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_agent.py
@@ -1456,15 +1456,22 @@ class TestAgentOperatorCapabilities:
assert op._supports_tool_approval() is True
- @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool approval
needs Airflow >= 3.3")
def test_capability_function_is_left_for_the_run_to_resolve(self,
fake_harness):
def build(ctx):
raise AssertionError("resolved before the run started")
+ op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm",
durable=True, capabilities=[build])
+
+ assert op.capabilities == [build]
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="Per-tool approval
needs Airflow >= 3.3")
+ def
test_tool_approval_check_leaves_a_capability_function_for_the_run(self,
fake_harness):
+ def build(ctx):
+ raise AssertionError("resolved before the run started")
+
op = AgentOperator(task_id="t", prompt="p", llm_conn_id="llm",
capabilities=[build])
assert op._supports_tool_approval() is True
- assert op.capabilities == [build]
@pytest.mark.skipif(
