Reworked fix to avoid test execution hanging

Author: aidbutlr

mcpgateway/plugins/__init__.py

@@ -23,7 +23,7 @@
 import logging
 import random
 import time
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+from typing import Any, Callable, Literal, Optional, TYPE_CHECKING, Union
 
 # Third-Party
 from cpex.framework import ObservabilityProvider, TenantPluginManager
@@ -102,9 +102,10 @@ class _TeamBindingChangeMsg(BaseModel):
 
 def _get_invalidation_hmac_key() -> Optional[bytes]:
     """Return the HMAC signing key for pub/sub messages, or None if unconfigured."""
-    from mcpgateway.config import settings  # pylint: disable=import-outside-toplevel
     from pydantic import SecretStr  # pylint: disable=import-outside-toplevel
 
+    from mcpgateway.config import settings  # pylint: disable=import-outside-toplevel
+
     secret = settings.jwt_secret_key
     if not secret:
         return None
@@ -491,21 +492,22 @@ async def _plugin_invalidation_listener() -> None:
     consecutive_failures = 0
 
     while True:
+        pubsub = None
         try:
             client = await _redis()
             if not client:
                 _logger.debug("Plugin invalidation listener: Redis unavailable, retrying in 10s")
                 await asyncio.sleep(10)
                 continue
 
-            async with client.pubsub() as pubsub:
-                await pubsub.subscribe(_REDIS_INVALIDATION_CHANNEL)
-                _logger.info("Plugin invalidation listener subscribed to %s", _REDIS_INVALIDATION_CHANNEL)
-                backoff = 1.0
-                consecutive_failures = 0
+            pubsub = client.pubsub()
+            await pubsub.subscribe(_REDIS_INVALIDATION_CHANNEL)
+            _logger.info("Plugin invalidation listener subscribed to %s", _REDIS_INVALIDATION_CHANNEL)
+            backoff = 1.0
+            consecutive_failures = 0
 
-                async for message in pubsub.listen():
-                    await _handle_invalidation_message(message)
+            async for message in pubsub.listen():
+                await _handle_invalidation_message(message)
 
         except asyncio.CancelledError:
             _logger.info("Plugin invalidation listener cancelled")
@@ -524,6 +526,14 @@ async def _plugin_invalidation_listener() -> None:
             )
             await asyncio.sleep(delay)
             backoff = min(max_backoff, backoff * 2)
+        finally:
+            # Always release the pubsub connection back to the pool.
+            # Using wait_for shields aclose() from an in-flight CancelledError:
+            if pubsub is not None:
+                try:
+                    await asyncio.wait_for(pubsub.aclose(), timeout=2.0)
+                except Exception:
+                    pass
 
 
 async def start_plugin_invalidation_listener() -> None: