Add unit tests for ServerClassificationService

- Implement tests for initialization, classification logic, leader election, polling decisions, and Redis state management.
- Cover various scenarios including hot/cold classification, tie-breaking logic, and service lifecycle management.
- Ensure comprehensive testing of the server classification algorithm and its integration with the GatewayService.

Signed-off-by: Lang-Akshay <akshay.shinde26@ibm.com>

Author: Lang-Akshay

mcpgateway/services/gateway_service.py

@@ -3351,10 +3351,14 @@ async def _check_single_gateway_health(self, gateway: DbGateway, user_email: Opt
 
         # Hot/cold classification: Check if this server should be health-checked now
         if self._classification_service:
-            should_check = await self._classification_service.should_poll_server(gateway_url, "health")
-            if not should_check:
-                logger.debug(f"Skipping health check for {SecurityValidator.sanitize_log_message(gateway_name)}: " f"not yet due based on hot/cold classification")
-                return
+            try:
+                should_check = await self._classification_service.should_poll_server(gateway_url, "health")
+                if not should_check:
+                    logger.debug(f"Skipping health check for {SecurityValidator.sanitize_log_message(gateway_name)}: " f"not yet due based on hot/cold classification")
+                    return
+            except Exception as e:
+                # Fail open: proceed with health check if classification check fails
+                logger.warning(f"Classification check failed for {gateway_name}, proceeding with health check (fail-open): {e}")
 
         # Create span for individual gateway health check
         with create_span(
@@ -3556,9 +3560,14 @@ def get_httpx_client_factory(
                     if settings.auto_refresh_servers:
                         # Hot/cold classification: Check if this server should have tools refreshed now
                         if self._classification_service:
-                            should_auto_refresh = await self._classification_service.should_poll_server(gateway_url, "tools")
-                            if not should_auto_refresh:
-                                logger.debug(f"Skipping auto-refresh for {SecurityValidator.sanitize_log_message(gateway_name)}: " f"not yet due based on hot/cold classification")
+                            try:
+                                should_auto_refresh = await self._classification_service.should_poll_server(gateway_url, "tools")
+                                if not should_auto_refresh:
+                                    logger.debug(f"Skipping auto-refresh for {SecurityValidator.sanitize_log_message(gateway_name)}: " f"not yet due based on hot/cold classification")
+                            except Exception as e:
+                                # Fail open: proceed with auto-refresh if classification check fails
+                                logger.warning(f"Classification check failed for {gateway_name}, proceeding with auto-refresh (fail-open): {e}")
+                                should_auto_refresh = True
                         else:
                             should_auto_refresh = True
 

mcpgateway/services/server_classification_service.py

@@ -227,15 +227,15 @@ def _classify_servers_from_pool(self, pool: MCPSessionPool, all_gateway_urls: Li
         Returns:
             ClassificationResult with hot/cold servers and metadata
         """
-        N = len(all_gateway_urls)
-        hot_cap = floor(0.20 * N)
+        total_servers = len(all_gateway_urls)
+        hot_cap = floor(0.20 * total_servers)
 
         # Step 3: Extract server usage from pooled sessions
         server_metrics: Dict[str, ServerUsageMetrics] = {}
 
         # Iterate over pool._pools (Dict[PoolKey, Queue[PooledSession]])
         # PoolKey = (user_identity, url, identity_hash, transport_type, gateway_id)
-        for pool_key, session_queue in pool._pools.items():
+        for pool_key, session_queue in pool._pools.items():  # pylint: disable=protected-access
             url = pool_key[1]  # Extract server URL from pool key
 
             if url not in server_metrics:
@@ -244,7 +244,7 @@ def _classify_servers_from_pool(self, pool: MCPSessionPool, all_gateway_urls: Li
             # Process each pooled session in the queue
             try:
                 # Access queue items (asyncio.Queue has internal _queue deque)
-                sessions_list = list(session_queue._queue) if hasattr(session_queue, "_queue") else []
+                sessions_list = list(session_queue._queue) if hasattr(session_queue, "_queue") else []  # pylint: disable=protected-access
 
                 for session in sessions_list:
                     # PooledSession has: last_used, use_count
@@ -258,7 +258,7 @@ def _classify_servers_from_pool(self, pool: MCPSessionPool, all_gateway_urls: Li
                 continue
 
         # Count active sessions from _active dict
-        for pool_key, active_set in pool._active.items():
+        for pool_key, active_set in pool._active.items():  # pylint: disable=protected-access
             url = pool_key[1]
             if url in server_metrics:
                 server_metrics[url].active_session_count += len(active_set)
@@ -293,7 +293,7 @@ def _classify_servers_from_pool(self, pool: MCPSessionPool, all_gateway_urls: Li
         return ClassificationResult(
             hot_servers=hot_servers,
             cold_servers=cold_servers,
-            metadata=ClassificationMetadata(N=N, hot_cap=hot_cap, hot_actual=hot_actual, eligible_count=eligible_count, timestamp=time.time(), underutilized_reason=underutilized_reason),
+            metadata=ClassificationMetadata(N=total_servers, hot_cap=hot_cap, hot_actual=hot_actual, eligible_count=eligible_count, timestamp=time.time(), underutilized_reason=underutilized_reason),
         )
 
     async def _get_all_gateway_urls(self) -> List[str]:
@@ -306,11 +306,11 @@ async def _get_all_gateway_urls(self) -> List[str]:
         from sqlalchemy import select
 
         # First-Party
-        from mcpgateway.db import DbGateway, SessionLocal
+        from mcpgateway.db import Gateway, SessionLocal
 
         try:
             with SessionLocal() as db:
-                result = db.execute(select(DbGateway.url).where(DbGateway.enabled == True))  # noqa: E712
+                result = db.execute(select(Gateway.url).where(Gateway.enabled is True))
                 urls = [row[0] for row in result]
                 return urls
         except Exception as e: