wehub-resource-sync
261 lines
9.0 KiB
Python
261 lines
9.0 KiB
Python
"""Unit tests for simple_after_model_modifier.
|
||
|
||
Covers the truth table of (partial, has_text, has_function_call) × role
|
||
and asserts that `end_invocation` is flipped exactly when expected.
|
||
|
||
end_invocation should be set to True iff ALL of:
|
||
- llm_response.content and parts present
|
||
- partial == False (or absent)
|
||
- role == "model"
|
||
- has_text is True
|
||
- has_function_call is False
|
||
|
||
In all other cases, end_invocation must not be flipped.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
from types import SimpleNamespace
|
||
|
||
import pytest
|
||
|
||
from agents.main import simple_after_model_modifier
|
||
|
||
|
||
class FakeInvocationContext:
|
||
"""Stub for ADK's private _invocation_context — only carries end_invocation."""
|
||
|
||
def __init__(self) -> None:
|
||
self.end_invocation = False
|
||
|
||
|
||
class FakeCallbackContext:
|
||
def __init__(self, agent_name: str = "SalesPipelineAgent") -> None:
|
||
self.agent_name = agent_name
|
||
self._invocation_context = FakeInvocationContext()
|
||
|
||
|
||
def _make_part(*, text: str | None = None, function_call: object | None = None):
|
||
# google.genai Part is a pydantic model with optional fields; SimpleNamespace
|
||
# is enough because the callback uses getattr() to read them.
|
||
part = SimpleNamespace()
|
||
if text is not None:
|
||
part.text = text
|
||
else:
|
||
part.text = None
|
||
if function_call is not None:
|
||
part.function_call = function_call
|
||
else:
|
||
part.function_call = None
|
||
return part
|
||
|
||
|
||
def _make_response(
|
||
*,
|
||
role: str = "model",
|
||
has_text: bool = False,
|
||
has_function_call: bool = False,
|
||
partial: bool = False,
|
||
with_parts: bool = True,
|
||
error_message: str | None = None,
|
||
finish_reason: object = "STOP",
|
||
):
|
||
"""Build a fake LlmResponse.
|
||
|
||
`finish_reason` defaults to "STOP" — the real terminal-response shape.
|
||
`simple_after_model_modifier` / `stop_on_terminal_text` only terminate
|
||
when finish_reason is STOP, to avoid premature termination on Gemini
|
||
thinking-mode chunks that arrive non-partial with `finish_reason=None`.
|
||
"""
|
||
parts = []
|
||
if with_parts:
|
||
parts.append(
|
||
_make_part(
|
||
text="hello" if has_text else None,
|
||
function_call=SimpleNamespace(name="get_weather")
|
||
if has_function_call
|
||
else None,
|
||
)
|
||
)
|
||
content = SimpleNamespace(role=role, parts=parts) if with_parts else None
|
||
response = SimpleNamespace(
|
||
content=content,
|
||
partial=partial,
|
||
error_message=error_message,
|
||
finish_reason=finish_reason,
|
||
turn_complete=None,
|
||
)
|
||
return response
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Terminal case — the ONLY combination that should flip end_invocation.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_flips_end_invocation_on_final_text_only_model_response():
|
||
ctx = FakeCallbackContext()
|
||
response = _make_response(
|
||
role="model", has_text=True, has_function_call=False, partial=False
|
||
)
|
||
|
||
result = simple_after_model_modifier(ctx, response)
|
||
|
||
assert result is None
|
||
assert ctx._invocation_context.end_invocation is True
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# partial × has_text × has_function_call truth table (role=model)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize(
|
||
("partial", "has_text", "has_function_call", "expected_end"),
|
||
[
|
||
# (partial=False already covered above for the True case)
|
||
(False, True, True, False), # text + function_call => must NOT terminate
|
||
(False, False, True, False), # function_call only => tool call pending
|
||
(False, False, False, False), # empty parts => nothing to terminate on
|
||
(True, True, False, False), # partial text => wait for turn_complete
|
||
(True, True, True, False), # partial text + fc => wait
|
||
(True, False, True, False), # partial fc => wait
|
||
(True, False, False, False), # partial empty => wait
|
||
],
|
||
)
|
||
def test_truth_table_model_role(partial, has_text, has_function_call, expected_end):
|
||
ctx = FakeCallbackContext()
|
||
response = _make_response(
|
||
role="model",
|
||
has_text=has_text,
|
||
has_function_call=has_function_call,
|
||
partial=partial,
|
||
)
|
||
|
||
simple_after_model_modifier(ctx, response)
|
||
|
||
assert ctx._invocation_context.end_invocation is expected_end
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# role != "model" => never terminate, even on final text-only responses.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.mark.parametrize("role", ["user", "tool", ""])
|
||
def test_non_model_role_never_terminates(role):
|
||
ctx = FakeCallbackContext()
|
||
response = _make_response(
|
||
role=role, has_text=True, has_function_call=False, partial=False
|
||
)
|
||
|
||
simple_after_model_modifier(ctx, response)
|
||
|
||
assert ctx._invocation_context.end_invocation is False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Non-SalesPipelineAgent agents should be a no-op entirely.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
# Note: the legacy `test_non_sales_pipeline_agent_is_noop` test asserted the
|
||
# SalesPipelineAgent name-gate that used to live in this callback. That gate
|
||
# was lifted out when the loop terminator became universal across every
|
||
# registered LlmAgent (see `shared_chat.stop_on_terminal_text`); the
|
||
# behavior coverage moved to `tests/python/test_stop_on_terminal_text.py`
|
||
# which exercises the truth table without any agent-name filter.
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Missing content / error_message paths — should not crash.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_no_content_no_parts_is_safe():
|
||
ctx = FakeCallbackContext()
|
||
response = SimpleNamespace(content=None, partial=False, error_message=None)
|
||
|
||
result = simple_after_model_modifier(ctx, response)
|
||
|
||
assert result is None
|
||
assert ctx._invocation_context.end_invocation is False
|
||
|
||
|
||
def test_error_message_only_is_safe():
|
||
ctx = FakeCallbackContext()
|
||
response = SimpleNamespace(
|
||
content=None, partial=False, error_message="something broke"
|
||
)
|
||
|
||
result = simple_after_model_modifier(ctx, response)
|
||
|
||
assert result is None
|
||
assert ctx._invocation_context.end_invocation is False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Defensive fallback — callback_context missing _invocation_context must not crash.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_missing_invocation_context_does_not_crash():
|
||
ctx = SimpleNamespace(agent_name="SalesPipelineAgent") # no _invocation_context
|
||
response = _make_response(
|
||
role="model", has_text=True, has_function_call=False, partial=False
|
||
)
|
||
|
||
# The callback must handle the missing attribute gracefully (see the
|
||
# getattr(..., None) guard) and not raise.
|
||
result = simple_after_model_modifier(ctx, response)
|
||
|
||
assert result is None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# error_message branch must log at WARNING (CR round 4 finding #2).
|
||
#
|
||
# Gemini surfaces quota/safety-filter/context-overflow errors via
|
||
# llm_response.error_message. The prior implementation silently returned
|
||
# None, making these failures invisible in the server log. The fix logs at
|
||
# WARNING with the agent name before returning.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_error_message_logs_warning_with_agent_name(caplog):
|
||
"""When llm_response.error_message is set (no content), the callback
|
||
must emit a WARNING that includes the agent name and the error text."""
|
||
ctx = FakeCallbackContext(agent_name="SalesPipelineAgent")
|
||
response = SimpleNamespace(
|
||
content=None,
|
||
partial=False,
|
||
error_message="RESOURCE_EXHAUSTED: quota exceeded",
|
||
)
|
||
|
||
with caplog.at_level(logging.WARNING, logger="agents.main"):
|
||
result = simple_after_model_modifier(ctx, response)
|
||
|
||
assert result is None
|
||
warnings = [
|
||
rec
|
||
for rec in caplog.records
|
||
if rec.levelno == logging.WARNING
|
||
and "error_message" in rec.getMessage()
|
||
and "SalesPipelineAgent" in rec.getMessage()
|
||
and "RESOURCE_EXHAUSTED" in rec.getMessage()
|
||
]
|
||
assert warnings, (
|
||
f"expected WARNING log with agent name and error text, got: "
|
||
f"{[r.getMessage() for r in caplog.records]}"
|
||
)
|
||
|
||
|
||
# Note: the legacy `test_error_message_on_non_sales_agent_is_noop` test
|
||
# asserted that error_message warnings only fired for SalesPipelineAgent.
|
||
# With the unified `stop_on_terminal_text`, the error_message branch
|
||
# warns for EVERY agent (still a no-op for `end_invocation`), so the
|
||
# old assertion no longer holds. Error-message logging is now covered
|
||
# generically by `test_handles_error_message_branch_without_crashing`
|
||
# in `tests/python/test_stop_on_terminal_text.py`.
|