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

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}")

gaia/testing/fixtures.py

@@ -0,0 +1,333 @@
+# Copyright(C) 2025-2026 Advanced Micro Devices, Inc. All rights reserved.
+# SPDX-License-Identifier: MIT
+
+"""Test fixtures and factories for GAIA agent testing."""
+
+import tempfile
+from contextlib import contextmanager
+from pathlib import Path
+from typing import List, Optional, Type, TypeVar
+
+from gaia.testing.mocks import MockLLMProvider, MockVLMClient
+
+# Type variable for agent classes
+AgentT = TypeVar("AgentT")
+
+
+@contextmanager
+def temp_directory():
+    """
+    Create a temporary directory for testing with automatic cleanup.
+
+    Yields a Path object to the temporary directory. The directory
+    and all its contents are automatically deleted when the context exits.
+
+    Example:
+        from gaia.testing import temp_directory
+
+        def test_file_processing():
+            with temp_directory() as tmp_dir:
+                # Create test files
+                test_file = tmp_dir / "test.txt"
+                test_file.write_text("Hello World")
+
+                # Test your agent
+                agent = MyAgent(data_dir=str(tmp_dir))
+                result = agent.process_file(str(test_file))
+
+                assert result["status"] == "success"
+
+            # Directory automatically deleted after context exits
+
+    Yields:
+        Path: Path to the temporary directory
+    """
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        yield Path(tmp_dir)
+
+
+@contextmanager
+def temp_file(
+    content: str = "",
+    suffix: str = ".txt",
+    prefix: str = "test_",
+):
+    """
+    Create a temporary file for testing with automatic cleanup.
+
+    Args:
+        content: Initial content to write to the file
+        suffix: File extension (default: .txt)
+        prefix: Filename prefix (default: test_)
+
+    Yields:
+        Path: Path to the temporary file
+
+    Example:
+        from gaia.testing import temp_file
+
+        def test_file_reading():
+            with temp_file(content="test data", suffix=".json") as tmp_path:
+                result = my_agent.read_file(str(tmp_path))
+                assert "test data" in result
+    """
+    with tempfile.NamedTemporaryFile(
+        mode="w",
+        suffix=suffix,
+        prefix=prefix,
+        delete=False,
+    ) as f:
+        f.write(content)
+        tmp_path = Path(f.name)
+
+    try:
+        yield tmp_path
+    finally:
+        tmp_path.unlink(missing_ok=True)
+
+
+def create_test_agent(
+    agent_class: Type[AgentT],
+    mock_responses: Optional[List[str]] = None,
+    mock_vlm_text: Optional[str] = None,
+    inject_mocks: bool = True,
+    **agent_kwargs,
+) -> AgentT:
+    """
+    Create an agent instance configured for testing.
+
+    Creates an agent with silent mode enabled and optionally injects
+    mock LLM and VLM providers for testing without real API calls.
+
+    Args:
+        agent_class: The Agent class to instantiate
+        mock_responses: List of LLM responses to return in sequence
+        mock_vlm_text: Text for mock VLM to return
+        inject_mocks: Whether to inject mock providers (default: True)
+        **agent_kwargs: Additional arguments passed to agent constructor
+
+    Returns:
+        Configured agent instance
+
+    Example:
+        from gaia.testing import create_test_agent
+        from my_agent import CustomerAgent
+
+        def test_customer_search():
+            agent = create_test_agent(
+                CustomerAgent,
+                mock_responses=[
+                    '{"tool": "search", "args": {"name": "John"}}',
+                    "Found 1 customer named John."
+                ],
+                db_url="sqlite:///:memory:",
+            )
+
+            result = agent.process_query("Find customer John")
+            assert "John" in result["answer"]
+
+    Note:
+        The agent is created with `skip_lemonade=True` and `silent_mode=True`
+        to avoid requiring a running Lemonade server during tests.
+    """
+    # Ensure testing-friendly defaults
+    agent_kwargs.setdefault("skip_lemonade", True)
+    agent_kwargs.setdefault("silent_mode", True)
+
+    # Create the agent
+    agent = agent_class(**agent_kwargs)
+
+    # Inject mocks if requested
+    if inject_mocks:
+        # Create and attach mock LLM
+        mock_llm = MockLLMProvider(responses=mock_responses)
+        agent._mock_llm = mock_llm  # pylint: disable=protected-access
+
+        # Inject into chat attribute if it exists
+        if hasattr(agent, "chat"):
+            agent.chat = mock_llm
+
+        # Create and attach mock VLM if text provided
+        if mock_vlm_text is not None:
+            mock_vlm = MockVLMClient(extracted_text=mock_vlm_text)
+            agent._mock_vlm = mock_vlm  # pylint: disable=protected-access
+
+            if hasattr(agent, "vlm"):
+                agent.vlm = mock_vlm
+
+    return agent
+
+
+class AgentTestContext:
+    """
+    Context manager for comprehensive agent testing.
+
+    Provides a clean testing environment with temporary directories,
+    mock providers, and automatic cleanup.
+
+    Example:
+        from gaia.testing import AgentTestContext
+        from my_agent import IntakeAgent
+
+        def test_intake_processing():
+            with AgentTestContext(IntakeAgent) as ctx:
+                # Create test file
+                test_file = ctx.create_file("form.txt", "Patient: John Doe")
+
+                # Process with agent
+                result = ctx.agent.process_query(f"Process {test_file}")
+
+                # Verify
+                assert ctx.mock_llm.was_called
+                assert "John" in result["answer"]
+    """
+
+    def __init__(
+        self,
+        agent_class: Type[AgentT],
+        mock_responses: Optional[List[str]] = None,
+        mock_vlm_text: Optional[str] = None,
+        **agent_kwargs,
+    ):
+        """
+        Initialize test context.
+
+        Args:
+            agent_class: Agent class to instantiate
+            mock_responses: LLM responses to use
+            mock_vlm_text: VLM extraction text to use
+            **agent_kwargs: Additional agent arguments
+        """
+        self.agent_class = agent_class
+        self.mock_responses = mock_responses
+        self.mock_vlm_text = mock_vlm_text
+        self.agent_kwargs = agent_kwargs
+
+        self.agent: Optional[AgentT] = None
+        self.mock_llm: Optional[MockLLMProvider] = None
+        self.mock_vlm: Optional[MockVLMClient] = None
+        self._temp_dir: Optional[Path] = None
+        self._temp_dir_context = None
+
+    def __enter__(self) -> "AgentTestContext":
+        """Enter context and set up test environment."""
+        # Create temporary directory
+        self._temp_dir_context = tempfile.TemporaryDirectory()
+        self._temp_dir = Path(self._temp_dir_context.name)
+
+        # Create mock providers
+        self.mock_llm = MockLLMProvider(responses=self.mock_responses)
+        if self.mock_vlm_text is not None:
+            self.mock_vlm = MockVLMClient(extracted_text=self.mock_vlm_text)
+
+        # Create agent with mocks
+        self.agent_kwargs.setdefault("skip_lemonade", True)
+        self.agent_kwargs.setdefault("silent_mode", True)
+
+        self.agent = self.agent_class(**self.agent_kwargs)
+
+        # Inject mocks
+        self.agent._mock_llm = self.mock_llm  # pylint: disable=protected-access
+        if hasattr(self.agent, "chat"):
+            self.agent.chat = self.mock_llm
+
+        if self.mock_vlm is not None:
+            self.agent._mock_vlm = self.mock_vlm  # pylint: disable=protected-access
+            if hasattr(self.agent, "vlm"):
+                self.agent.vlm = self.mock_vlm
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context and clean up."""
+        # Clean up agent if it has cleanup method
+        if self.agent is not None:
+            if hasattr(self.agent, "close"):
+                try:
+                    self.agent.close()
+                except Exception:
+                    pass
+            if hasattr(self.agent, "close_db"):
+                try:
+                    self.agent.close_db()
+                except Exception:
+                    pass
+            if hasattr(self.agent, "stop_all_watchers"):
+                try:
+                    self.agent.stop_all_watchers()
+                except Exception:
+                    pass
+
+        # Clean up temp directory
+        if self._temp_dir_context is not None:
+            self._temp_dir_context.cleanup()
+
+        return False  # Don't suppress exceptions
+
+    @property
+    def temp_dir(self) -> Path:
+        """Get the temporary directory path."""
+        if self._temp_dir is None:
+            raise RuntimeError("Context not entered. Use 'with' statement.")
+        return self._temp_dir
+
+    def create_file(
+        self,
+        name: str,
+        content: str = "",
+        subdir: Optional[str] = None,
+    ) -> Path:
+        """
+        Create a file in the temporary directory.
+
+        Args:
+            name: Filename
+            content: File content
+            subdir: Optional subdirectory to create file in
+
+        Returns:
+            Path to created file
+        """
+        if subdir:
+            file_dir = self.temp_dir / subdir
+            file_dir.mkdir(parents=True, exist_ok=True)
+        else:
+            file_dir = self.temp_dir
+
+        file_path = file_dir / name
+        file_path.write_text(content)
+        return file_path
+
+    def create_directory(self, name: str) -> Path:
+        """
+        Create a subdirectory in the temporary directory.
+
+        Args:
+            name: Directory name
+
+        Returns:
+            Path to created directory
+        """
+        dir_path = self.temp_dir / name
+        dir_path.mkdir(parents=True, exist_ok=True)
+        return dir_path
+
+    def set_llm_responses(self, responses: List[str]) -> None:
+        """
+        Set new LLM responses.
+
+        Args:
+            responses: New list of responses
+        """
+        if self.mock_llm:
+            self.mock_llm.set_responses(responses)
+
+    def set_vlm_text(self, text: str) -> None:
+        """
+        Set new VLM extraction text.
+
+        Args:
+            text: New extraction text
+        """
+        if self.mock_vlm:
+            self.mock_vlm.set_extracted_text(text)