Fixing test cases for pooling

Signed-off-by: Mohan Lakshmaiah <mohalaks@in.ibm.com>

Author: MohanLakshmaiah

mcpgateway/cache/session_registry.py

@@ -1235,7 +1235,12 @@ def _db_cleanup() -> int:
                     transport = session_data['transport']
                     try:
                         if not await transport.is_connected():
-                            await self.remove_session(session_id)
+                            if pooled:
+                                # For pooled sessions, remove from registry but don't disconnect
+                                await self.remove_session_from_registry_only(session_id)
+                            else:
+                                # For non-pooled sessions, full removal with disconnect
+                                await self.remove_session(session_id)
                             continue
 
                         # Refresh session in database

mcpgateway/main.py

@@ -3967,13 +3967,13 @@ async def websocket_endpoint(websocket: WebSocket):
     Args:
         websocket: The WebSocket connection instance.
     """
+    transport = None
+    proxy_user = None
+    token = None
     try:
         # Authenticate WebSocket connection
         if settings.mcp_client_auth_enabled or settings.auth_required:
             # Extract auth from query params or headers
-            token = None
-            proxy_user = None
-
             # Try to get token from query parameter
             if "token" in websocket.query_params:
                 token = websocket.query_params["token"]
@@ -4001,9 +4001,6 @@ async def websocket_endpoint(websocket: WebSocket):
                     await websocket.close(code=1008, reason="Invalid authentication")
                     return
 
-        await websocket.accept()
-        logger.info("WebSocket connection accepted")
-
         # Identify user and server for pooling key
         user_id = proxy_user or "anonymous"
         server_id = websocket.query_params.get("server_id", "default-server")
@@ -4013,21 +4010,16 @@ async def websocket_endpoint(websocket: WebSocket):
         transport = None
         if await should_use_session_pooling(server_id):
             # Use existing or create pooled session
-            # Note: WebSocket transport needs the actual WebSocket object, so pooling works differently
             transport = WebSocketTransport(websocket, pooled=True, pool_key=f"{user_id}:{server_id}")
             await transport.connect()
             await session_registry.add_session(transport.session_id, transport, pooled=True)
-            logger.info(
-                f"Created pooled WebSocket session for user={user_id}, server={server_id}, session={transport.session_id}"
-            )
+            logger.info(f"Created pooled WebSocket session for user={user_id}, server={server_id}, session={transport.session_id}")
         else:
             # Fallback: create new transport
             transport = WebSocketTransport(websocket)
             await transport.connect()
             await session_registry.add_session(transport.session_id, transport)
-            logger.info(
-                f"Created new WebSocket session for user={user_id}, server={server_id}, session={transport.session_id}"
-            )
+            logger.info(f"Created new WebSocket session for user={user_id}, server={server_id}, session={transport.session_id}")
 
         while True:
             try:
@@ -4058,10 +4050,8 @@ async def websocket_endpoint(websocket: WebSocket):
                 break
     except WebSocketDisconnect:
         logger.info("WebSocket disconnected")
-        # Cleanup pooled session if needed
         if transport and hasattr(transport, '_pooled') and transport._pooled:
-            # For pooled sessions, we don't immediately remove from registry
-            # They get cleaned up by the pool's background task
+            # For pooled sessions, we don't immediately remove from registry. They get cleaned up by the pool's background task
             pass
         else:
             # For non-pooled sessions, remove from registry

mcpgateway/services/server_service.py

@@ -15,6 +15,7 @@
 import asyncio
 from datetime import datetime, timezone
 from typing import Any, AsyncGenerator, Dict, List, Optional
+import builtins
 
 # Third-Party
 import httpx
@@ -51,7 +52,7 @@ class ServerNotFoundError(ServerError):
     """Raised when a requested server is not found."""
 
 
-class PermissionError(ServerError):
+class PermissionError(builtins.PermissionError, ServerError):
     """Raised when a user does not have permission to perform an action on a server."""
 
 
@@ -647,7 +648,7 @@ async def get_server(self, db: Session, server_id: str) -> ServerRead:
             raise ServerNotFoundError(f"Server not found: {server_id}")
 
         try:
-            effective_strategy = await self.get_session_strategy(db, server_id)
+            effective_strategy = await self.get_session_strategy(db, server_id, server=server)
             logger.debug(f"Server {server_id} effective session strategy: {effective_strategy}")
         except Exception as e:
             logger.warning(f"Could not determine session strategy for server {server_id}: {e}")
@@ -998,13 +999,14 @@ async def delete_server(self, db: Session, server_id: str, user_email: Optional[
             if not server:
                 raise ServerNotFoundError(f"Server not found: {server_id}")
 
-            # Check ownership if user_email provided
+            # Always perform ownership check if user_email is provided
             if user_email:
                 # First-Party
                 from mcpgateway.services.permission_service import PermissionService  # pylint: disable=import-outside-toplevel
 
                 permission_service = PermissionService(db)
-                if not await permission_service.check_resource_ownership(user_email, server):
+                can_delete = await permission_service.check_resource_ownership(user_email, server)
+                if not can_delete:
                     raise PermissionError("Only the owner can delete this server")
 
             server_info = {"id": server.id, "name": server.name}
@@ -1147,7 +1149,7 @@ async def _notify_server_deleted(self, server_info: Dict[str, Any]) -> None:
         }
         await self._publish_event(event)
 
-    async def get_session_strategy(self, db: Session, server_id: str) -> str:
+    async def get_session_strategy(self, db: Session, server_id: str, server: Optional[DbServer] = None) -> str:
         """Determine effective session strategy for server.
 
         This method resolves the session strategy for a specific server, taking into account:
@@ -1180,7 +1182,12 @@ async def get_session_strategy(self, db: Session, server_id: str) -> str:
             >>> result == settings.session_pool_strategy
             True
         """
-        server = db.get(DbServer, server_id)
+        # server = db.get(DbServer, server_id)
+        # if not server:
+        #     raise ServerNotFoundError(f"Server not found: {server_id}")
+        # Allow callers to pass an already-loaded server object to avoid repeated DB lookups.
+        if server is None:
+            server = db.get(DbServer, server_id)
         if not server:
             raise ServerNotFoundError(f"Server not found: {server_id}")
 

mcpgateway/transports/websocket_transport.py

@@ -128,8 +128,9 @@ async def connect(self) -> None:
             >>> mock_ws.accept.called
             True
         """
-        await self._websocket.accept()
-        self._connected = True
+        if not self._connected:
+            await self._websocket.accept()
+            self._connected = True
 
         # Start ping task
         if settings.websocket_ping_interval > 0:

tests/unit/mcpgateway/services/test_resource_ownership.py

@@ -202,18 +202,37 @@ async def test_delete_server_non_owner_denied(self, server_service, mock_db_sess
         mock_server = MagicMock(spec=Server)
         mock_server.id = "server-1"
         mock_server.owner_email = "owner@example.com"
-
+        mock_server.team_id = None
+        mock_server.visibility = "private"
+        mock_server.name = "Test Server"
+        mock_server.session_pooling_strategy = "inherit"
+        
         mock_db_session.get.return_value = mock_server
+        mock_db_session.rollback = MagicMock()
+        mock_db_session.commit = MagicMock()
+        mock_db_session.delete = MagicMock()
 
         with patch("mcpgateway.services.permission_service.PermissionService") as mock_perm_service_class:
             mock_perm_service = mock_perm_service_class.return_value
             mock_perm_service.check_resource_ownership = AsyncMock(return_value=False)
-
-            with pytest.raises(PermissionError, match="Only the owner can delete this server"):
-                await server_service.delete_server(mock_db_session, "server-1", user_email="other@example.com")
-
+            server_service._notify_server_deleted = AsyncMock()
+
+            try:
+                with pytest.raises(PermissionError, match="Only the owner can delete this server"):
+                    await server_service.delete_server(mock_db_session, "server-1", user_email="other@example.com")
+            except AssertionError as e:
+                # This will help us understand if we're getting a different error message
+                print(f"Test failed because: {e}")
+                raise
+            except Exception as e:
+                print(f"Unexpected error: {e}")
+                raise
+
+            # Verify the expectations
             mock_db_session.delete.assert_not_called()
-
+            mock_db_session.rollback.assert_called_once()
+            mock_perm_service.check_resource_ownership.assert_called_once_with("other@example.com", mock_server)
+            mock_db_session.commit.assert_not_called()
 
 class TestToolServiceOwnership:
     """Test ownership checks in ToolService delete/update methods."""