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

tests/unit/a2a/google_sdk/test_remote_agent_connection.py

@@ -0,0 +1,413 @@
+"""Tests for the remote agent connection module."""
+
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch, call
+
+from a2a.client import Client
+from a2a.types import AgentCard, Message, Task, TaskState, TaskStatus, TaskQueryParams
+
+from aixtools.a2a.google_sdk.remote_agent_connection import (
+    RemoteAgentConnection,
+    is_in_terminal_state,
+    is_in_terminal_or_interrupted_state,
+)
+
+
+class TestRemoteAgentConnection(unittest.IsolatedAsyncioTestCase):
+    """Tests for the RemoteAgentConnection class."""
+
+    def setUp(self):
+        self.mock_card = MagicMock(spec=AgentCard)
+        self.mock_client = AsyncMock(spec=Client)
+        self.connection = RemoteAgentConnection(self.mock_card, self.mock_client)
+
+    def test_get_agent_card(self):
+        """Test that get_agent_card returns the stored card."""
+        result = self.connection.get_agent_card()
+        self.assertEqual(result, self.mock_card)
+
+    async def test_send_message_returns_message(self):
+        """Test send_message when it receives a Message response."""
+        mock_message = MagicMock(spec=Message)
+        mock_task = MagicMock(spec=Task)
+        
+        # Mock the async generator to yield a message
+        async def mock_generator():
+            yield mock_message
+            yield (mock_task,)
+        
+        self.mock_client.send_message.return_value = mock_generator()
+        
+        test_message = MagicMock(spec=Message)
+        result = await self.connection.send_message(test_message)
+        
+        self.assertEqual(result, mock_message)
+        self.mock_client.send_message.assert_called_once_with(test_message)
+
+    async def test_send_message_returns_terminal_task(self):
+        """Test send_message when it receives a task in terminal state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.completed
+        
+        # Mock the async generator to yield a terminal task
+        async def mock_generator():
+            yield (mock_task,)
+        
+        self.mock_client.send_message.return_value = mock_generator()
+        
+        test_message = MagicMock(spec=Message)
+        result = await self.connection.send_message(test_message)
+        
+        self.assertEqual(result, mock_task)
+
+    async def test_send_message_returns_interrupted_task(self):
+        """Test send_message when it receives a task in interrupted state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.input_required
+        
+        # Mock the async generator to yield an interrupted task
+        async def mock_generator():
+            yield (mock_task,)
+        
+        self.mock_client.send_message.return_value = mock_generator()
+        
+        test_message = MagicMock(spec=Message)
+        result = await self.connection.send_message(test_message)
+        
+        self.assertEqual(result, mock_task)
+
+    async def test_send_message_returns_last_task(self):
+        """Test send_message returns the last task when no terminal/interrupted state is reached."""
+        mock_task1 = MagicMock(spec=Task)
+        mock_task1.status = MagicMock(spec=TaskStatus)
+        mock_task1.status.state = TaskState.working
+        
+        mock_task2 = MagicMock(spec=Task)
+        mock_task2.status = MagicMock(spec=TaskStatus)
+        mock_task2.status.state = TaskState.working
+        
+        # Mock the async generator to yield multiple running tasks
+        async def mock_generator():
+            yield (mock_task1,)
+            yield (mock_task2,)
+        
+        self.mock_client.send_message.return_value = mock_generator()
+        
+        test_message = MagicMock(spec=Message)
+        result = await self.connection.send_message(test_message)
+        
+        self.assertEqual(result, mock_task2)
+
+    async def test_send_message_handles_exception(self):
+        """Test send_message properly handles and re-raises exceptions."""
+        test_error = Exception("Test error")
+        self.mock_client.send_message.side_effect = test_error
+        
+        test_message = MagicMock(spec=Message)
+        
+        with self.assertRaises(Exception) as context:
+            await self.connection.send_message(test_message)
+        
+        self.assertEqual(context.exception, test_error)
+
+    async def test_send_message_returns_none_when_no_events(self):
+        """Test send_message returns None when no events are yielded."""
+        # Mock empty async generator
+        async def mock_generator():
+            return
+            yield  # This line will never be reached
+        
+        self.mock_client.send_message.return_value = mock_generator()
+        
+        test_message = MagicMock(spec=Message)
+        result = await self.connection.send_message(test_message)
+        
+        self.assertIsNone(result)
+
+    async def test_send_message_with_polling_returns_message(self):
+        """Test send_message_with_polling when send_message returns a Message."""
+        mock_message = MagicMock(spec=Message)
+        
+        # Mock send_message to return a Message
+        self.connection.send_message = AsyncMock(return_value=mock_message)
+        
+        test_message = MagicMock(spec=Message)
+        result = await self.connection.send_message_with_polling(test_message)
+        
+        self.assertEqual(result, mock_message)
+        self.connection.send_message.assert_called_once_with(test_message)
+
+    async def test_send_message_with_polling_returns_terminal_task(self):
+        """Test send_message_with_polling when send_message returns a task in terminal state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.completed
+        
+        # Mock send_message to return a terminal task
+        self.connection.send_message = AsyncMock(return_value=mock_task)
+        
+        test_message = MagicMock(spec=Message)
+        result = await self.connection.send_message_with_polling(test_message)
+        
+        self.assertEqual(result, mock_task)
+        self.connection.send_message.assert_called_once_with(test_message)
+
+    async def test_send_message_with_polling_polls_until_terminal(self):
+        """Test send_message_with_polling polls until task reaches terminal state."""
+        # Create initial task in working state
+        initial_task = MagicMock(spec=Task)
+        initial_task.id = "task123"
+        initial_task.status = MagicMock(spec=TaskStatus)
+        initial_task.status.state = TaskState.working
+        
+        # Create intermediate task still in working state
+        intermediate_task = MagicMock(spec=Task)
+        intermediate_task.status = MagicMock(spec=TaskStatus)
+        intermediate_task.status.state = TaskState.working
+        
+        # Create final task in completed state
+        final_task = MagicMock(spec=Task)
+        final_task.status = MagicMock(spec=TaskStatus)
+        final_task.status.state = TaskState.completed
+        
+        # Mock send_message to return initial working task
+        self.connection.send_message = AsyncMock(return_value=initial_task)
+        
+        # Mock get_task to return intermediate task first, then final task
+        self.mock_client.get_task.side_effect = [intermediate_task, final_task]
+        
+        test_message = MagicMock(spec=Message)
+        
+        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+            result = await self.connection.send_message_with_polling(
+                test_message, sleep_time=0.1, max_iter=10
+            )
+        
+        self.assertEqual(result, final_task)
+        self.connection.send_message.assert_called_once_with(test_message)
+        
+        # Verify get_task was called with correct parameters
+        expected_calls = [
+            call(TaskQueryParams(id="task123")),
+            call(TaskQueryParams(id="task123"))
+        ]
+        self.mock_client.get_task.assert_has_calls(expected_calls)
+        
+        # Verify sleep was called twice (once for each polling iteration)
+        self.assertEqual(mock_sleep.call_count, 2)
+        mock_sleep.assert_called_with(0.1)
+
+    async def test_send_message_with_polling_polls_until_interrupted(self):
+        """Test send_message_with_polling polls until task reaches interrupted state."""
+        # Create initial task in working state
+        initial_task = MagicMock(spec=Task)
+        initial_task.id = "task456"
+        initial_task.status = MagicMock(spec=TaskStatus)
+        initial_task.status.state = TaskState.working
+        
+        # Create final task in input_required state (interrupted)
+        final_task = MagicMock(spec=Task)
+        final_task.status = MagicMock(spec=TaskStatus)
+        final_task.status.state = TaskState.input_required
+        
+        # Mock send_message to return initial working task
+        self.connection.send_message = AsyncMock(return_value=initial_task)
+        
+        # Mock get_task to return final task
+        self.mock_client.get_task.return_value = final_task
+        
+        test_message = MagicMock(spec=Message)
+        
+        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+            result = await self.connection.send_message_with_polling(
+                test_message, sleep_time=0.05, max_iter=5
+            )
+        
+        self.assertEqual(result, final_task)
+        self.mock_client.get_task.assert_called_once_with(TaskQueryParams(id="task456"))
+        mock_sleep.assert_called_once_with(0.05)
+
+    async def test_send_message_with_polling_timeout_exception(self):
+        """Test send_message_with_polling raises exception when max_iter is reached."""
+        # Create initial task in working state
+        initial_task = MagicMock(spec=Task)
+        initial_task.id = "task789"
+        initial_task.status = MagicMock(spec=TaskStatus)
+        initial_task.status.state = TaskState.working
+        
+        # Create task that stays in working state
+        working_task = MagicMock(spec=Task)
+        working_task.status = MagicMock(spec=TaskStatus)
+        working_task.status.state = TaskState.working
+        
+        # Mock send_message to return initial working task
+        self.connection.send_message = AsyncMock(return_value=initial_task)
+        
+        # Mock get_task to always return working task
+        self.mock_client.get_task.return_value = working_task
+        
+        test_message = MagicMock(spec=Message)
+        
+        with patch('asyncio.sleep', new_callable=AsyncMock):
+            with self.assertRaises(Exception) as context:
+                await self.connection.send_message_with_polling(
+                    test_message, sleep_time=0.1, max_iter=3
+                )
+        
+        expected_timeout = 3 * 0.1  # max_iter * sleep_time
+        self.assertIn(f"Task did not complete in {expected_timeout} seconds", str(context.exception))
+        
+        # Verify get_task was called max_iter times
+        self.assertEqual(self.mock_client.get_task.call_count, 3)
+
+    async def test_send_message_with_polling_no_task_returned(self):
+        """Test send_message_with_polling raises ValueError when send_message returns None."""
+        # Mock send_message to return None
+        self.connection.send_message = AsyncMock(return_value=None)
+        
+        test_message = MagicMock(spec=Message)
+        
+        with self.assertRaises(ValueError) as context:
+            await self.connection.send_message_with_polling(test_message)
+        
+        self.assertEqual(str(context.exception), "No task or message returned from send_message")
+        self.connection.send_message.assert_called_once_with(test_message)
+
+    async def test_send_message_with_polling_custom_parameters(self):
+        """Test send_message_with_polling with custom sleep_time and max_iter."""
+        # Create initial task in working state
+        initial_task = MagicMock(spec=Task)
+        initial_task.id = "custom_task"
+        initial_task.status = MagicMock(spec=TaskStatus)
+        initial_task.status.state = TaskState.working
+        
+        # Create final task in completed state
+        final_task = MagicMock(spec=Task)
+        final_task.status = MagicMock(spec=TaskStatus)
+        final_task.status.state = TaskState.completed
+        
+        # Mock send_message to return initial working task
+        self.connection.send_message = AsyncMock(return_value=initial_task)
+        
+        # Mock get_task to return final task
+        self.mock_client.get_task.return_value = final_task
+        
+        test_message = MagicMock(spec=Message)
+        
+        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+            result = await self.connection.send_message_with_polling(
+                test_message, sleep_time=0.5, max_iter=100
+            )
+        
+        self.assertEqual(result, final_task)
+        mock_sleep.assert_called_once_with(0.5)
+
+    async def test_send_message_with_polling_default_parameters(self):
+        """Test send_message_with_polling uses default parameters correctly."""
+        # Create initial task in working state
+        initial_task = MagicMock(spec=Task)
+        initial_task.id = "default_task"
+        initial_task.status = MagicMock(spec=TaskStatus)
+        initial_task.status.state = TaskState.working
+        
+        # Create final task in failed state (terminal)
+        final_task = MagicMock(spec=Task)
+        final_task.status = MagicMock(spec=TaskStatus)
+        final_task.status.state = TaskState.failed
+        
+        # Mock send_message to return initial working task
+        self.connection.send_message = AsyncMock(return_value=initial_task)
+        
+        # Mock get_task to return final task
+        self.mock_client.get_task.return_value = final_task
+        
+        test_message = MagicMock(spec=Message)
+        
+        with patch('asyncio.sleep', new_callable=AsyncMock) as mock_sleep:
+            result = await self.connection.send_message_with_polling(test_message)
+        
+        self.assertEqual(result, final_task)
+        # Default sleep_time is 0.2
+        mock_sleep.assert_called_once_with(0.2)
+
+
+class TestHelperFunctions(unittest.TestCase):
+    """Tests for helper functions in the remote agent connection module."""
+
+    def test_is_in_terminal_state_completed(self):
+        """Test is_in_terminal_state returns True for completed state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.completed
+        
+        result = is_in_terminal_state(mock_task)
+        self.assertTrue(result)
+
+    def test_is_in_terminal_state_canceled(self):
+        """Test is_in_terminal_state returns True for canceled state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.canceled
+        
+        result = is_in_terminal_state(mock_task)
+        self.assertTrue(result)
+
+    def test_is_in_terminal_state_failed(self):
+        """Test is_in_terminal_state returns True for failed state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.failed
+        
+        result = is_in_terminal_state(mock_task)
+        self.assertTrue(result)
+
+    def test_is_in_terminal_state_running(self):
+        """Test is_in_terminal_state returns False for running state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.working
+        
+        result = is_in_terminal_state(mock_task)
+        self.assertFalse(result)
+
+    def test_is_in_terminal_or_interrupted_state_input_required(self):
+        """Test is_in_terminal_or_interrupted_state returns True for input_required state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.input_required
+        
+        result = is_in_terminal_or_interrupted_state(mock_task)
+        self.assertTrue(result)
+
+    def test_is_in_terminal_or_interrupted_state_unknown(self):
+        """Test is_in_terminal_or_interrupted_state returns True for unknown state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.unknown
+        
+        result = is_in_terminal_or_interrupted_state(mock_task)
+        self.assertTrue(result)
+
+    def test_is_in_terminal_or_interrupted_state_completed(self):
+        """Test is_in_terminal_or_interrupted_state returns True for terminal states."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.completed
+        
+        result = is_in_terminal_or_interrupted_state(mock_task)
+        self.assertTrue(result)
+
+    def test_is_in_terminal_or_interrupted_state_running(self):
+        """Test is_in_terminal_or_interrupted_state returns False for running state."""
+        mock_task = MagicMock(spec=Task)
+        mock_task.status = MagicMock(spec=TaskStatus)
+        mock_task.status.state = TaskState.working
+        
+        result = is_in_terminal_or_interrupted_state(mock_task)
+        self.assertFalse(result)
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/unit/a2a/google_sdk/test_utils.py

@@ -0,0 +1,208 @@
+"""Tests for the A2A utils module."""
+
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import httpx
+from a2a.client import ClientConfig, ClientFactory, A2ACardResolver
+from a2a.server.agent_execution import RequestContext
+from a2a.types import AgentCard
+
+from aixtools.a2a.google_sdk.utils import (
+    _AgentCardResolver,
+    get_a2a_clients,
+    get_session_id_tuple,
+)
+from aixtools.a2a.google_sdk.remote_agent_connection import RemoteAgentConnection
+from aixtools.context import DEFAULT_USER_ID, DEFAULT_SESSION_ID
+
+
+class TestAgentCardResolver(unittest.IsolatedAsyncioTestCase):
+    """Tests for the _AgentCardResolver class."""
+
+    def setUp(self):
+        self.mock_client = AsyncMock(spec=httpx.AsyncClient)
+        self.resolver = _AgentCardResolver(self.mock_client)
+
+    @patch("aixtools.a2a.google_sdk.utils.ClientFactory")
+    def test_init(self, mock_client_factory_class):
+        """Test _AgentCardResolver initialization."""
+        mock_factory = MagicMock(spec=ClientFactory)
+        mock_client_factory_class.return_value = mock_factory
+        
+        resolver = _AgentCardResolver(self.mock_client)
+        
+        # Verify ClientFactory was created with correct config
+        mock_client_factory_class.assert_called_once()
+        call_args = mock_client_factory_class.call_args[0][0]
+        self.assertIsInstance(call_args, ClientConfig)
+        self.assertEqual(call_args.httpx_client, self.mock_client)
+        
+        # Verify attributes are set
+        self.assertEqual(resolver._httpx_client, self.mock_client)
+        self.assertEqual(resolver._a2a_client_factory, mock_factory)
+        self.assertEqual(resolver.clients, {})
+
+    @patch("aixtools.a2a.google_sdk.utils.RemoteAgentConnection")
+    def test_register_agent_card(self, mock_connection_class):
+        """Test registering an agent card."""
+        mock_card = MagicMock(spec=AgentCard)
+        mock_card.name = "test_agent"
+        
+        mock_client = MagicMock()
+        mock_factory = MagicMock(spec=ClientFactory)
+        mock_factory.create.return_value = mock_client
+        self.resolver._a2a_client_factory = mock_factory
+        
+        mock_connection = MagicMock(spec=RemoteAgentConnection)
+        mock_connection_class.return_value = mock_connection
+        
+        self.resolver.register_agent_card(mock_card)
+        
+        # Verify client was created
+        mock_factory.create.assert_called_once_with(mock_card)
+        
+        # Verify RemoteAgentConnection was created
+        mock_connection_class.assert_called_once_with(mock_card, mock_client)
+        
+        # Verify connection was stored
+        self.assertEqual(self.resolver.clients["test_agent"], mock_connection)
+
+    @patch("aixtools.a2a.google_sdk.utils.A2ACardResolver")
+    async def test_retrieve_card(self, mock_resolver_class):
+        """Test retrieving a card from an address."""
+        mock_resolver = AsyncMock(spec=A2ACardResolver)
+        mock_resolver_class.return_value = mock_resolver
+        
+        mock_card = MagicMock(spec=AgentCard)
+        mock_card.name = "test_agent"
+        mock_resolver.get_agent_card.return_value = mock_card
+        
+        with patch.object(self.resolver, 'register_agent_card') as mock_register:
+            await self.resolver.retrieve_card("http://test.com")
+            
+            # Verify resolver was created correctly (it tries the first card path)
+            mock_resolver_class.assert_called_with(self.mock_client, "http://test.com", "/.well-known/agent-card.json")
+            
+            # Verify card was retrieved
+            mock_resolver.get_agent_card.assert_called_once()
+            
+            # Verify card was registered
+            mock_register.assert_called_once_with(mock_card)
+
+    async def test_get_a2a_clients(self):
+        """Test getting A2A clients for multiple hosts."""
+        agent_hosts = ["http://agent1.com", "http://agent2.com"]
+        
+        with patch.object(self.resolver, 'retrieve_card') as mock_retrieve:
+            mock_retrieve.return_value = None  # Mock async function
+            
+            # Set up some mock clients
+            mock_connection1 = MagicMock(spec=RemoteAgentConnection)
+            mock_connection2 = MagicMock(spec=RemoteAgentConnection)
+            self.resolver.clients = {
+                "agent1": mock_connection1,
+                "agent2": mock_connection2
+            }
+            
+            result = await self.resolver.get_a2a_clients(agent_hosts)
+            
+            # Verify retrieve_card was called for each host
+            self.assertEqual(mock_retrieve.call_count, 2)
+            mock_retrieve.assert_any_call("http://agent1.com")
+            mock_retrieve.assert_any_call("http://agent2.com")
+            
+            # Verify result contains the clients
+            self.assertEqual(result, self.resolver.clients)
+
+
+class TestGetA2AClients(unittest.IsolatedAsyncioTestCase):
+    """Tests for the get_a2a_clients function."""
+
+    @patch("aixtools.a2a.google_sdk.utils._AgentCardResolver")
+    @patch("aixtools.a2a.google_sdk.utils.httpx.AsyncClient")
+    async def test_get_a2a_clients(self, mock_client_class, mock_resolver_class):
+        """Test the get_a2a_clients function."""
+        mock_client = AsyncMock()
+        mock_client_class.return_value = mock_client
+        
+        mock_resolver = AsyncMock(spec=_AgentCardResolver)
+        mock_resolver_class.return_value = mock_resolver
+        
+        mock_clients = {"agent1": MagicMock(), "agent2": MagicMock()}
+        mock_resolver.get_a2a_clients.return_value = mock_clients
+        
+        ctx = ("user123", "session456")
+        agent_hosts = ["http://agent1.com", "http://agent2.com"]
+        
+        result = await get_a2a_clients(ctx, agent_hosts)
+        
+        # Verify httpx client was created with correct headers
+        expected_headers = {
+            "user-id": "user123",
+            "session-id": "session456",
+        }
+        mock_client_class.assert_called_once_with(headers=expected_headers, timeout=60.0)
+        
+        # Verify resolver was created with the client
+        mock_resolver_class.assert_called_once_with(mock_client)
+        
+        # Verify get_a2a_clients was called with the hosts
+        mock_resolver.get_a2a_clients.assert_called_once_with(agent_hosts)
+        
+        # Verify result
+        self.assertEqual(result, mock_clients)
+
+
+class TestGetSessionIdTuple(unittest.TestCase):
+    """Tests for the get_session_id_tuple function."""
+
+    def test_get_session_id_tuple_with_headers(self):
+        """Test getting session ID tuple when headers are present."""
+        mock_context = MagicMock(spec=RequestContext)
+        mock_context.call_context.state = {
+            "headers": {
+                "user-id": "test_user",
+                "session-id": "test_session"
+            }
+        }
+        
+        result = get_session_id_tuple(mock_context)
+        
+        self.assertEqual(result, ("test_user", "test_session"))
+
+    def test_get_session_id_tuple_with_partial_headers(self):
+        """Test getting session ID tuple when only some headers are present."""
+        mock_context = MagicMock(spec=RequestContext)
+        mock_context.call_context.state = {
+            "headers": {
+                "user-id": "test_user"
+                # session-id is missing
+            }
+        }
+        
+        result = get_session_id_tuple(mock_context)
+        
+        self.assertEqual(result, ("test_user", DEFAULT_SESSION_ID))
+
+    def test_get_session_id_tuple_no_headers(self):
+        """Test getting session ID tuple when no headers are present."""
+        mock_context = MagicMock(spec=RequestContext)
+        mock_context.call_context.state = {}
+        
+        result = get_session_id_tuple(mock_context)
+        
+        self.assertEqual(result, (DEFAULT_USER_ID, DEFAULT_SESSION_ID))
+
+    def test_get_session_id_tuple_empty_headers(self):
+        """Test getting session ID tuple when headers dict is empty."""
+        mock_context = MagicMock(spec=RequestContext)
+        mock_context.call_context.state = {"headers": {}}
+        
+        result = get_session_id_tuple(mock_context)
+        
+        self.assertEqual(result, (DEFAULT_USER_ID, DEFAULT_SESSION_ID))
+
+
+if __name__ == '__main__':
+    unittest.main()