amd-gaia 0.14.1__py3-none-any.whl → 0.14.3__py3-none-any.whl

gaia/testing/__init__.py

@@ -0,0 +1,87 @@
+# Copyright(C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+"""
+GAIA Testing Utilities
+
+Provides mocks, fixtures, and assertions for testing GAIA agents
+without requiring real LLM/VLM services or databases.
+
+Example:
+    from gaia.testing import (
+        MockLLMProvider,
+        MockVLMClient,
+        create_test_agent,
+        temp_directory,
+        assert_llm_called,
+    )
+
+    def test_my_agent():
+        agent = create_test_agent(
+            MyAgent,
+            mock_responses=["I'll search for that", "Found results"],
+        )
+
+        result = agent.process_query("Find data")
+
+        assert_llm_called(agent._mock_llm)
+        assert "results" in result["answer"].lower()
+"""
+
+# Assertions
+from gaia.testing.assertions import (
+    assert_agent_completed,
+    assert_llm_called,
+    assert_llm_prompt_contains,
+    assert_no_errors,
+    assert_result_has_keys,
+    assert_result_value,
+    assert_tool_args,
+    assert_tool_called,
+    assert_vlm_called,
+)
+
+# Fixtures
+from gaia.testing.fixtures import (
+    AgentTestContext,
+    create_test_agent,
+    temp_directory,
+    temp_file,
+)
+
+# Mocks
+from gaia.testing.mocks import (
+    MockLLMProvider,
+    MockToolExecutor,
+    MockVLMClient,
+)
+
+# Re-export temp_db from database.testing for convenience
+try:
+    from gaia.database.testing import temp_db
+except ImportError:
+    # Database module may not be available
+    temp_db = None
+
+__all__ = [
+    # Mocks
+    "MockLLMProvider",
+    "MockVLMClient",
+    "MockToolExecutor",
+    # Fixtures
+    "create_test_agent",
+    "temp_directory",
+    "temp_file",
+    "temp_db",
+    "AgentTestContext",
+    # Assertions
+    "assert_llm_called",
+    "assert_llm_prompt_contains",
+    "assert_vlm_called",
+    "assert_tool_called",
+    "assert_tool_args",
+    "assert_result_has_keys",
+    "assert_result_value",
+    "assert_agent_completed",
+    "assert_no_errors",
+]

gaia/testing/assertions.py

@@ -0,0 +1,330 @@
+# Copyright(C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+"""Assertion helpers for testing GAIA agents."""
+
+from typing import Any, Dict, List, Optional, Union
+
+from gaia.testing.mocks import MockLLMProvider, MockToolExecutor, MockVLMClient
+
+
+def assert_llm_called(
+    mock_llm: MockLLMProvider,
+    times: Optional[int] = None,
+    min_times: Optional[int] = None,
+    max_times: Optional[int] = None,
+) -> None:
+    """
+    Assert that the mock LLM was called.
+
+    Args:
+        mock_llm: MockLLMProvider instance
+        times: Exact number of expected calls (optional)
+        min_times: Minimum number of calls (optional)
+        max_times: Maximum number of calls (optional)
+
+    Raises:
+        AssertionError: If call count doesn't match expectations
+
+    Example:
+        from gaia.testing import MockLLMProvider, assert_llm_called
+
+        mock_llm = MockLLMProvider(responses=["Hello"])
+        mock_llm.generate("Test")
+
+        assert_llm_called(mock_llm)  # At least once
+        assert_llm_called(mock_llm, times=1)  # Exactly once
+        assert_llm_called(mock_llm, min_times=1, max_times=5)  # Range
+    """
+    call_count = mock_llm.call_count
+
+    if times is not None:
+        assert call_count == times, (
+            f"Expected LLM to be called {times} time(s), "
+            f"but was called {call_count} time(s)"
+        )
+    else:
+        if min_times is None and max_times is None:
+            # Just check it was called at least once
+            assert call_count > 0, "Expected LLM to be called at least once"
+
+        if min_times is not None:
+            assert call_count >= min_times, (
+                f"Expected LLM to be called at least {min_times} time(s), "
+                f"but was called {call_count} time(s)"
+            )
+
+        if max_times is not None:
+            assert call_count <= max_times, (
+                f"Expected LLM to be called at most {max_times} time(s), "
+                f"but was called {call_count} time(s)"
+            )
+
+
+def assert_llm_prompt_contains(
+    mock_llm: MockLLMProvider,
+    text: str,
+    call_index: int = -1,
+) -> None:
+    """
+    Assert that an LLM prompt contains specific text.
+
+    Args:
+        mock_llm: MockLLMProvider instance
+        text: Text that should be in the prompt
+        call_index: Which call to check (-1 = last call, 0 = first call)
+
+    Raises:
+        AssertionError: If text not found in prompt
+
+    Example:
+        assert_llm_prompt_contains(mock_llm, "customer")
+        assert_llm_prompt_contains(mock_llm, "search", call_index=0)
+    """
+    assert mock_llm.call_history, "LLM was never called"
+
+    call = mock_llm.call_history[call_index]
+    prompt = call.get("prompt", "")
+
+    assert (
+        text in prompt
+    ), f"Expected prompt to contain '{text}', but prompt was:\n{prompt[:500]}"
+
+
+def assert_vlm_called(
+    mock_vlm: MockVLMClient,
+    times: Optional[int] = None,
+) -> None:
+    """
+    Assert that the mock VLM was called.
+
+    Args:
+        mock_vlm: MockVLMClient instance
+        times: Exact number of expected calls (optional)
+
+    Raises:
+        AssertionError: If call count doesn't match
+
+    Example:
+        assert_vlm_called(mock_vlm)  # At least once
+        assert_vlm_called(mock_vlm, times=2)  # Exactly twice
+    """
+    call_count = mock_vlm.call_count
+
+    if times is not None:
+        assert call_count == times, (
+            f"Expected VLM to be called {times} time(s), "
+            f"but was called {call_count} time(s)"
+        )
+    else:
+        assert call_count > 0, "Expected VLM to be called at least once"
+
+
+def assert_tool_called(
+    executor: MockToolExecutor,
+    tool_name: str,
+    times: Optional[int] = None,
+) -> None:
+    """
+    Assert that a specific tool was called.
+
+    Args:
+        executor: MockToolExecutor instance
+        tool_name: Name of the tool
+        times: Exact number of expected calls (optional)
+
+    Raises:
+        AssertionError: If tool wasn't called or count doesn't match
+
+    Example:
+        from gaia.testing import MockToolExecutor, assert_tool_called
+
+        executor = MockToolExecutor()
+        executor.execute("search", {"query": "test"})
+
+        assert_tool_called(executor, "search")
+        assert_tool_called(executor, "search", times=1)
+    """
+    calls = executor.get_tool_calls(tool_name)
+
+    if times is not None:
+        assert len(calls) == times, (
+            f"Expected tool '{tool_name}' to be called {times} time(s), "
+            f"but was called {len(calls)} time(s)"
+        )
+    else:
+        assert len(calls) > 0, (
+            f"Expected tool '{tool_name}' to be called, but it was never called. "
+            f"Tools called: {executor.tool_names_called}"
+        )
+
+
+def assert_tool_args(
+    executor: MockToolExecutor,
+    tool_name: str,
+    expected_args: Dict[str, Any],
+    call_index: int = 0,
+) -> None:
+    """
+    Assert that a tool was called with specific arguments.
+
+    Args:
+        executor: MockToolExecutor instance
+        tool_name: Name of the tool
+        expected_args: Expected arguments (subset matching)
+        call_index: Which call to check (0 = first call)
+
+    Raises:
+        AssertionError: If arguments don't match
+
+    Example:
+        executor.execute("search", {"query": "test", "limit": 10})
+        assert_tool_args(executor, "search", {"query": "test"})
+    """
+    actual_args = executor.get_tool_args(tool_name, call_index)
+
+    assert (
+        actual_args is not None
+    ), f"Tool '{tool_name}' was not called (or call_index {call_index} out of range)"
+
+    for key, expected_value in expected_args.items():
+        assert key in actual_args, (
+            f"Expected argument '{key}' not found in tool call. "
+            f"Actual args: {actual_args}"
+        )
+        assert actual_args[key] == expected_value, (
+            f"Argument '{key}' mismatch. "
+            f"Expected: {expected_value}, Actual: {actual_args[key]}"
+        )
+
+
+def assert_result_has_keys(
+    result: Dict[str, Any],
+    keys: List[str],
+) -> None:
+    """
+    Assert that a result dictionary has specific keys.
+
+    Args:
+        result: Result dictionary to check
+        keys: List of required keys
+
+    Raises:
+        AssertionError: If any key is missing
+
+    Example:
+        result = agent.process_query("test")
+        assert_result_has_keys(result, ["answer", "steps_taken"])
+    """
+    assert isinstance(
+        result, dict
+    ), f"Expected result to be dict, got {type(result).__name__}"
+
+    missing_keys = [key for key in keys if key not in result]
+    if missing_keys:
+        raise AssertionError(
+            f"Result missing required keys: {missing_keys}. "
+            f"Available keys: {list(result.keys())}"
+        )
+
+
+def assert_result_value(
+    result: Dict[str, Any],
+    key: str,
+    expected: Any,
+) -> None:
+    """
+    Assert that a result has a specific value for a key.
+
+    Args:
+        result: Result dictionary
+        key: Key to check
+        expected: Expected value
+
+    Raises:
+        AssertionError: If value doesn't match
+
+    Example:
+        assert_result_value(result, "status", "success")
+    """
+    assert key in result, f"Key '{key}' not found in result: {list(result.keys())}"
+    actual = result[key]
+    assert (
+        actual == expected
+    ), f"Value mismatch for key '{key}'. Expected: {expected}, Actual: {actual}"
+
+
+def assert_agent_completed(
+    result: Union[Dict[str, Any], str],
+    has_answer: bool = True,
+) -> None:
+    """
+    Assert that an agent completed processing successfully.
+
+    Args:
+        result: Result from agent.process_query()
+        has_answer: Whether to check for an 'answer' key
+
+    Raises:
+        AssertionError: If agent didn't complete properly
+
+    Example:
+        result = agent.process_query("test")
+        assert_agent_completed(result)
+    """
+    # Handle string results (some agents return strings directly)
+    if isinstance(result, str):
+        assert len(result) > 0, "Agent returned empty string"
+        return
+
+    assert isinstance(
+        result, dict
+    ), f"Expected result to be dict or str, got {type(result).__name__}"
+
+    # Check for error indicators
+    if "error" in result and result["error"]:
+        raise AssertionError(f"Agent returned error: {result['error']}")
+
+    if "status" in result and result["status"] == "error":
+        error_msg = result.get("message", result.get("error", "Unknown error"))
+        raise AssertionError(f"Agent returned error status: {error_msg}")
+
+    # Check for answer if required
+    if has_answer:
+        assert "answer" in result or "response" in result or "result" in result, (
+            "Agent result missing answer/response/result key. "
+            f"Keys present: {list(result.keys())}"
+        )
+
+
+def assert_no_errors(result: Dict[str, Any]) -> None:
+    """
+    Assert that a result contains no errors.
+
+    Args:
+        result: Result dictionary
+
+    Raises:
+        AssertionError: If result contains error indicators
+
+    Example:
+        result = agent.process_query("test")
+        assert_no_errors(result)
+    """
+    if not isinstance(result, dict):
+        return  # Non-dict results don't have error keys
+
+    # Check various error patterns
+    if "error" in result and result["error"]:
+        raise AssertionError(f"Result contains error: {result['error']}")
+
+    if "errors" in result and result["errors"]:
+        raise AssertionError(f"Result contains errors: {result['errors']}")
+
+    if result.get("status") == "error":
+        msg = result.get("message", result.get("error", "Unknown"))
+        raise AssertionError(f"Result has error status: {msg}")
+
+    if result.get("success") is False:
+        msg = result.get("message", result.get("error", "Unknown"))
+        raise AssertionError(f"Result indicates failure: {msg}")