chore: import upstream snapshot with attribution

showcase/integrations/google-adk/tests/python/test_after_model_modifier.py

@@ -0,0 +1,260 @@
+"""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`.