aixtools 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_agent_executor.py

@@ -0,0 +1,188 @@
+"""Tests for the Pydantic AI adapter agent executor module."""
+
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from a2a.server.agent_execution import RequestContext
+from a2a.server.events import EventQueue
+from a2a.types import Message, TaskState, TaskStatusUpdateEvent
+from pydantic_ai import Agent
+
+from aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor import (
+    AgentParameters,
+    RunOutput,
+    PydanticAgentExecutor,
+    _task_failed_event,
+)
+
+
+class TestAgentParameters(unittest.TestCase):
+    """Tests for the AgentParameters model."""
+
+    def test_agent_parameters_creation(self):
+        """Test creating AgentParameters with valid data."""
+        params = AgentParameters(
+            system_prompt="Test prompt",
+            mcp_servers=["server1", "server2"]
+        )
+        
+        self.assertEqual(params.system_prompt, "Test prompt")
+        self.assertEqual(params.mcp_servers, ["server1", "server2"])
+
+
+class TestRunOutput(unittest.TestCase):
+    """Tests for the RunOutput model."""
+
+    def test_run_output_creation(self):
+        """Test creating RunOutput with valid data."""
+        output = RunOutput(
+            is_task_failed=False,
+            is_task_in_progress=True,
+            is_input_required=False,
+            output="Test output",
+            created_artifacts_paths=["path1", "path2"]
+        )
+        
+        self.assertFalse(output.is_task_failed)
+        self.assertTrue(output.is_task_in_progress)
+        self.assertFalse(output.is_input_required)
+        self.assertEqual(output.output, "Test output")
+        self.assertEqual(output.created_artifacts_paths, ["path1", "path2"])
+
+
+class TestTaskFailedEvent(unittest.TestCase):
+    """Tests for the _task_failed_event function."""
+
+    def test_task_failed_event_creation(self):
+        """Test creating a task failed event."""
+        text = "Test error message"
+        context_id = "ctx123"
+        task_id = "task456"
+        
+        event = _task_failed_event(text, context_id, task_id)
+        
+        self.assertIsInstance(event, TaskStatusUpdateEvent)
+        self.assertEqual(event.status.state, TaskState.failed)
+        self.assertEqual(event.context_id, context_id)
+        self.assertEqual(event.task_id, task_id)
+        self.assertTrue(event.final)
+
+    def test_task_failed_event_with_none_ids(self):
+        """Test creating a task failed event with None IDs."""
+        text = "Test error message"
+        
+        event = _task_failed_event(text, "default_ctx", "default_task")
+        
+        self.assertIsInstance(event, TaskStatusUpdateEvent)
+        self.assertEqual(event.status.state, TaskState.failed)
+        self.assertEqual(event.context_id, "default_ctx")
+        self.assertEqual(event.task_id, "default_task")
+        self.assertTrue(event.final)
+
+
+class TestPydanticAgentExecutor(unittest.IsolatedAsyncioTestCase):
+    """Tests for the PydanticAgentExecutor class."""
+
+    def setUp(self):
+        self.agent_params = AgentParameters(
+            system_prompt="Test system prompt",
+            mcp_servers=["test_server"]
+        )
+        self.executor = PydanticAgentExecutor(self.agent_params)
+
+    def test_init(self):
+        """Test PydanticAgentExecutor initialization."""
+        self.assertEqual(self.executor._agent_parameters, self.agent_params)
+        self.assertIsNotNone(self.executor.history_storage)
+
+    @patch("aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor.get_message_text")
+    @patch("aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor.get_file_parts")
+    def test_convert_message_to_pydantic_parts_text_only(self, mock_get_file_parts, mock_get_message_text):
+        """Test converting message with text only."""
+        mock_get_message_text.return_value = "Test message"
+        mock_get_file_parts.return_value = []
+        
+        mock_message = MagicMock(spec=Message)
+        mock_message.parts = []
+        session_tuple = ("user1", "session1")
+        
+        result = self.executor._convert_message_to_pydantic_parts(session_tuple, mock_message)
+        
+        self.assertEqual(result, "Test message")
+        mock_get_message_text.assert_called_once_with(mock_message)
+        mock_get_file_parts.assert_called_once_with([])
+
+    @patch("aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor.build_user_input")
+    @patch("aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor.get_message_text")
+    @patch("aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor.get_file_parts")
+    def test_convert_message_to_pydantic_parts_with_files(self, mock_get_file_parts, mock_get_message_text, mock_build_user_input):
+        """Test converting message with files."""
+        mock_get_message_text.return_value = "Test message"
+        mock_file_part = MagicMock()
+        mock_file_part.uri = "/test/path.txt"
+        mock_get_file_parts.return_value = [mock_file_part]
+        mock_build_user_input.return_value = ["processed", "input"]
+        
+        mock_message = MagicMock(spec=Message)
+        mock_message.parts = [mock_file_part]
+        session_tuple = ("user1", "session1")
+        
+        result = self.executor._convert_message_to_pydantic_parts(session_tuple, mock_message)
+        
+        self.assertEqual(result, ["processed", "input"])
+        # Verify that build_user_input was called (exact args may vary due to filtering)
+        mock_build_user_input.assert_called_once()
+
+    def test_execute_no_message_validation(self):
+        """Test that execute validates message presence."""
+        # This is a simple validation test without complex integration
+        mock_context = MagicMock(spec=RequestContext)
+        mock_context.message = None
+        
+        # The actual validation happens in the execute method
+        # We test the validation logic separately
+        self.assertIsNone(mock_context.message)
+
+    async def test_cancel_not_supported(self):
+        """Test that cancel raises an exception."""
+        mock_context = MagicMock(spec=RequestContext)
+        mock_event_queue = AsyncMock(spec=EventQueue)
+        
+        with self.assertRaises(Exception) as context:
+            await self.executor.cancel(mock_context, mock_event_queue)
+        
+        self.assertIn("cancel not supported", str(context.exception))
+
+    @patch("aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor.get_agent")
+    @patch("aixtools.a2a.google_sdk.pydantic_ai_adapter.agent_executor.get_configured_mcp_servers")
+    def test_build_agent(self, mock_get_mcp_servers, mock_get_agent):
+        """Test building an agent."""
+        mock_agent = MagicMock(spec=Agent)
+        mock_get_agent.return_value = mock_agent
+        mock_get_mcp_servers.return_value = ["mcp1", "mcp2"]
+        
+        session_tuple = ("user1", "session1")
+        result = self.executor._build_agent(session_tuple)
+        
+        self.assertEqual(result, mock_agent)
+        mock_get_mcp_servers.assert_called_once_with(session_tuple, ["test_server"])
+        mock_get_agent.assert_called_once_with(
+            system_prompt="Test system prompt",
+            toolsets=["mcp1", "mcp2"],
+            output_type=RunOutput
+        )
+
+    def test_task_creation_logic(self):
+        """Test task creation logic validation."""
+        # Test the basic logic without complex integration
+        mock_context = MagicMock(spec=RequestContext)
+        mock_context.current_task = None
+        mock_context.message = MagicMock(spec=Message)
+        
+        # Verify that when current_task is None, we have a message to work with
+        self.assertIsNone(mock_context.current_task)
+        self.assertIsNotNone(mock_context.message)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/unit/a2a/google_sdk/pydantic_ai_adapter/test_storage.py

@@ -0,0 +1,156 @@
+"""Tests for the Pydantic AI adapter storage module."""
+
+import unittest
+from unittest.mock import MagicMock
+
+from aixtools.a2a.google_sdk.pydantic_ai_adapter.storage import (
+    PydanticAiAgentHistoryStorage,
+    InMemoryHistoryStorage,
+)
+
+
+class TestInMemoryHistoryStorage(unittest.TestCase):
+    """Tests for the InMemoryHistoryStorage class."""
+
+    def setUp(self):
+        self.storage = InMemoryHistoryStorage()
+
+    def test_init(self):
+        """Test InMemoryHistoryStorage initialization."""
+        self.assertEqual(self.storage.storage, {})
+
+    def test_get_nonexistent_task(self):
+        """Test getting history for a task that doesn't exist."""
+        result = self.storage.get("nonexistent_task")
+        self.assertIsNone(result)
+
+    def test_store_and_get_messages(self):
+        """Test storing and retrieving messages."""
+        task_id = "test_task_1"
+        # Use simple mock objects that can be stored
+        messages = [MagicMock(), MagicMock()]
+        
+        # Store the messages
+        self.storage.store(task_id, messages)
+        
+        # Retrieve the messages
+        result = self.storage.get(task_id)
+        
+        self.assertIsNotNone(result)
+        self.assertEqual(result, messages)
+        self.assertEqual(len(result), 2)
+
+    def test_store_overwrites_existing(self):
+        """Test that storing overwrites existing messages for the same task."""
+        task_id = "test_task_3"
+        
+        # Store initial messages
+        initial_messages = [MagicMock()]
+        self.storage.store(task_id, initial_messages)
+        
+        # Store new messages (should overwrite)
+        new_messages = [MagicMock(), MagicMock()]
+        self.storage.store(task_id, new_messages)
+        
+        # Retrieve and verify new messages are stored
+        result = self.storage.get(task_id)
+        
+        self.assertIsNotNone(result)
+        self.assertEqual(result, new_messages)
+        self.assertEqual(len(result), 2)
+        self.assertNotEqual(result, initial_messages)
+
+    def test_multiple_tasks(self):
+        """Test storing and retrieving messages for multiple tasks."""
+        task1_id = "task_1"
+        task2_id = "task_2"
+        
+        task1_messages = [MagicMock()]
+        task2_messages = [MagicMock(), MagicMock()]
+        
+        # Store messages for both tasks
+        self.storage.store(task1_id, task1_messages)
+        self.storage.store(task2_id, task2_messages)
+        
+        # Retrieve and verify both tasks' messages
+        result1 = self.storage.get(task1_id)
+        result2 = self.storage.get(task2_id)
+        
+        self.assertIsNotNone(result1)
+        self.assertIsNotNone(result2)
+        self.assertEqual(result1, task1_messages)
+        self.assertEqual(result2, task2_messages)
+        self.assertNotEqual(result1, result2)
+
+    def test_store_empty_list(self):
+        """Test storing an empty list of messages."""
+        task_id = "empty_task"
+        empty_messages = []
+        
+        self.storage.store(task_id, empty_messages)
+        result = self.storage.get(task_id)
+        
+        self.assertIsNotNone(result)
+        self.assertEqual(result, empty_messages)
+        self.assertEqual(len(result), 0)
+
+    def test_get_after_multiple_stores(self):
+        """Test that get returns the most recent store for a task."""
+        task_id = "update_task"
+        
+        # Store multiple times
+        messages1 = [MagicMock()]
+        messages2 = [MagicMock()]
+        messages3 = [MagicMock(), MagicMock()]
+        
+        self.storage.store(task_id, messages1)
+        self.storage.store(task_id, messages2)
+        self.storage.store(task_id, messages3)
+        
+        result = self.storage.get(task_id)
+        
+        self.assertIsNotNone(result)
+        self.assertEqual(result, messages3)
+
+    def test_storage_isolation(self):
+        """Test that different storage instances are isolated."""
+        storage1 = InMemoryHistoryStorage()
+        storage2 = InMemoryHistoryStorage()
+        
+        task_id = "isolation_test"
+        messages1 = [MagicMock()]
+        messages2 = [MagicMock()]
+        
+        storage1.store(task_id, messages1)
+        storage2.store(task_id, messages2)
+        
+        result1 = storage1.get(task_id)
+        result2 = storage2.get(task_id)
+        
+        self.assertIsNotNone(result1)
+        self.assertIsNotNone(result2)
+        self.assertEqual(result1, messages1)
+        self.assertEqual(result2, messages2)
+        self.assertNotEqual(result1, result2)
+
+
+class TestPydanticAiAgentHistoryStorageInterface(unittest.TestCase):
+    """Tests for the PydanticAiAgentHistoryStorage abstract interface."""
+
+    def test_cannot_instantiate_abstract_class(self):
+        """Test that the abstract base class cannot be instantiated."""
+        with self.assertRaises(TypeError):
+            PydanticAiAgentHistoryStorage()
+
+    def test_inmemory_implements_interface(self):
+        """Test that InMemoryHistoryStorage properly implements the interface."""
+        storage = InMemoryHistoryStorage()
+        
+        # Verify it's an instance of the abstract base class
+        self.assertIsInstance(storage, PydanticAiAgentHistoryStorage)
+        
+        # Verify it has the required methods
+        self.assertTrue(hasattr(storage, 'get'))
+        self.assertTrue(hasattr(storage, 'store'))
+        self.assertTrue(callable(getattr(storage, 'get')))
+        self.assertTrue(callable(getattr(storage, 'store')))

tests/unit/a2a/google_sdk/test_card.py

@@ -0,0 +1,114 @@
+"""Tests for the A2A card module."""
+
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+from a2a.client import A2ACardResolver
+from a2a.types import AgentCard
+
+from aixtools.a2a.google_sdk.card import get_agent_card
+
+
+class TestCard(unittest.IsolatedAsyncioTestCase):
+    """Tests for the A2A card module."""
+
+    def setUp(self):
+        self.test_agent_host = "http://localhost:9999"
+
+    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
+    async def test_get_agent_card_success(self, mock_resolver_class):
+        """Test successful retrieval of agent card."""
+        # Setup
+        mock_client = AsyncMock(spec=httpx.AsyncClient)
+        
+        mock_resolver = AsyncMock(spec=A2ACardResolver)
+        mock_resolver_class.return_value = mock_resolver
+        
+        mock_card = MagicMock(spec=AgentCard)
+        mock_card.model_dump_json.return_value = '{"test": "data"}'
+        mock_resolver.get_agent_card.return_value = mock_card
+        
+        # Call the function
+        result = await get_agent_card(mock_client, self.test_agent_host)
+        
+        # Verify the result
+        self.assertEqual(result, mock_card)
+        
+        # Verify the resolver was created correctly
+        mock_resolver_class.assert_called_once_with(
+            httpx_client=mock_client,
+            base_url=self.test_agent_host
+        )
+        
+        # Verify the card was fetched
+        mock_resolver.get_agent_card.assert_called_once()
+        
+        # Verify the URL was set
+        self.assertEqual(result.url, self.test_agent_host)
+
+    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
+    async def test_get_agent_card_failure(self, mock_resolver_class):
+        """Test handling of errors when retrieving agent card."""
+        # Setup
+        mock_client = AsyncMock(spec=httpx.AsyncClient)
+        
+        mock_resolver = AsyncMock(spec=A2ACardResolver)
+        mock_resolver_class.return_value = mock_resolver
+        
+        # Make the resolver raise an exception
+        mock_resolver.get_agent_card.side_effect = Exception("Failed to fetch card")
+        
+        # Call the function and expect an exception
+        with self.assertRaises(RuntimeError) as context:
+            await get_agent_card(mock_client, self.test_agent_host)
+        
+        self.assertIn("Failed to fetch the public agent card", str(context.exception))
+
+    @patch("aixtools.a2a.google_sdk.card.logger")
+    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
+    async def test_get_agent_card_logging(self, mock_resolver_class, mock_logger):
+        """Test that proper logging occurs during card retrieval."""
+        # Setup
+        mock_client = AsyncMock(spec=httpx.AsyncClient)
+        
+        mock_resolver = AsyncMock(spec=A2ACardResolver)
+        mock_resolver_class.return_value = mock_resolver
+        
+        mock_card = MagicMock(spec=AgentCard)
+        mock_card.model_dump_json.return_value = '{"test": "data"}'
+        mock_resolver.get_agent_card.return_value = mock_card
+        
+        # Call the function
+        await get_agent_card(mock_client, self.test_agent_host)
+        
+        # Verify logging calls
+        mock_logger.info.assert_called()
+        self.assertEqual(mock_logger.info.call_count, 2)  # Two info calls in the function
+
+    @patch("aixtools.a2a.google_sdk.card.logger")
+    @patch("aixtools.a2a.google_sdk.card.A2ACardResolver")
+    async def test_get_agent_card_error_logging(self, mock_resolver_class, mock_logger):
+        """Test that errors are properly logged."""
+        # Setup
+        mock_client = AsyncMock(spec=httpx.AsyncClient)
+        
+        mock_resolver = AsyncMock(spec=A2ACardResolver)
+        mock_resolver_class.return_value = mock_resolver
+        
+        test_error = Exception("Test error")
+        mock_resolver.get_agent_card.side_effect = test_error
+        
+        # Call the function and expect an exception
+        with self.assertRaises(RuntimeError):
+            await get_agent_card(mock_client, self.test_agent_host)
+        
+        # Verify error logging
+        mock_logger.error.assert_called_once()
+        args, kwargs = mock_logger.error.call_args
+        self.assertIn("Critical error fetching public agent card", args[0])
+        self.assertTrue(kwargs.get('exc_info'))
+
+
+if __name__ == '__main__':
+    unittest.main()