pygeai 0.2.7b51__py3-none-any.whl → 0.2.7b52__py3-none-any.whl

pygeai/tests/proxy/test_integration.py

@@ -0,0 +1,304 @@
+import unittest
+from unittest.mock import Mock, patch, MagicMock, AsyncMock
+import asyncio
+import uuid
+import json
+import sys
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../../..')))
+from pygeai.proxy.config import ProxySettingsManager
+from pygeai.proxy.tool import ProxiedTool
+from pygeai.proxy.clients import ProxyClient, ToolProxyData, ToolProxyJob, ToolProxyJobResult
+from pygeai.proxy.managers import ServerManager
+from pygeai.proxy.servers import A2AServer, MCPServer
+from types import SimpleNamespace
+
+
+class TestProxyIntegration(unittest.IsolatedAsyncioTestCase):
+    """
+    python -m unittest pygeai.tests.proxy.test_integration.TestProxyIntegration
+    """
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.settings = ProxySettingsManager()
+        self.test_uuid = uuid.uuid4()
+        self.test_affinity = uuid.uuid4()
+
+    def test_tool_to_proxy_data_integration(self):
+        """Test integration between ProxiedTool and ToolProxyData."""
+        # Create a tool
+        tool = ProxiedTool(
+            server_name="test_server",
+            name="test_tool",
+            description="Test tool description",
+            public_prefix="public.prefix",
+            input_schema={
+                "type": "object",
+                "properties": {
+                    "param1": {"type": "string"},
+                    "param2": {"type": "number"}
+                }
+            }
+        )
+        
+        # Create proxy data with the tool
+        proxy_data = ToolProxyData(
+            id=self.test_uuid,
+            name="Test Proxy",
+            description="Test Description",
+            affinity=self.test_affinity,
+            tools=[tool]
+        )
+        
+        # Convert to dictionary
+        data_dict = proxy_data.to_dict()
+        
+        # Verify the tool data is correctly included
+        self.assertEqual(len(data_dict["tools"]), 1)
+        tool_dict = data_dict["tools"][0]
+        
+        self.assertEqual(tool_dict["name"], "test_server__test_tool")
+        self.assertEqual(tool_dict["publicName"], "public.prefix.test_server__test_tool")
+        self.assertEqual(tool_dict["server"], "test_server")
+        self.assertEqual(tool_dict["description"], "Test tool description")
+        
+        # Verify the input schema is correctly formatted
+        input_schema = json.loads(tool_dict["inputSchema"])
+        self.assertEqual(input_schema["function"]["parameters"]["type"], "object")
+
+    def test_job_processing_integration(self):
+        """Test integration between job creation and result processing."""
+        # Create a job
+        job = ToolProxyJob(
+            id=uuid.uuid4(),
+            proxy_id=self.test_uuid,
+            proxy_status="active",
+            job_status="pending",
+            input=json.dumps({
+                "function": {
+                    "name": "test_server__test_tool",
+                    "arguments": '{"param1": "value1", "param2": 42}'
+                }
+            }),
+            server="test_server"
+        )
+        
+        # Simulate job execution
+        job.job_status = "completed"
+        job.output = "Execution completed successfully"
+        
+        # Create result
+        result = ToolProxyJobResult(success=True, job=job)
+        result_dict = result.to_dict()
+        
+        # Verify result structure
+        self.assertTrue(result_dict["success"])
+        self.assertEqual(result_dict["job"]["jobStatus"], "completed")
+        self.assertEqual(result_dict["job"]["output"], "Execution completed successfully")
+
+    @patch('pygeai.proxy.managers.ProxyClient')
+    async def test_server_manager_integration(self, mock_client_class):
+        """Test integration between ServerManager and its components."""
+        # Configure settings
+        self.settings.set_proxy_id(self.test_uuid)
+        self.settings.set_proxy_name("Test Proxy")
+        self.settings.set_proxy_description("Test Description")
+        self.settings.set_proxy_affinity(self.test_affinity)
+        
+        # Mock settings methods
+        self.settings.get_current_alias = Mock(return_value="default")
+        self.settings.get_api_key = Mock(return_value="test_api_key")
+        self.settings.get_base_url = Mock(return_value="https://api.example.com")
+        
+        # Create server configuration
+        servers_cfg = [
+            {
+                "name": "test_mcp_server",
+                "type": "mcp",
+                "command": "test_command",
+                "args": ["arg1", "arg2"]
+            }
+        ]
+        
+        # Create manager
+        manager = ServerManager(servers_cfg, self.settings)
+        
+        # Add a mock server and tool
+        tool = ProxiedTool(
+            server_name="test_mcp_server",
+            name="test_tool",
+            description="Test tool",
+            public_prefix="public.prefix",
+            input_schema={"type": "object"}
+        )
+        manager.tools["test_mcp_server__test_tool"] = tool
+        
+        # Mock client
+        mock_client = Mock()
+        mock_client.register.return_value = {"status": "registered"}
+        mock_client_class.return_value = mock_client
+        
+        # Add a mock server and tool
+        mock_server = AsyncMock()
+        mock_server.name = "test_mcp_server"
+        mock_server.exit_stack = AsyncMock()
+        mock_server.list_tools.return_value = []
+        
+        manager.servers["test_mcp_server"] = mock_server
+        
+        # Test client initialization
+        client = await manager._initialize_client()
+        
+        # Verifica que el mock fue llamado
+        mock_client.register.assert_called_once()
+        self.assertEqual(client, mock_client)
+        
+        # Verify registration data - get the actual arguments passed to register
+        call_args = mock_client.register.call_args
+        self.assertIsNotNone(call_args)
+        
+        # The register method is called with keyword arguments, so we need to access them properly
+        proxy_data = call_args.kwargs.get('proxy_data')
+        self.assertIsNotNone(proxy_data)
+        self.assertEqual(proxy_data.id, self.test_uuid)
+        self.assertEqual(proxy_data.name, "Test Proxy")
+        self.assertEqual(proxy_data.description, "Test Description")
+        self.assertEqual(proxy_data.affinity, self.test_affinity)
+        self.assertEqual(len(proxy_data.tools), 1)
+
+    @patch('pygeai.proxy.managers.Console')
+    async def test_tool_execution_integration(self, mock_console):
+        """Test integration of tool execution flow."""
+        # Create manager
+        servers_cfg = []
+        manager = ServerManager(servers_cfg, self.settings)
+        
+        # Add mock server
+        mock_server = AsyncMock()
+        mock_server.execute_tool.return_value = "execution_result"
+        manager.servers["test_server"] = mock_server
+        
+        # Add tool
+        tool = ProxiedTool(
+            server_name="test_server",
+            name="test_tool",
+            description="Test tool",
+            public_prefix="public.prefix",
+            input_schema={"type": "object"}
+        )
+        manager.tools["test_server__test_tool"] = tool
+        
+        # Test tool execution
+        result = await manager.execute_tool("test_server", "test_server__test_tool", {"param": "value"})
+        
+        # Verify execution
+        self.assertEqual(result, "execution_result")
+        mock_server.execute_tool.assert_called_once_with("test_tool", {"param": "value"}, 2, 1.0)
+
+    def test_function_call_parsing_integration(self):
+        """Test integration of function call parsing."""
+        # Create manager
+        servers_cfg = []
+        manager = ServerManager(servers_cfg, self.settings)
+        
+        # Test function call parsing
+        function_call_json = json.dumps({
+            "function": {
+                "name": "test_server__test_tool",
+                "arguments": '{"param1": "value1", "param2": 42}'
+            }
+        })
+        
+        name, arguments = manager.extract_function_call_info(function_call_json)
+        
+        # Verify parsing
+        self.assertEqual(name, "test_server__test_tool")
+        self.assertEqual(arguments, '{"param1": "value1", "param2": 42}')
+        
+        # Test with parsed arguments
+        parsed_args = json.loads(arguments)
+        self.assertEqual(parsed_args["param1"], "value1")
+        self.assertEqual(parsed_args["param2"], 42)
+
+    @patch('pygeai.proxy.clients.requests.Session')
+    def test_client_request_integration(self, mock_session):
+        """Test integration of client request handling."""
+        # Create client
+        client = ProxyClient("test_api_key", "https://api.example.com", self.test_uuid)
+        
+        # Mock response
+        mock_response = Mock()
+        mock_response.json.return_value = {"status": "success", "data": "test_data"}
+        mock_response.raise_for_status.return_value = None
+        mock_session.return_value.request.return_value = mock_response
+        
+        # Test request
+        result = client._make_request("GET", "/test")
+        
+        # Verify request
+        self.assertEqual(result, {"status": "success", "data": "test_data"})
+        mock_session.return_value.request.assert_called_once()
+
+    def test_settings_integration(self):
+        """Test integration of settings management."""
+        # Set various settings
+        self.settings.set_proxy_id(self.test_uuid)
+        self.settings.set_proxy_name("Test Proxy")
+        self.settings.set_proxy_description("Test Description")
+        self.settings.set_proxy_affinity(self.test_affinity)
+        
+        # Get complete configuration
+        config = self.settings.get_proxy_config()
+        
+        # Verify configuration
+        expected = {
+            "id": str(self.test_uuid),
+            "name": "Test Proxy",
+            "description": "Test Description",
+            "affinity": str(self.test_affinity)
+        }
+        self.assertEqual(config, expected)
+
+    def test_tool_formatting_integration(self):
+        """Test integration of tool formatting for different scenarios."""
+        # Test public tool
+        public_tool = ProxiedTool(
+            server_name="test_server",
+            name="public_tool",
+            description="Public tool description",
+            public_prefix="public.prefix",
+            input_schema={"type": "object"}
+        )
+        
+        # Test private tool
+        private_tool = ProxiedTool(
+            server_name="test_server",
+            name="private_tool",
+            description="Private tool description",
+            public_prefix=None,
+            input_schema={"type": "object"}
+        )
+        
+        # Create proxy data with both tools
+        proxy_data = ToolProxyData(
+            id=self.test_uuid,
+            name="Test Proxy",
+            tools=[public_tool, private_tool]
+        )
+        
+        # Convert to dictionary
+        data_dict = proxy_data.to_dict()
+        
+        # Verify public tool has publicName
+        public_tool_dict = data_dict["tools"][0]
+        self.assertIn("publicName", public_tool_dict)
+        self.assertEqual(public_tool_dict["publicName"], "public.prefix.test_server__public_tool")
+        
+        # Verify private tool doesn't have publicName
+        private_tool_dict = data_dict["tools"][1]
+        self.assertNotIn("publicName", private_tool_dict)
+
+
+if __name__ == '__main__':
+    unittest.main() 

pygeai/tests/proxy/test_managers.py

@@ -0,0 +1,312 @@
+import unittest
+from unittest.mock import Mock, patch, MagicMock, AsyncMock
+import asyncio
+import json
+import uuid
+from pygeai.proxy.managers import ServerManager
+from pygeai.proxy.config import ProxySettingsManager
+from pygeai.proxy.clients import ProxyClient, ToolProxyData, ToolProxyJobResult
+from pygeai.proxy.tool import ProxiedTool
+
+
+class TestServerManager(unittest.IsolatedAsyncioTestCase):
+    """
+    python -m unittest pygeai.tests.proxy.test_managers.TestServerManager
+    """
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.servers_cfg = [
+            {
+                "name": "test_mcp_server",
+                "type": "mcp",
+                "command": "test_command",
+                "args": ["arg1", "arg2"]
+            },
+            {
+                "name": "test_a2a_server",
+                "type": "a2a",
+                "url": "https://test.a2a.com",
+                "headers": {"Authorization": "Bearer token"}
+            }
+        ]
+        self.settings = ProxySettingsManager()
+        self.manager = ServerManager(self.servers_cfg, self.settings)
+
+    def test_initialization(self):
+        """Test manager initialization."""
+        self.assertEqual(self.manager.servers_cfg, self.servers_cfg)
+        self.assertEqual(self.manager.settings, self.settings)
+        self.assertEqual(self.manager.servers, {})
+        self.assertEqual(self.manager.tools, {})
+
+    @patch('pygeai.proxy.managers.MCPServer')
+    @patch('pygeai.proxy.managers.A2AServer')
+    @patch('pygeai.proxy.managers.Console')
+    async def test_initialize_servers_success(self, mock_console, mock_a2a_server_class, mock_mcp_server_class):
+        """Test successful server initialization."""
+        # Mock server instances
+        mock_mcp_server = AsyncMock()
+        mock_mcp_server.name = "test_mcp_server"
+        mock_mcp_server.exit_stack = AsyncMock()
+        mock_mcp_server.list_tools.return_value = []
+        
+        mock_a2a_server = AsyncMock()
+        mock_a2a_server.name = "test_a2a_server"
+        mock_a2a_server.exit_stack = AsyncMock()
+        mock_a2a_server.list_tools.return_value = []
+        
+        mock_mcp_server_class.return_value = mock_mcp_server
+        mock_a2a_server_class.return_value = mock_a2a_server
+        
+        await self.manager._initialize_servers()
+        
+        # Verify servers were created and initialized
+        mock_mcp_server_class.assert_called_once_with(
+            "test_mcp_server", self.servers_cfg[0], self.settings
+        )
+        mock_a2a_server_class.assert_called_once_with(
+            "test_a2a_server", self.servers_cfg[1], self.settings
+        )
+        
+        mock_mcp_server.initialize.assert_called_once()
+        mock_a2a_server.initialize.assert_called_once()
+
+    @patch('pygeai.proxy.managers.MCPServer')
+    @patch('pygeai.proxy.managers.Console')
+    async def test_initialize_servers_invalid_type(self, mock_console, mock_mcp_server_class):
+        """Test server initialization with invalid server type."""
+        invalid_config = [{"name": "invalid_server", "type": "invalid_type"}]
+        manager = ServerManager(invalid_config, self.settings)
+        
+        with self.assertRaises(ValueError, msg="Invalid server type: invalid_type"):
+            await manager._initialize_servers()
+
+    @patch('pygeai.proxy.managers.MCPServer')
+    @patch('pygeai.proxy.managers.Console')
+    async def test_initialize_servers_initialization_error(self, mock_console, mock_mcp_server_class):
+        """Test server initialization error."""
+        mock_server = AsyncMock()
+        mock_server.name = "test_server"
+        mock_server.initialize.side_effect = Exception("Initialization failed")
+        mock_mcp_server_class.return_value = mock_server
+        
+        with self.assertRaises(Exception):
+            await self.manager._initialize_servers()
+
+    @patch('pygeai.proxy.managers.ProxyClient')
+    @patch('pygeai.proxy.managers.Console')
+    async def test_initialize_client_success(self, mock_console, mock_client_class):
+        """Test successful client initialization."""
+        # Mock settings methods
+        self.settings.get_current_alias = Mock(return_value="default")
+        self.settings.get_api_key = Mock(return_value="test_api_key")
+        self.settings.get_base_url = Mock(return_value="https://api.example.com")
+        self.settings.get_proxy_id = Mock(return_value=uuid.uuid4())
+        self.settings.get_proxy_name = Mock(return_value="Test Proxy")
+        self.settings.get_proxy_description = Mock(return_value="Test Description")
+        self.settings.get_proxy_affinity = Mock(return_value=uuid.uuid4())
+        
+        # Mock client
+        mock_client = Mock()
+        mock_client.register.return_value = {"status": "success"}
+        mock_client_class.return_value = mock_client
+        
+        # Add some tools to the manager
+        tool = ProxiedTool(
+            server_name="test_server",
+            name="test_tool",
+            description="Test tool",
+            public_prefix="public.prefix",
+            input_schema={"type": "object"}
+        )
+        self.manager.tools["test_server__test_tool"] = tool
+        
+        result = await self.manager._initialize_client()
+        
+        self.assertEqual(result, mock_client)
+        mock_client.register.assert_called_once()
+
+    @patch('pygeai.proxy.managers.ProxyClient')
+    @patch('pygeai.proxy.managers.Console')
+    async def test_initialize_client_connection_error(self, mock_console, mock_client_class):
+        """Test client initialization with connection error."""
+        # Mock settings methods
+        self.settings.get_current_alias = Mock(return_value="default")
+        self.settings.get_api_key = Mock(return_value="test_api_key")
+        self.settings.get_base_url = Mock(return_value="https://api.example.com")
+        self.settings.get_proxy_id = Mock(return_value=uuid.uuid4())
+        
+        # Mock client to raise connection error
+        mock_client = Mock()
+        mock_client.register.side_effect = ConnectionError("Connection failed")
+        mock_client_class.return_value = mock_client
+        
+        with self.assertRaises(ConnectionError):
+            await self.manager._initialize_client()
+
+    async def test_execute_tool_server_not_found(self):
+        """Test executing tool with non-existent server."""
+        with self.assertRaises(RuntimeError, msg="Server non_existent not found"):
+            await self.manager.execute_tool("non_existent", "test_tool", {})
+
+    async def test_execute_tool_tool_not_found(self):
+        """Test executing non-existent tool."""
+        # Add a server but no tools
+        self.manager.servers["test_server"] = Mock()
+        
+        with self.assertRaises(RuntimeError, msg="Tool non_existent not found"):
+            await self.manager.execute_tool("test_server", "non_existent", {})
+
+    @patch('pygeai.proxy.managers.Console')
+    async def test_execute_tool_success(self, mock_console):
+        """Test successful tool execution."""
+        # Mock server
+        mock_server = AsyncMock()
+        mock_server.execute_tool.return_value = "success"
+        self.manager.servers["test_server"] = mock_server
+        
+        # Add tool
+        tool = ProxiedTool(
+            server_name="test_server",
+            name="test_tool",
+            description="Test tool",
+            public_prefix="public.prefix",
+            input_schema={"type": "object"}
+        )
+        self.manager.tools["test_server__test_tool"] = tool
+        
+        result = await self.manager.execute_tool("test_server", "test_server__test_tool", {"param": "value"})
+        
+        self.assertEqual(result, "success")
+        mock_server.execute_tool.assert_called_once_with("test_tool", {"param": "value"}, 2, 1.0)
+
+    @patch('pygeai.proxy.managers.Console')
+    async def test_execute_tool_with_custom_retries(self, mock_console):
+        """Test tool execution with custom retry parameters."""
+        mock_server = AsyncMock()
+        mock_server.execute_tool.return_value = "success"
+        self.manager.servers["test_server"] = mock_server
+        
+        tool = ProxiedTool(
+            server_name="test_server",
+            name="test_tool",
+            description="Test tool",
+            public_prefix="public.prefix",
+            input_schema={"type": "object"}
+        )
+        self.manager.tools["test_server__test_tool"] = tool
+        
+        await self.manager.execute_tool("test_server", "test_server__test_tool", {}, retries=5, delay=2.0)
+        
+        mock_server.execute_tool.assert_called_once_with("test_tool", {}, 5, 2.0)
+
+    @patch('pygeai.proxy.managers.Console')
+    async def test_execute_tool_execution_error(self, mock_console):
+        """Test tool execution with error."""
+        mock_server = AsyncMock()
+        mock_server.execute_tool.side_effect = RuntimeError("Tool execution failed")
+        self.manager.servers["test_server"] = mock_server
+        
+        tool = ProxiedTool(
+            server_name="test_server",
+            name="test_tool",
+            description="Test tool",
+            public_prefix="public.prefix",
+            input_schema={"type": "object"}
+        )
+        self.manager.tools["test_server__test_tool"] = tool
+        
+        with self.assertRaises(Exception, msg="Failed to execute tool test_server__test_tool on server test_server: Tool execution failed"):
+            await self.manager.execute_tool("test_server", "test_server__test_tool", {})
+
+    def test_extract_function_call_info_success(self):
+        """Test successful function call info extraction."""
+        function_call_json = json.dumps({
+            "function": {
+                "name": "test_function",
+                "arguments": '{"param": "value"}'
+            }
+        })
+        
+        name, arguments = self.manager.extract_function_call_info(function_call_json)
+        
+        self.assertEqual(name, "test_function")
+        self.assertEqual(arguments, '{"param": "value"}')
+
+    def test_extract_function_call_info_invalid_json(self):
+        """Test function call info extraction with invalid JSON."""
+        name, arguments = self.manager.extract_function_call_info("invalid json")
+        
+        self.assertIsNone(name)
+        self.assertIsNone(arguments)
+
+    def test_extract_function_call_info_missing_keys(self):
+        """Test function call info extraction with missing keys."""
+        function_call_json = json.dumps({"other_key": "value"})
+        
+        name, arguments = self.manager.extract_function_call_info(function_call_json)
+        
+        self.assertIsNone(name)
+        self.assertIsNone(arguments)
+
+    @patch('pygeai.proxy.managers.ServerManager._initialize_servers')
+    @patch('pygeai.proxy.managers.ServerManager._initialize_client')
+    @patch('pygeai.proxy.managers.Console')
+    async def test_start_success(self, mock_console, mock_init_client, mock_init_servers):
+        """Test successful start method."""
+        # Mock client
+        mock_client = Mock()
+        mock_client.dequeue.return_value = []
+        mock_init_client.return_value = mock_client
+        
+        # Mock settings
+        self.settings.get_current_alias = Mock(return_value="default")
+        self.settings.get_proxy_id = Mock(return_value=uuid.uuid4())
+        
+        # Create a task to run start method
+        task = asyncio.create_task(self.manager.start())
+        
+        # Let it run for a short time
+        await asyncio.sleep(0.1)
+        
+        # Cancel the task
+        task.cancel()
+        
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+        
+        # Verify initialization was called
+        mock_init_servers.assert_called_once()
+        mock_init_client.assert_called_once()
+
+    @patch('pygeai.proxy.managers.ServerManager._initialize_servers')
+    @patch('pygeai.proxy.managers.ServerManager._initialize_client')
+    @patch('pygeai.proxy.managers.Console')
+    async def test_start_client_initialization_error(self, mock_console, mock_init_client, mock_init_servers):
+        """Test start method with client initialization error."""
+        mock_init_client.side_effect = ConnectionError("Connection failed")
+        
+        # Create a task to run start method
+        task = asyncio.create_task(self.manager.start())
+        
+        # Let it run for a short time
+        await asyncio.sleep(0.1)
+        
+        # Cancel the task
+        task.cancel()
+        
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+        
+        # Verify initialization was called
+        mock_init_servers.assert_called_once()
+        mock_init_client.assert_called_once()
+
+
+if __name__ == '__main__':
+    unittest.main()