agent-starter-pack 0.3.3__py3-none-any.whl → 0.21.0__py3-none-any.whl

agent_starter_pack/deployment_targets/agent_engine/tests/integration/test_agent_engine_app.py

@@ -0,0 +1,484 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+{%- if cookiecutter.agent_name == "adk_live" %}
+
+import asyncio
+import json
+import logging
+import subprocess
+import sys
+import threading
+import time
+from collections.abc import Iterator
+from typing import Any
+
+import pytest
+import requests
+from websockets.asyncio.client import connect
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+WS_URL = "ws://127.0.0.1:8000/ws"
+FEEDBACK_URL = "http://127.0.0.1:8000/feedback"
+
+
+def log_output(pipe: Any, log_func: Any) -> None:
+    """Log the output from the given pipe."""
+    for line in iter(pipe.readline, ""):
+        log_func(line.strip())
+
+
+def start_server() -> subprocess.Popen[str]:
+    """Start the server using expose_app in local mode."""
+    command = [
+        sys.executable,
+        "-m",
+        "uvicorn",
+        "app.app_utils.expose_app:app",
+        "--host",
+        "0.0.0.0",
+        "--port",
+        "8000",
+    ]
+    process = subprocess.Popen(
+        command,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        bufsize=1,
+        encoding="utf-8",
+    )
+
+    # Start threads to log stdout and stderr in real-time
+    threading.Thread(
+        target=log_output, args=(process.stdout, logger.info), daemon=True
+    ).start()
+    threading.Thread(
+        target=log_output, args=(process.stderr, logger.error), daemon=True
+    ).start()
+
+    return process
+
+
+def wait_for_server(timeout: int = 60, interval: int = 1) -> bool:
+    """Wait for the server to be ready."""
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            response = requests.get("http://127.0.0.1:8000/docs", timeout=10)
+            if response.status_code == 200:
+                logger.info("Server is ready")
+                return True
+        except Exception:
+            pass
+        time.sleep(interval)
+    logger.error(f"Server did not become ready within {timeout} seconds")
+    return False
+
+
+@pytest.fixture(scope="module")
+def server_fixture(request: Any) -> Iterator[subprocess.Popen[str]]:
+    """Pytest fixture to start and stop the server for testing."""
+    logger.info("Starting server process")
+    server_process = start_server()
+    if not wait_for_server():
+        pytest.fail("Server failed to start")
+    logger.info("Server process started")
+
+    def stop_server() -> None:
+        logger.info("Stopping server process")
+        server_process.terminate()
+        try:
+            server_process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            logger.warning("Server process did not terminate, killing it")
+            server_process.kill()
+            server_process.wait()
+        logger.info("Server process stopped")
+
+    request.addfinalizer(stop_server)
+    yield server_process
+
+
+@pytest.mark.asyncio
+async def test_websocket_audio_input(server_fixture: subprocess.Popen[str]) -> None:
+    """Test websocket with audio input in local mode."""
+
+    async def send_message(websocket: Any, message: dict[str, Any]) -> None:
+        """Helper to send JSON messages."""
+        await websocket.send(json.dumps(message))
+
+    async def receive_message(websocket: Any, timeout: float = 5.0) -> dict[str, Any]:
+        """Helper to receive messages with timeout."""
+        try:
+            response = await asyncio.wait_for(websocket.recv(), timeout=timeout)
+            if isinstance(response, bytes):
+                return json.loads(response.decode())
+            if isinstance(response, str):
+                return json.loads(response)
+            return response
+        except asyncio.TimeoutError as exc:
+            raise TimeoutError(
+                f"No response received within {timeout} seconds"
+            ) from exc
+
+    try:
+        await asyncio.sleep(2)
+
+        async with connect(WS_URL, ping_timeout=10, close_timeout=10) as websocket:
+            try:
+                # Wait for setupComplete
+                setup_response = await receive_message(websocket, timeout=10.0)
+                assert "setupComplete" in setup_response
+                logger.info("Received setupComplete")
+
+                # Send dummy audio chunk with user_id
+                dummy_audio = bytes([0] * 1024)
+                audio_msg = {
+                    "user_id": "test-user",
+                    "realtimeInput": {
+                        "mediaChunks": [
+                            {
+                                "mimeType": "audio/pcm;rate=16000",
+                                "data": dummy_audio.hex(),
+                            }
+                        ]
+                    },
+                }
+                await send_message(websocket, audio_msg)
+                logger.info("Sent audio chunk")
+
+                # Send text message to complete the turn (matching frontend format)
+                text_msg = {
+                    "content": {
+                        "role": "user",
+                        "parts": [{"text": "Test audio"}],
+                    }
+                }
+                await send_message(websocket, text_msg)
+                logger.info("Sent text completion")
+
+                # Collect responses
+                responses = []
+                for _ in range(10):
+                    try:
+                        response = await receive_message(websocket, timeout=5.0)
+                        responses.append(response)
+                        logger.info(f"Received: {response}")
+
+                        if isinstance(response, dict) and response.get("turn_complete"):
+                            break
+                    except TimeoutError:
+                        break
+
+                # Verify we got responses
+                assert len(responses) > 0, "No responses received"
+
+                logger.info(f"Audio test passed. Received {len(responses)} responses")
+
+            finally:
+                await websocket.close()
+
+    except Exception as e:
+        logger.error(f"Audio test failed: {e}")
+        raise
+
+
+def test_feedback_endpoint(server_fixture: subprocess.Popen[str]) -> None:
+    """Test the feedback endpoint."""
+    feedback_data = {
+        "score": 5,
+        "text": "Great response!",
+        "user_id": "test-user-123",
+        "session_id": "test-session-123",
+        "log_type": "feedback",
+    }
+
+    response = requests.post(FEEDBACK_URL, json=feedback_data, timeout=10)
+    assert response.status_code == 200
+    assert response.json() == {"status": "success"}
+    logger.info("Feedback endpoint test passed")
+{% else %}
+
+# mypy: disable-error-code="arg-type"
+{%- if cookiecutter.is_a2a %}
+
+import os
+
+import pytest
+
+from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
+from tests.helpers import (
+    build_get_request,
+    build_post_request,
+    poll_task_completion,
+)
+{%- elif cookiecutter.is_adk %}
+
+import logging
+
+import pytest
+from google.adk.events.event import Event
+
+from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
+{%- else %}
+
+import logging
+
+import pytest
+
+from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
+{%- endif %}
+{%- if cookiecutter.is_a2a %}
+
+
+@pytest.fixture
+def agent_app() -> AgentEngineApp:
+    """Fixture to create and set up AgentEngineApp instance"""
+    from {{cookiecutter.agent_directory}}.agent_engine_app import agent_engine
+
+    agent_engine.set_up()
+    return agent_engine
+{%- else %}
+
+
+@pytest.fixture
+def agent_app() -> AgentEngineApp:
+    """Fixture to create and set up AgentEngineApp instance"""
+    from {{cookiecutter.agent_directory}}.agent_engine_app import agent_engine
+
+    agent_engine.set_up()
+    return agent_engine
+{% endif %}
+{%- if cookiecutter.is_a2a %}
+
+
+@pytest.mark.asyncio
+async def test_agent_on_message_send(agent_app: AgentEngineApp) -> None:
+    """Test complete A2A message workflow from send to task completion with artifacts."""
+    # Send message
+    message_data = {
+        "message": {
+            "messageId": f"msg-{os.urandom(8).hex()}",
+            "content": [{"text": "What is the capital of France?"}],
+            "role": "ROLE_USER",
+        },
+    }
+    response = await agent_app.on_message_send(
+        request=build_post_request(message_data),
+        context=None,
+    )
+
+    # Verify task creation
+    assert "task" in response and "id" in response["task"], (
+        "Expected task with ID in response"
+    )
+
+    # Poll for completion
+    final_response = await poll_task_completion(agent_app, response["task"]["id"])
+
+    # Verify artifacts
+    assert final_response.get("artifacts"), "Expected artifacts in completed task"
+    artifact = final_response["artifacts"][0]
+    assert artifact.get("parts") and artifact["parts"][0].get("text"), (
+        "Expected artifact with text content"
+    )
+
+
+@pytest.mark.asyncio
+async def test_agent_card(agent_app: AgentEngineApp) -> None:
+    """Test agent card retrieval and validation of required A2A fields."""
+    response = await agent_app.handle_authenticated_agent_card(
+        request=build_get_request(None),
+        context=None,
+    )
+
+    # Verify core agent card fields
+    assert response.get("name") == "root_agent", "Expected agent name 'root_agent'"
+    assert response.get("protocolVersion") == "0.3.0", "Expected protocol version 0.3.0"
+    assert response.get("preferredTransport") == "HTTP+JSON", (
+        "Expected HTTP+JSON transport"
+    )
+
+    # Verify capabilities
+    capabilities = response.get("capabilities", {})
+    assert capabilities.get("streaming") is False, "Expected streaming disabled"
+
+    # Verify skills
+    skills = response.get("skills", [])
+    assert len(skills) > 0, "Expected at least one skill"
+    for skill in skills:
+        assert all(key in skill for key in ["id", "name", "description"]), (
+            "Expected id, name, and description in each skill"
+        )
+
+    # Verify extended card support
+    assert response.get("supportsAuthenticatedExtendedCard") is True, (
+        "Expected supportsAuthenticatedExtendedCard to be True"
+    )
+{% elif cookiecutter.is_adk %}
+
+@pytest.mark.asyncio
+async def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
+    """
+    Integration test for the agent stream query functionality.
+    Tests that the agent returns valid streaming responses.
+    """
+    # Create message and events for the async_stream_query
+    message = "What's the weather in San Francisco?"
+    events = []
+    async for event in agent_app.async_stream_query(message=message, user_id="test"):
+        events.append(event)
+    assert len(events) > 0, "Expected at least one chunk in response"
+
+    # Check for valid content in the response
+    has_text_content = False
+    for event in events:
+        validated_event = Event.model_validate(event)
+        content = validated_event.content
+        if (
+            content is not None
+            and content.parts
+            and any(part.text for part in content.parts)
+        ):
+            has_text_content = True
+            break
+
+    assert has_text_content, "Expected at least one event with text content"
+
+
+def test_agent_feedback(agent_app: AgentEngineApp) -> None:
+    """
+    Integration test for the agent feedback functionality.
+    Tests that feedback can be registered successfully.
+    """
+    feedback_data = {
+        "score": 5,
+        "text": "Great response!",
+        "user_id": "test-user-456",
+        "session_id": "test-session-456",
+    }
+
+    # Should not raise any exceptions
+    agent_app.register_feedback(feedback_data)
+
+    # Test invalid feedback
+    with pytest.raises(ValueError):
+        invalid_feedback = {
+            "score": "invalid",  # Score must be numeric
+            "text": "Bad feedback",
+            "user_id": "test-user-789",
+            "session_id": "test-session-789",
+        }
+        agent_app.register_feedback(invalid_feedback)
+
+    logging.info("All assertions passed for agent feedback test")
+{% else %}
+
+def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
+    """
+    Integration test for the agent stream query functionality.
+    Tests that the agent returns valid streaming responses.
+    """
+    input_dict = {
+        "messages": [
+            {"type": "human", "content": "Test message"},
+        ],
+        "user_id": "test-user",
+        "session_id": "test-session",
+    }
+
+    events = list(agent_app.stream_query(input=input_dict))
+
+    assert len(events) > 0, "Expected at least one chunk in response"
+
+    # Verify each event is a tuple of message and metadata
+    for event in events:
+        assert isinstance(event, list), "Event should be a list"
+        assert len(event) == 2, "Event should contain message and metadata"
+        message, _ = event
+
+        # Verify message structure
+        assert isinstance(message, dict), "Message should be a dictionary"
+        assert message["type"] == "constructor"
+        assert "kwargs" in message, "Constructor message should have kwargs"
+
+    # Verify at least one message has content
+    has_content = False
+    for event in events:
+        message = event[0]
+        if message.get("type") == "constructor" and "content" in message["kwargs"]:
+            has_content = True
+            break
+    assert has_content, "At least one message should have content"
+
+
+def test_agent_query(agent_app: AgentEngineApp) -> None:
+    """
+    Integration test for the agent query functionality.
+    Tests that the agent returns valid responses.
+    """
+    input_dict = {
+        "messages": [
+            {"type": "human", "content": "Test message"},
+        ],
+        "user_id": "test-user",
+        "session_id": "test-session",
+    }
+
+    response = agent_app.query(input=input_dict)
+
+    # Basic response validation
+    assert isinstance(response, dict), "Response should be a dictionary"
+    assert "messages" in response, "Response should contain messages"
+    assert len(response["messages"]) > 0, "Response should have at least one message"
+
+    # Validate last message is AI response with content
+    message = response["messages"][-1]
+    kwargs = message["kwargs"]
+    assert kwargs["type"] == "ai", "Last message should be AI response"
+    assert len(kwargs["content"]) > 0, "AI message content should not be empty"
+
+    logging.info("All assertions passed for agent query test")
+
+
+def test_agent_feedback(agent_app: AgentEngineApp) -> None:
+    """
+    Integration test for the agent feedback functionality.
+    Tests that feedback can be registered successfully.
+    """
+    feedback_data = {
+        "score": 5,
+        "text": "Great response!",
+        "user_id": "test-user-456",
+        "session_id": "test-session-456",
+    }
+
+    # Should not raise any exceptions
+    agent_app.register_feedback(feedback_data)
+
+    # Test invalid feedback
+    with pytest.raises(ValueError):
+        invalid_feedback = {
+            "score": "invalid",  # Score must be numeric
+            "text": "Bad feedback",
+            "user_id": "test-user-789",
+            "session_id": "test-session-789",
+        }
+        agent_app.register_feedback(invalid_feedback)
+
+    logging.info("All assertions passed for agent feedback test")
+{% endif %}{% endif %}

agent_starter_pack/deployment_targets/agent_engine/tests/load_test/README.md

@@ -0,0 +1,84 @@
+{%- if cookiecutter.agent_name == "adk_live" %}
+# WebSocket Load Testing for Remote Agent Engine
+
+This directory provides a comprehensive load testing framework for your Agent Engine application using WebSocket connections, leveraging the power of [Locust](http://locust.io), a leading open-source load testing tool.
+
+The load test simulates realistic user interactions by:
+- Establishing WebSocket connections
+- Sending audio chunks in the proper `realtimeInput` format
+- Sending text messages to complete turns
+- Collecting and measuring responses until `turn_complete`
+
+## Load Testing with Remote Agent Engine
+
+**1. Start the Expose App in Remote Mode:**
+
+Launch the expose app server in a separate terminal, pointing to your deployed agent engine:
+
+```bash
+uv run python -m app.app_utils.expose_app --mode remote --remote-id <your-agent-engine-id>
+```
+
+Or if you have `deployment_metadata.json` in your project root:
+
+```bash
+uv run python -m app.app_utils.expose_app --mode remote
+```
+
+**2. Execute the Load Test:**
+
+Using another terminal tab, trigger the Locust load test:
+
+```bash
+uv run --with locust==2.31.1 --with websockets locust -f tests/load_test/load_test.py \
+-H http://127.0.0.1:8000 \
+--headless \
+-t 30s -u 1 -r 1 \
+--csv=tests/load_test/.results/results \
+--html=tests/load_test/.results/report.html
+```
+
+This command initiates a 30-second load test with 1 concurrent user.
+
+**Results:**
+
+Comprehensive CSV and HTML reports detailing the load test performance will be generated and saved in the `tests/load_test/.results` directory.
+{%- else %}
+# Robust Load Testing for Generative AI Applications
+
+This directory provides a comprehensive load testing framework for your Generative AI application, leveraging the power of [Locust](http://locust.io), a leading open-source load testing tool.
+
+##  Load Testing
+
+Before running load tests, ensure you have deployed the backend remotely.
+
+Follow these steps to execute load tests:
+
+**1. Deploy the Backend Remotely:**
+   ```bash
+   gcloud config set project <your-dev-project-id>
+   make deploy
+   ```
+
+**2. Create a Virtual Environment for Locust:**
+   It's recommended to use a separate terminal tab and create a virtual environment for Locust to avoid conflicts with your application's Python environment.
+
+   ```bash
+   python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1
+   ```
+
+**3. Execute the Load Test:**
+   Trigger the Locust load test with the following command:
+
+   ```bash
+   export _AUTH_TOKEN=$(gcloud auth print-access-token -q)
+   locust -f tests/load_test/load_test.py \
+   --headless \
+   -t 30s -u 5 -r 2 \
+   --csv=tests/load_test/.results/results \
+   --html=tests/load_test/.results/report.html
+   ```
+
+   This command initiates a 30-second load test, simulating 2 users spawning per second, reaching a maximum of 10 concurrent users.
+{%- endif %}
+