extended integration tests coverage

Author: snuryyeva

tests/integration/endpoints/test_a2a_integration.py

@@ -122,8 +122,9 @@ def _jsonrpc_body(
     context_id: Optional[str] = None,
     metadata: Optional[dict[str, Any]] = None,
     parts: Optional[list[dict[str, Any]]] = None,
+    method: str = "message/send",
 ) -> dict[str, Any]:
-    """Build a ``message/send`` JSON-RPC request body.
+    """Build a JSON-RPC request body for message/send or message/stream.
 
     Parameters:
         text: Text content for the default single text part. Ignored when
@@ -132,6 +133,7 @@ def _jsonrpc_body(
         metadata: Optional message metadata (e.g. ``model``, ``vector_store_ids``).
         parts: Optional explicit parts list, overriding the default single
             text part. Pass ``[]`` to build a message with no input content.
+        method: JSON-RPC method, either "message/send" or "message/stream".
     """
     message: dict[str, Any] = {
         "messageId": str(uuid.uuid4()),
@@ -145,7 +147,7 @@ def _jsonrpc_body(
     return {
         "jsonrpc": "2.0",
         "id": "1",
-        "method": "message/send",
+        "method": method,
         "params": {"message": message},
     }
 
@@ -210,18 +212,70 @@ async def test_a2a_jsonrpc_forbidden_without_action(
 
 
 @pytest.mark.asyncio
-async def test_message_send_returns_failed_task_when_ogx_is_unreachable(
+@pytest.mark.parametrize(
+    "method",
+    ["message/send", "message/stream"],
+)
+async def test_ogx_unreachable_handles_gracefully(
     mock_ogx_client: Any,
     test_auth: AuthTuple,
+    method: str,
 ) -> None:
-    """``message/send`` returns a failed task when OGX cannot be reached."""
+    """Both message/send and message/stream handle OGX unreachable gracefully."""
+    from starlette.responses import StreamingResponse
+
     mock_ogx_client.openai.list.side_effect = ApiException(status=None)
 
-    result = _jsonrpc_result(
-        await _post_a2a(_jsonrpc_body("What is a pod?"), test_auth)
+    response = await _post_a2a(
+        _jsonrpc_body("What is a pod?", method=method), test_auth
+    )
+
+    if method == "message/send":
+        result = _jsonrpc_result(response)
+        assert _task_state(result) == "failed"
+    else:
+        assert isinstance(response, StreamingResponse)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "method",
+    ["message/send", "message/stream"],
+)
+async def test_message_without_text_short_circuits(
+    test_auth: AuthTuple,
+    mocker: MockerFixture,
+    method: str,
+) -> None:
+    """Message without text short-circuits without calling agent.
+
+    Note: the parts list itself can't be empty -- the a2a-sdk's own
+    ``new_task()`` raises ``ValueError`` for an empty ``parts`` list before
+    our handler ever runs (see ``a2a.utils.task.new_task``), which surfaces
+    as a JSON-RPC -32603 Internal Error rather than our own
+    ``input_required`` handling. To reach our ``if not user_input`` branch
+    for real, the message must carry a non-text part (e.g. ``DataPart``) so
+    ``RequestContext.get_user_input()`` -- which only extracts ``TextPart``
+    content -- returns an empty string while ``parts`` stays non-empty.
+    """
+    from starlette.responses import StreamingResponse
+
+    build_agent = mocker.patch("app.endpoints.a2a.build_agent")
+
+    response = await _post_a2a(
+        _jsonrpc_body(
+            "", parts=[{"kind": "data", "data": {"foo": "bar"}}], method=method
+        ),
+        test_auth,
     )
 
-    assert _task_state(result) == "failed"
+    if method == "message/send":
+        result = _jsonrpc_result(response)
+        assert _task_state(result) == "input-required"
+    else:
+        assert isinstance(response, StreamingResponse)
+
+    build_agent.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -239,10 +293,13 @@ async def test_follow_up_reuses_one_conversation(
     context_id = _context_id(first)
     assert context_id
 
-    await _post_a2a(
-        _jsonrpc_body("What is its population?", context_id=context_id),
-        test_auth,
+    second = _jsonrpc_result(
+        await _post_a2a(
+            _jsonrpc_body("What is its population?", context_id=context_id),
+            test_auth,
+        )
     )
+    assert _context_id(second) == context_id
 
     mock_ogx_client.conversations.create.assert_awaited_once()
     assert build_agent.call_args.args[1].conversation == "conv_" + "a" * 48
@@ -299,32 +356,90 @@ async def test_message_metadata_sets_vector_store_ids(
 
 
 @pytest.mark.asyncio
-async def test_message_send_without_text_returns_input_required(
-    mock_ogx_client: Any,  # pylint: disable=unused-argument
+async def test_a2a_health_check_returns_healthy_status() -> None:
+    """Health check endpoint returns expected structure and healthy status."""
+    result = await a2a_endpoint.a2a_health_check()
+
+    assert result["status"] == "healthy"
+    assert result["service"] == "lightspeed-a2a"
+    assert "version" in result
+    assert result["version"]
+    assert "a2a_sdk_version" in result
+    assert result["a2a_sdk_version"]
+    assert "timestamp" in result
+    assert result["timestamp"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "endpoint_path",
+    [
+        "/.well-known/agent.json",
+        "/.well-known/agent-card.json",
+    ],
+)
+async def test_agent_card_endpoints_return_same_card(
+    test_auth: AuthTuple,
+    endpoint_path: str,
+) -> None:
+    """Both agent card endpoints return identical agent card based on _AGENT_CARD_CONFIG."""
+    agent_card = await a2a_endpoint.get_agent_card(auth=test_auth)
+
+    assert agent_card.name == _AGENT_CARD_CONFIG["name"]
+    assert agent_card.description == _AGENT_CARD_CONFIG["description"]
+    assert agent_card.protocol_version == _AGENT_CARD_CONFIG["protocolVersion"]
+
+    assert agent_card.provider is not None
+    assert (
+        agent_card.provider.organization
+        == _AGENT_CARD_CONFIG["provider"]["organization"]
+    )
+    assert agent_card.provider.url == _AGENT_CARD_CONFIG["provider"]["url"]
+
+    assert len(agent_card.skills) == len(_AGENT_CARD_CONFIG["skills"])
+    assert agent_card.skills[0].id == _AGENT_CARD_CONFIG["skills"][0]["id"]
+    assert agent_card.skills[0].name == _AGENT_CARD_CONFIG["skills"][0]["name"]
+    assert (
+        agent_card.skills[0].description
+        == _AGENT_CARD_CONFIG["skills"][0]["description"]
+    )
+
+    assert agent_card.capabilities is not None
+    assert (
+        agent_card.capabilities.streaming
+        == _AGENT_CARD_CONFIG["capabilities"]["streaming"]
+    )
+    assert (
+        agent_card.capabilities.push_notifications
+        == _AGENT_CARD_CONFIG["capabilities"]["pushNotifications"]
+    )
+    assert (
+        agent_card.capabilities.state_transition_history
+        == _AGENT_CARD_CONFIG["capabilities"]["stateTransitionHistory"]
+    )
+
+
+@pytest.mark.asyncio
+async def test_message_stream_with_metadata_returns_streaming_response(
+    mock_ogx_client: Any,
     test_auth: AuthTuple,
     mocker: MockerFixture,
 ) -> None:
-    """A message with no text part short-circuits to an ``input_required`` task.
+    """message/stream with metadata returns StreamingResponse."""
+    from starlette.responses import StreamingResponse
 
-    Note: the parts list itself can't be empty -- the a2a-sdk's own
-    ``new_task()`` raises ``ValueError`` for an empty ``parts`` list before
-    our handler ever runs (see ``a2a.utils.task.new_task``), which surfaces
-    as a JSON-RPC -32603 Internal Error rather than our own
-    ``input_required`` handling. To reach our ``if not user_input`` branch
-    for real, the message must carry a non-text part (e.g. ``DataPart``) so
-    ``RequestContext.get_user_input()`` -- which only extracts ``TextPart``
-    content -- returns an empty string while ``parts`` stays non-empty.
-    """
-    build_agent = mocker.patch("app.endpoints.a2a.build_agent")
+    mock_ogx_client.openai.list.return_value = make_openai_models_list_response(
+        make_openai_model(model_id="together/llama3.1", provider_id="together"),
+    )
+    _install_agent(mocker, "Response with metadata")
 
-    result = _jsonrpc_result(
-        await _post_a2a(
-            _jsonrpc_body("", parts=[{"kind": "data", "data": {"foo": "bar"}}]),
-            test_auth,
-        )
+    response = await _post_a2a(
+        _jsonrpc_body(
+            "What is a pod?",
+            method="message/stream",
+            metadata={"model": "llama3.1", "provider": "together"},
+        ),
+        test_auth,
     )
 
-    # TaskState serializes with a hyphen on the wire ("input-required"),
-    # unlike the Python enum member name (TaskState.input_required).
-    assert _task_state(result) == "input-required"
-    build_agent.assert_not_called()
+    assert isinstance(response, StreamingResponse)