This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 42dfb20b812 Assert on structured log events in common.ai tests (#74443)
42dfb20b812 is described below
commit 42dfb20b81288b4aac802c4040ac307bd4877575
Author: Wei Lee <[email protected]>
AuthorDate: Fri Oct 9 14:49:42 2026 +0900
Assert on structured log events in common.ai tests (#74443)
---
.../unit/common/ai/managed_agents/test_failover.py | 7 +++--
.../tests/unit/common/ai/operators/test_agent.py | 31 +++++++++++++++-------
.../ai/tests/unit/common/ai/operators/test_llm.py | 2 +-
.../ai/tests/unit/common/ai/policies/test_retry.py | 18 +++++++++----
.../ai/tests/unit/common/ai/sandbox/test_base.py | 2 +-
.../ai/tests/unit/common/ai/sandbox/test_modal.py | 12 ++++++---
.../unit/common/ai/toolsets/test_managed_agent.py | 9 ++++---
.../tests/unit/common/ai/toolsets/test_sandbox.py | 21 ++++++++++-----
.../unit/common/ai/utils/test_toolset_base.py | 4 +--
9 files changed, 71 insertions(+), 35 deletions(-)
diff --git
a/providers/common/ai/tests/unit/common/ai/managed_agents/test_failover.py
b/providers/common/ai/tests/unit/common/ai/managed_agents/test_failover.py
index 8a7ba89a650..e9a652ba39f 100644
--- a/providers/common/ai/tests/unit/common/ai/managed_agents/test_failover.py
+++ b/providers/common/ai/tests/unit/common/ai/managed_agents/test_failover.py
@@ -164,8 +164,11 @@ class TestInvoke:
group = FailoverManagedAgentClient([FakeClient("a",
raises=RuntimeError("503")), FakeClient("b")])
with caplog.at_level("INFO"):
group.invoke(REQUEST)
- assert "a on fake.cloud (member 0) failed; failing over to b on
fake.cloud (member 1)" in caplog.text
- assert "served by standby b on fake.cloud (member 1)" in caplog.text
+ assert (
+ "Managed agent a on fake.cloud (member 0) failed; failing over to
b on fake.cloud (member 1)"
+ in caplog.messages
+ )
+ assert "Managed agent request served by standby b on fake.cloud
(member 1)" in caplog.messages
class TestMetrics:
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 bbfcf8ad514..d878ed546b2 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
@@ -486,7 +486,7 @@ class TestAgentOperatorToolsetTemplating:
with caplog.at_level("INFO"):
op.render_template_fields(self.CONTEXT)
- assert "Rendered toolset sql-tenant_acme" in caplog.text
+ assert "Rendered toolset sql-tenant_acme" in caplog.messages
def test_rendering_twice_logs_once(self, caplog):
"""@task.agent renders a second time; by then the id no longer
changes."""
@@ -501,7 +501,7 @@ class TestAgentOperatorToolsetTemplating:
op.render_template_fields(self.CONTEXT)
op.render_template_fields(self.CONTEXT)
- assert caplog.text.count("Rendered toolset sql-tenant_acme") == 1
+ assert caplog.messages.count("Rendered toolset sql-tenant_acme") == 1
def test_toolset_without_template_fields_is_left_as_is(self):
toolset = FunctionToolset()
@@ -1665,13 +1665,17 @@ class TestAgentOperatorDurable:
op._log_durable_summary(counter)
assert (
- "replayed 0 cached steps (0 model, 0 tool), cached 3 new steps (2
model, 1 tool)" in caplog.text
+ "Durable: replayed 0 cached steps (0 model, 0 tool), cached 3 new
steps (2 model, 1 tool)"
+ in caplog.messages
)
assert (
- "3 tool results were not cached, and a retry runs them again:
run_query (x2), get_schema"
- in caplog.text
+ "Durable: 3 tool results were not cached, and a retry runs them
again: run_query (x2), get_schema"
+ in caplog.messages
)
- assert "1 model responses were not cached, and a retry re-runs them"
in caplog.text
+ assert (
+ "Durable: 1 model responses were not cached, and a retry re-runs
them and every step after "
+ "the first of them"
+ ) in caplog.messages
def test_durable_summary_has_no_warning_when_everything_was_cached(self,
caplog):
counter = DurableStepCounter()
@@ -1681,7 +1685,10 @@ class TestAgentOperatorDurable:
with caplog.at_level("INFO"):
op._log_durable_summary(counter)
- assert "cached 1 new steps (1 model, 0 tool)" in caplog.text
+ assert (
+ "Durable: replayed 0 cached steps (0 model, 0 tool), cached 1 new
steps (1 model, 0 tool)"
+ in caplog.messages
+ )
assert not [r for r in caplog.records if r.levelname == "WARNING"]
@patch("airflow.providers.common.ai.operators.agent.AgentOperator._build_durable_storage")
@@ -1716,8 +1723,14 @@ class TestAgentOperatorDurable:
with caplog.at_level("INFO"), pytest.raises(RuntimeError,
match="downstream failure"):
op.execute(context=_make_context())
- assert "cached 2 new steps (2 model, 0 tool)" in caplog.text
- assert "1 tool results were not cached, and a retry runs them again:
send_email" in caplog.text
+ assert (
+ "Durable: replayed 0 cached steps (0 model, 0 tool), cached 2 new
steps (2 model, 0 tool)"
+ in caplog.messages
+ )
+ assert (
+ "Durable: 1 tool results were not cached, and a retry runs them
again: send_email"
+ in caplog.messages
+ )
@patch("pydantic_ai.models.wrapper.infer_model", side_effect=lambda m: m)
@patch("pydantic_ai.models.infer_model", autospec=True)
diff --git a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
index 51d79fb8baa..123b8b4cdd3 100644
--- a/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
+++ b/providers/common/ai/tests/unit/common/ai/operators/test_llm.py
@@ -404,7 +404,7 @@ class TestLLMOperatorConfidenceGate:
output = op.execute(context)
assert Summary.model_validate(output).text == "t"
- assert "'decision' was not pushed to XCom" in caplog.text
+ assert "No task instance in the context; 'decision' was not pushed to
XCom." in caplog.messages
@patch("airflow.providers.common.ai.operators.llm.PydanticAIHook",
autospec=True)
def test_resolved_model_name_pushed_to_xcom(self, mock_hook_cls,
make_mock_run_result):
diff --git a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
index bd34472b8de..303114e2c24 100644
--- a/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
+++ b/providers/common/ai/tests/unit/common/ai/policies/test_retry.py
@@ -20,6 +20,7 @@ import copy
import json
import logging
import math
+import re
import traceback
import warnings
from datetime import timedelta
@@ -990,8 +991,8 @@ class TestFallbackBehaviour:
assert decision.action == RetryAction.DEFAULT
assert decision.reason == "classifier answer not applied
(model_error); task retry settings apply"
- assert "answered 'not_a_category', which is not a configured category"
in caplog.text
- assert "KeyError" not in caplog.text
+ assert "Classifier answered 'not_a_category', which is not a
configured category" in caplog.messages
+ assert not [r for r in caplog.records if r.exc_info and r.exc_info[0]
is KeyError]
class TestOnUncertain:
@@ -1162,7 +1163,11 @@ class TestOnUncertain:
assert decision.action == RetryAction.RETRY
assert decision.retry_delay == timedelta(seconds=5)
assert decision.reason == "classifier answer not applied
(model_error); rule"
- assert "decided nothing" in caplog.text
+ assert any(
+ r.name == "airflow.providers.common.ai.policies.retry"
+ and re.match(r"MagicMock decided nothing", r.getMessage())
+ for r in caplog.records
+ )
def test_fallback_policy_raising_falls_to_the_outer_rules(self, caplog):
"""A third-party policy that blows up must not take the classifier's
rules floor with it."""
@@ -1180,7 +1185,7 @@ class TestOnUncertain:
assert decision.action == RetryAction.RETRY
assert decision.reason == "classifier answer not applied
(model_error); rule"
- assert "fallback_policy failed" in caplog.text
+ assert "fallback_policy failed, using fallback rules" in
caplog.messages
@patch(HOOK, autospec=True)
def
test_nested_classifiers_both_down_still_reach_the_outer_fail_rule(self,
mock_hook_cls):
@@ -1343,7 +1348,10 @@ class TestLLMRetryPolicy:
decision = policy.evaluate(ValueError("x"), try_number=1,
max_tries=3)
assert decision.action == RetryAction.DEFAULT
- assert "use ClassifierRetryPolicy" in caplog.text
+ assert (
+ "This model cannot answer ErrorClassification. If it is a decision
model, "
+ "use ClassifierRetryPolicy, which asks it a typed question."
+ ) in caplog.messages
def test_no_warning_for_any_instructions(self):
with warnings.catch_warnings():
diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
b/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
index 82e43e34c2a..02bff63dd97 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_base.py
@@ -549,7 +549,7 @@ class TestAttachablePolicy:
attached = backend.attach("sb", owner="o", holder="h")
assert attached.network is None
- assert "network stamp this backend cannot read" in caplog.text
+ assert "Sandbox sb carries a network stamp this backend cannot read:
'open-ish'" in caplog.messages
def test_release_clears_only_the_callers_claim(self):
backend = TaggedBackend({"sb": {OWNER_TAG: "o", HOLDER_TAG:
"someone_else"}})
diff --git a/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
b/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
index be57f1db80b..348f3865563 100644
--- a/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
+++ b/providers/common/ai/tests/unit/common/ai/sandbox/test_modal.py
@@ -20,6 +20,7 @@ import builtins
import importlib
import json
import logging
+import re
import sys
import threading
@@ -510,8 +511,11 @@ class TestAddressAllowlist:
backend, fake, SandboxSpec(allow_egress_to=["pypi.org"],
allow_egress_to_cidrs=["1.1.1.1/32"])
)
- assert "port 443 is admitted to ANY address" in caplog.text
- assert "1.1.1.1/32" in caplog.text
+ assert any(
+ r.name == "airflow.providers.common.ai.sandbox.modal"
+ and re.match(r".*port 443 is admitted to ANY
address.*1\.1\.1\.1/32", r.getMessage())
+ for r in caplog.records
+ )
def test_an_address_list_alone_does_not_warn(self, backend, fake, caplog):
caplog.set_level(logging.WARNING,
logger="airflow.providers.common.ai.sandbox.modal")
@@ -1351,7 +1355,7 @@ class TestLivenessProbe:
assert result.timed_out is True
assert result.sandbox_terminated is False
- assert "Could not check whether Modal sandbox" in caplog.text
+ assert "Could not check whether Modal sandbox sb-0 is alive" in
caplog.messages
# The handle is still cached, so the next command needs no lookup and
no create.
fake.Sandbox.from_id_error = AssertionError("a kept handle must not be
looked up again")
sandbox.process = FakeProcess(returncode=0)
@@ -1400,7 +1404,7 @@ class TestDestroy:
backend.destroy(handle)
- assert "Could not terminate" in caplog.text
+ assert "Could not terminate Modal sandbox sb-0" in caplog.messages
def test_can_be_called_while_a_command_is_outstanding(self, backend, fake):
"""
diff --git
a/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py
index 9bb03c1c27b..ef156869859 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_managed_agent.py
@@ -233,8 +233,7 @@ class TestCallTool:
async def test_logs_the_agent_it_consulted(self, caplog):
with caplog.at_level("INFO"):
await call(FakeManagedAgentToolset())
- assert "specialist-1" in caplog.text
- assert "fake.cloud" in caplog.text
+ assert "Consulted managed agent specialist-1 on fake.cloud" in
caplog.messages
@pytest.mark.asyncio
@mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats.incr",
autospec=True)
@@ -313,8 +312,10 @@ class TestManagedAgentToolset:
client = FakeClient(ref_raises=RuntimeError("Connection 'standby' not
found"))
with caplog.at_level("WARNING"):
assert await call(ManagedAgentToolset(client, tool_name="ask"),
"q") == "the answer: q"
- assert "identity could not be resolved" in caplog.text
- assert "Connection 'standby' not found" in caplog.text
+ [record] = [
+ r for r in caplog.records if r.getMessage() == "Managed agent
identity could not be resolved"
+ ]
+ assert str(record.exc_info[1]) == "Connection 'standby' not found"
@pytest.mark.asyncio
@mock.patch("airflow.providers.common.ai.toolsets.managed_agent.Stats.incr",
autospec=True)
diff --git a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
index 48d9f7edc5c..4023628b0dd 100644
--- a/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
+++ b/providers/common/ai/tests/unit/common/ai/toolsets/test_sandbox.py
@@ -17,6 +17,7 @@
from __future__ import annotations
import asyncio
+import re
import threading
import time
from types import SimpleNamespace
@@ -689,7 +690,7 @@ class TestLifecycle:
async with ts:
await _call(ts, "run_command", {"command": "x"})
- assert "may need manual cleanup" in caplog.text
+ assert "Failed to destroy sandbox box-1 on backend rec; it may need
manual cleanup" in caplog.messages
@pytest.mark.asyncio
async def test_reentry_creates_a_fresh_sandbox(self):
@@ -959,7 +960,11 @@ class TestExports:
async with ts:
await _call(ts, "run_command", {"command": "x"})
- assert "Could not remove" in caplog.text
+ assert any(
+ r.name == "airflow.providers.common.ai.toolsets.sandbox"
+ and re.match(r"Could not remove .*out\.bin\.\w+\.partial after a
failed export", r.getMessage())
+ for r in caplog.records
+ )
@pytest.mark.asyncio
async def test_a_local_destination_gets_its_directories(self, tmp_path):
@@ -1005,7 +1010,7 @@ class TestExports:
await _call(ts, "run_command", {"command": "x"})
assert (tmp_path / "out.bin").read_bytes() == b"payload"
- assert "Failed to destroy sandbox box-1" in caplog.text
+ assert "Failed to destroy sandbox box-1 on backend rec; it may need
manual cleanup" in caplog.messages
@pytest.mark.asyncio
async def
test_a_run_that_never_used_its_sandbox_cannot_deliver_its_files(self, tmp_path):
@@ -1237,8 +1242,10 @@ class TestAttachMode:
async with ts:
backend.tags_errors = [RuntimeError("tags service down")] * 3
- assert "after 3 attempts" in caplog.text
- assert "held by me" in caplog.text
+ assert (
+ "Failed to release sandbox sb-1 on backend attachable after 3
attempts; it stays marked as "
+ "held by me, and only that holder, or the task that created the
sandbox, can use it"
+ ) in caplog.messages
assert backend.tags["sb-1"][HOLDER_TAG] == "me"
@pytest.mark.asyncio
@@ -1251,7 +1258,7 @@ class TestAttachMode:
backend.tags_errors = [RuntimeError("blip"), RuntimeError("blip")]
assert HOLDER_TAG not in backend.tags["sb-1"]
- assert "Failed to release" not in caplog.text
+ assert not any("Failed to release" in m for m in caplog.messages)
@pytest.mark.asyncio
async def test_releasing_a_sandbox_that_has_ended_is_quiet(self, caplog):
@@ -1263,7 +1270,7 @@ class TestAttachMode:
async with ts:
del backend.tags["sb-1"]
- assert "Failed to release" not in caplog.text
+ assert not any("Failed to release" in m for m in caplog.messages)
@pytest.mark.asyncio
async def
test_the_sandbox_is_released_but_not_destroyed_when_a_call_raises(self):
diff --git
a/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
b/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
index 30f085a42f9..df0f364ede8 100644
--- a/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
+++ b/providers/common/ai/tests/unit/common/ai/utils/test_toolset_base.py
@@ -159,7 +159,7 @@ class TestMasking:
_call(_ScriptedToolset(ApprovalRequired(metadata={"reason":
registered_secret})))
assert caught.value.metadata == {"reason": "***"}
- assert "failed" not in caplog.text
+ assert not any("failed" in m for m in caplog.messages)
def test_masks_the_exceptions_inside_an_exception_group(self,
registered_secret):
cause = ConnectionError(f"socket closed by {registered_secret}")
@@ -276,4 +276,4 @@ class TestWithMasking:
with pytest.raises(ValueError, match="boom"):
_call(MaskingToolset(wrapped=_ScriptedToolset(ValueError("boom"))))
- assert caplog.text.count("Tool t failed") == 1
+ assert caplog.messages.count("Tool t failed") == 1
