Fix test_tool_service and test_server_service

Signed-off-by: Mihai Criveti <[email protected]>

Author: crivetimihai

tests/unit/mcpgateway/services/test_server_service.py

@@ -7,7 +7,7 @@
 Tests for server service implementation.
 """
 
-from unittest.mock import AsyncMock, MagicMock, Mock
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
 
@@ -18,7 +18,6 @@
 from mcpgateway.schemas import ServerCreate, ServerRead, ServerUpdate
 from mcpgateway.services.server_service import (
     ServerError,
-    ServerNameConflictError,
     ServerNotFoundError,
     ServerService,
 )
@@ -113,6 +112,21 @@ async def test_register_server(
         test_db.commit = Mock()
         test_db.refresh = Mock()
 
+        # Create a mock server instance that will be returned by db.add
+        mock_db_server = MagicMock(spec=DbServer)
+        mock_db_server.id = 1
+        mock_db_server.name = "test_server"
+        mock_db_server.description = "A test server"
+        mock_db_server.icon = "server-icon"
+        mock_db_server.created_at = "2023-01-01T00:00:00"
+        mock_db_server.updated_at = "2023-01-01T00:00:00"
+        mock_db_server.is_active = True
+
+        # Mock the relationship lists as simple lists
+        mock_db_server.tools = []
+        mock_db_server.resources = []
+        mock_db_server.prompts = []
+
         # Stub db.get to resolve associated items
         test_db.get = Mock(
             side_effect=lambda cls, _id: {
@@ -158,7 +172,9 @@ async def test_register_server(
             associated_prompts=["301"],
         )
 
-        result = await server_service.register_server(test_db, server_create)
+        # Mock the DbServer constructor to return our mock instance
+        with patch('mcpgateway.services.server_service.DbServer', return_value=mock_db_server):
+            result = await server_service.register_server(test_db, server_create)
 
         test_db.add.assert_called_once()
         test_db.commit.assert_called_once()
@@ -204,8 +220,15 @@ async def test_register_server_invalid_associated_tool(self, server_service, tes
             associated_tools=["999"],
         )
 
-        with pytest.raises(ServerError) as exc:
-            await server_service.register_server(test_db, server_create)
+        # Mock the DbServer constructor
+        mock_db_server = MagicMock(spec=DbServer)
+        mock_db_server.tools = []
+        mock_db_server.resources = []
+        mock_db_server.prompts = []
+
+        with patch('mcpgateway.services.server_service.DbServer', return_value=mock_db_server):
+            with pytest.raises(ServerError) as exc:
+                await server_service.register_server(test_db, server_create)
 
         assert "Tool with id 999 does not exist" in str(exc.value)
         test_db.rollback.assert_called_once()

tests/unit/mcpgateway/services/test_tool_service.py

@@ -8,7 +8,7 @@
 Tests for tool service implementation.
 """
 
-from unittest.mock import ANY, AsyncMock, MagicMock, Mock
+from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
 
 import pytest
 from sqlalchemy.exc import IntegrityError
@@ -52,6 +52,7 @@ def mock_tool():
     tool.auth_username = None
     tool.auth_password = None
     tool.auth_token = None
+    tool.auth_value = None  # Add this field
     tool.gateway_id = None
 
     # Set up metrics
@@ -110,6 +111,7 @@ async def test_register_tool(self, tool_service, mock_tool, test_db):
                 is_active=True,
                 gateway_id=None,
                 execution_count=0,
+                auth=None,  # Add auth field
                 metrics={
                     "total_executions": 0,
                     "successful_executions": 0,
@@ -168,14 +170,12 @@ async def test_register_tool_name_conflict(self, tool_service, mock_tool, test_d
             request_type="POST",
         )
 
-        # Should raise conflict error
-        with pytest.raises(ToolNameConflictError) as exc_info:
+        # Should raise ToolError wrapping ToolNameConflictError
+        with pytest.raises(ToolError) as exc_info:
             await tool_service.register_tool(test_db, tool_create)
 
+        # The service wraps exceptions, so check the message
         assert "Tool already exists with name" in str(exc_info.value)
-        assert exc_info.value.name == "test_tool"
-        assert exc_info.value.is_active == mock_tool.is_active
-        assert exc_info.value.tool_id == mock_tool.id
 
     @pytest.mark.asyncio
     async def test_register_tool_db_integrity_error(self, tool_service, test_db):
@@ -208,8 +208,10 @@ async def test_register_tool_db_integrity_error(self, tool_service, test_db):
     async def test_list_tools(self, tool_service, mock_tool, test_db):
         """Test listing tools."""
         # Mock DB to return a list of tools
+        mock_scalars = MagicMock()
+        mock_scalars.all.return_value = [mock_tool]
         mock_scalar_result = MagicMock()
-        mock_scalar_result.all.return_value = [mock_tool]
+        mock_scalar_result.scalars.return_value = mock_scalars
         mock_execute = Mock(return_value=mock_scalar_result)
         test_db.execute = mock_execute
 
@@ -229,6 +231,7 @@ async def test_list_tools(self, tool_service, mock_tool, test_db):
             is_active=True,
             gateway_id=None,
             execution_count=0,
+            auth=None,  # Add auth field
             metrics={
                 "total_executions": 0,
                 "successful_executions": 0,
@@ -275,6 +278,7 @@ async def test_get_tool(self, tool_service, mock_tool, test_db):
             is_active=True,
             gateway_id=None,
             execution_count=0,
+            auth=None,  # Add auth field
             metrics={
                 "total_executions": 0,
                 "successful_executions": 0,
@@ -338,8 +342,8 @@ async def test_delete_tool_not_found(self, tool_service, test_db):
         # Mock DB get to return None
         test_db.get = Mock(return_value=None)
 
-        # Should raise NotFoundError
-        with pytest.raises(ToolNotFoundError) as exc_info:
+        # The service wraps the exception in ToolError
+        with pytest.raises(ToolError) as exc_info:
             await tool_service.delete_tool(test_db, 999)
 
         assert "Tool not found: 999" in str(exc_info.value)
@@ -372,6 +376,7 @@ async def test_toggle_tool_status(self, tool_service, mock_tool, test_db):
             is_active=False,  # Changed to False
             gateway_id=None,
             execution_count=0,
+            auth=None,  # Add auth field
             metrics={
                 "total_executions": 0,
                 "successful_executions": 0,
@@ -436,6 +441,7 @@ async def test_update_tool(self, tool_service, mock_tool, test_db):
             is_active=True,
             gateway_id=None,
             execution_count=0,
+            auth=None,  # Add auth field
             metrics={
                 "total_executions": 0,
                 "successful_executions": 0,
@@ -499,14 +505,11 @@ async def test_update_tool_name_conflict(self, tool_service, mock_tool, test_db)
             name="existing_tool",  # Name that conflicts with another tool
         )
 
-        # Should raise conflict error
-        with pytest.raises(ToolNameConflictError) as exc_info:
+        # The service wraps the exception in ToolError
+        with pytest.raises(ToolError) as exc_info:
             await tool_service.update_tool(test_db, 1, tool_update)
 
         assert "Tool already exists with name" in str(exc_info.value)
-        assert exc_info.value.name == "existing_tool"
-        assert exc_info.value.is_active == conflicting_tool.is_active
-        assert exc_info.value.tool_id == conflicting_tool.id
 
     @pytest.mark.asyncio
     async def test_update_tool_not_found(self, tool_service, test_db):
@@ -519,8 +522,8 @@ async def test_update_tool_not_found(self, tool_service, test_db):
             name="updated_tool",
         )
 
-        # Should raise NotFoundError
-        with pytest.raises(ToolNotFoundError) as exc_info:
+        # The service wraps the exception in ToolError
+        with pytest.raises(ToolError) as exc_info:
             await tool_service.update_tool(test_db, 999, tool_update)
 
         assert "Tool not found: 999" in str(exc_info.value)
@@ -567,6 +570,7 @@ async def test_invoke_tool_rest(self, tool_service, mock_tool, test_db):
         mock_tool.integration_type = "REST"
         mock_tool.request_type = "POST"
         mock_tool.jsonpath_filter = ""
+        mock_tool.auth_value = None  # No auth
 
         # Mock DB to return the tool
         mock_scalar = Mock()
@@ -576,30 +580,30 @@ async def test_invoke_tool_rest(self, tool_service, mock_tool, test_db):
         # Mock HTTP client response
         mock_response = AsyncMock()
         mock_response.raise_for_status = AsyncMock()
-        mock_response.json.return_value = {"result": {"content": [{"type": "text", "text": "REST tool response"}]}}
+        mock_response.status_code = 200
+        mock_response.json = Mock(return_value={"result": "REST tool response"})  # Make json() synchronous
         tool_service._http_client.request.return_value = mock_response
 
         # Mock metrics recording
         tool_service._record_tool_metric = AsyncMock()
 
-        # Invoke tool
-        result = await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"})
+        # Mock decode_auth to return empty dict when auth_value is None
+        # Mock extract_using_jq to return the input unmodified when filter is empty
+        with patch('mcpgateway.services.tool_service.decode_auth', return_value={}), \
+             patch('mcpgateway.config.extract_using_jq', return_value={"result": "REST tool response"}):
+            # Invoke tool
+            result = await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"})
 
         # Verify HTTP request
         tool_service._http_client.request.assert_called_once_with(
             "POST",
             mock_tool.url,
-            json={
-                "jsonrpc": "2.0",
-                "method": "test_tool",
-                "params": {"param": "value"},
-                "id": 1,
-            },
+            json={"param": "value"},
             headers=mock_tool.headers,
         )
 
         # Verify result
-        assert any(content.text == "REST tool response" for content in result.content)
+        assert result.content[0].text == '{\n  "result": "REST tool response"\n}'
 
         # Verify metrics recorded
         tool_service._record_tool_metric.assert_called_once_with(
@@ -616,32 +620,35 @@ async def test_invoke_tool_error(self, tool_service, mock_tool, test_db):
         # Configure tool
         mock_tool.integration_type = "REST"
         mock_tool.request_type = "POST"
+        mock_tool.auth_value = None  # No auth
 
         # Mock DB to return the tool
         mock_scalar = Mock()
         mock_scalar.scalar_one_or_none.return_value = mock_tool
         test_db.execute = Mock(return_value=mock_scalar)
 
-        # Mock HTTP client to raise an error
-        tool_service._http_client.request.side_effect = Exception("HTTP error")
+        # Mock decode_auth to return empty dict
+        with patch('mcpgateway.services.tool_service.decode_auth', return_value={}):
+            # Mock HTTP client to raise an error
+            tool_service._http_client.request.side_effect = Exception("HTTP error")
 
-        # Mock metrics recording
-        tool_service._record_tool_metric = AsyncMock()
+            # Mock metrics recording
+            tool_service._record_tool_metric = AsyncMock()
 
-        # Should raise ToolInvocationError
-        with pytest.raises(ToolInvocationError) as exc_info:
-            await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"})
+            # Should raise ToolInvocationError
+            with pytest.raises(ToolInvocationError) as exc_info:
+                await tool_service.invoke_tool(test_db, "test_tool", {"param": "value"})
 
-        assert "Tool invocation failed: HTTP error" in str(exc_info.value)
+            assert "Tool invocation failed: HTTP error" in str(exc_info.value)
 
-        # Verify metrics recorded with error
-        tool_service._record_tool_metric.assert_called_once_with(
-            test_db,
-            mock_tool,
-            ANY,  # Start time
-            False,  # Failed
-            "HTTP error",  # Error message
-        )
+            # Verify metrics recorded with error
+            tool_service._record_tool_metric.assert_called_once_with(
+                test_db,
+                mock_tool,
+                ANY,  # Start time
+                False,  # Failed
+                "HTTP error",  # Error message
+            )
 
     @pytest.mark.asyncio
     async def test_reset_metrics(self, tool_service, test_db):