80 lines
2.8 KiB
Python
80 lines
2.8 KiB
Python
"""SheetExtractorAgent fallback ladder tests (bare-list wrap + compact retry)."""
|
|
|
|
from unittest.mock import patch
|
|
|
|
from backend.agents.base import AgentScope, AgentUsage
|
|
from backend.agents.extractors import SheetExtractorAgent, _wrap_bare_list
|
|
|
|
|
|
def _scope():
|
|
return AgentScope(
|
|
scope_id="sheet:4",
|
|
payload={"page": {"page_number": 4, "base64": "QUJD"}},
|
|
)
|
|
|
|
|
|
def _objects(n=2):
|
|
return [
|
|
{
|
|
"object_id": f"obj-{i}",
|
|
"object_type": "equipment",
|
|
"category": "mechanical",
|
|
"name": f"RTU-{i}",
|
|
"source_text": f"RTU-{i}",
|
|
"confidence": "high",
|
|
}
|
|
for i in range(n)
|
|
]
|
|
|
|
|
|
def test_wrap_bare_list_builds_sheet_envelope():
|
|
wrapped = _wrap_bare_list(_objects(3), page_number=4)
|
|
assert wrapped["sheet"] == {}
|
|
assert len(wrapped["objects"]) == 3
|
|
|
|
|
|
def test_wrap_bare_list_passes_dicts_and_none_through():
|
|
assert _wrap_bare_list({"sheet": {}, "objects": []}, 1) == {"sheet": {}, "objects": []}
|
|
assert _wrap_bare_list(None, 1) is None
|
|
|
|
|
|
def test_run_accepts_bare_list_response():
|
|
agent = SheetExtractorAgent(usage=AgentUsage())
|
|
with patch("backend.agents.extractors.call_json",
|
|
return_value=_objects(5)) as mock_call:
|
|
result = agent.run(_scope())
|
|
assert not result.error
|
|
assert len(result.artifacts) == 1
|
|
sheet = result.artifacts[0]
|
|
assert sheet["page_number"] == 4
|
|
assert len(sheet["assertions"]) == 5
|
|
# No compact retry needed when the first call yields data.
|
|
assert mock_call.call_count == 1
|
|
# Reasoning knobs are forwarded (None when config is blank in tests).
|
|
assert "reasoning_effort" in mock_call.call_args.kwargs
|
|
assert "reasoning_max_tokens" in mock_call.call_args.kwargs
|
|
|
|
|
|
def test_run_compact_retry_after_hard_failure():
|
|
agent = SheetExtractorAgent(usage=AgentUsage())
|
|
with patch("backend.agents.extractors.call_json",
|
|
side_effect=[None, {"sheet": {"sheet_number": "A102"},
|
|
"objects": _objects(2)}]) as mock_call:
|
|
result = agent.run(_scope())
|
|
assert not result.error
|
|
assert result.artifacts[0]["sheet_number"] == "A102"
|
|
assert mock_call.call_count == 2
|
|
# Second call carried the compact suffix.
|
|
assert "COMPACT RETRY" in mock_call.call_args_list[1].kwargs["user_text"]
|
|
|
|
|
|
def test_run_returns_empty_sheet_after_both_attempts_miss():
|
|
agent = SheetExtractorAgent(usage=AgentUsage())
|
|
with patch("backend.agents.extractors.call_json", return_value=None) as mock_call:
|
|
result = agent.run(_scope())
|
|
# Coverage ladder: no text layer to rescue the page -> empty sheet,
|
|
# but no hard failure (the ladder replaced the old raise).
|
|
assert not result.error
|
|
assert result.artifacts[0]["assertions"] == []
|
|
assert mock_call.call_count == 2
|