Added doctests to utility modules (#400)

Signed-off-by: Manav Gupta <[email protected]>

Author: manavgup

mcpgateway/cache/resource_cache.py

@@ -10,6 +10,29 @@
 - TTL-based expiration
 - Maximum size limit with LRU eviction
 - Thread-safe operations
+
+Doctest examples
+----------------
+>>> from mcpgateway.cache.resource_cache import ResourceCache
+>>> cache = ResourceCache(max_size=2, ttl=1)
+>>> cache.set('a', 1)
+>>> cache.get('a')
+1
+>>> import time
+>>> time.sleep(1.1)
+>>> cache.get('a') is None
+True
+>>> cache.set('a', 1)
+>>> cache.set('b', 2)
+>>> cache.set('c', 3)  # LRU eviction
+>>> sorted(cache._cache.keys())
+['b', 'c']
+>>> cache.delete('b')
+>>> cache.get('b') is None
+True
+>>> cache.clear()
+>>> cache.get('a') is None
+True
 """
 
 # Standard
@@ -32,13 +55,36 @@ class CacheEntry:
 
 
 class ResourceCache:
-    """Resource content cache with TTL expiration.
+    """
+    Resource content cache with TTL expiration.
 
     Attributes:
         max_size: Maximum number of entries
         ttl: Time-to-live in seconds
         _cache: Cache storage
         _lock: Async lock for thread safety
+
+    Doctest:
+    >>> from mcpgateway.cache.resource_cache import ResourceCache
+    >>> cache = ResourceCache(max_size=2, ttl=1)
+    >>> cache.set('a', 1)
+    >>> cache.get('a')
+    1
+    >>> import time
+    >>> time.sleep(1.1)
+    >>> cache.get('a') is None
+    True
+    >>> cache.set('a', 1)
+    >>> cache.set('b', 2)
+    >>> cache.set('c', 3)  # LRU eviction
+    >>> sorted(cache._cache.keys())
+    ['b', 'c']
+    >>> cache.delete('b')
+    >>> cache.get('b') is None
+    True
+    >>> cache.clear()
+    >>> cache.get('a') is None
+    True
     """
 
     def __init__(self, max_size: int = 1000, ttl: int = 3600):
@@ -65,13 +111,25 @@ async def shutdown(self) -> None:
         self.clear()
 
     def get(self, key: str) -> Optional[Any]:
-        """Get value from cache.
+        """
+        Get value from cache.
 
         Args:
             key: Cache key
 
         Returns:
             Cached value or None if not found/expired
+
+        Doctest:
+        >>> from mcpgateway.cache.resource_cache import ResourceCache
+        >>> cache = ResourceCache(max_size=2, ttl=1)
+        >>> cache.set('a', 1)
+        >>> cache.get('a')
+        1
+        >>> import time
+        >>> time.sleep(1.1)
+        >>> cache.get('a') is None
+        True
         """
         if key not in self._cache:
             return None
@@ -89,11 +147,19 @@ def get(self, key: str) -> Optional[Any]:
         return entry.value
 
     def set(self, key: str, value: Any) -> None:
-        """Set value in cache.
+        """
+        Set value in cache.
 
         Args:
             key: Cache key
             value: Value to cache
+
+        Doctest:
+        >>> from mcpgateway.cache.resource_cache import ResourceCache
+        >>> cache = ResourceCache(max_size=2, ttl=1)
+        >>> cache.set('a', 1)
+        >>> cache.get('a')
+        1
         """
         now = time.time()
 
@@ -107,15 +173,34 @@ def set(self, key: str, value: Any) -> None:
         self._cache[key] = CacheEntry(value=value, expires_at=now + self.ttl, last_access=now)
 
     def delete(self, key: str) -> None:
-        """Delete value from cache.
+        """
+        Delete value from cache.
 
         Args:
             key: Cache key to delete
+
+        Doctest:
+        >>> from mcpgateway.cache.resource_cache import ResourceCache
+        >>> cache = ResourceCache()
+        >>> cache.set('a', 1)
+        >>> cache.delete('a')
+        >>> cache.get('a') is None
+        True
         """
         self._cache.pop(key, None)
 
     def clear(self) -> None:
-        """Clear all cached entries."""
+        """
+        Clear all cached entries.
+
+        Doctest:
+        >>> from mcpgateway.cache.resource_cache import ResourceCache
+        >>> cache = ResourceCache()
+        >>> cache.set('a', 1)
+        >>> cache.clear()
+        >>> cache.get('a') is None
+        True
+        """
         self._cache.clear()
 
     async def _cleanup_loop(self) -> None:

mcpgateway/cache/session_registry.py

@@ -7,6 +7,21 @@
 
 This module provides a registry for SSE sessions with support for distributed deployment
 using Redis or SQLAlchemy as optional backends for shared state between workers.
+
+Doctest examples (memory backend only)
+--------------------------------------
+>>> from mcpgateway.cache.session_registry import SessionRegistry
+>>> class DummyTransport:
+...     pass
+>>> reg = SessionRegistry(backend='memory')
+>>> import asyncio
+>>> asyncio.run(reg.add_session('sid', DummyTransport()))
+>>> t = asyncio.run(reg.get_session('sid'))
+>>> isinstance(t, DummyTransport)
+True
+>>> asyncio.run(reg.remove_session('sid'))
+>>> asyncio.run(reg.get_session('sid')) is None
+True
 """
 
 # Standard
@@ -106,7 +121,8 @@ def __init__(
 
 
 class SessionRegistry(SessionBackend):
-    """Registry for SSE sessions with optional distributed state.
+    """
+    Registry for SSE sessions with optional distributed state.
 
     Supports three backend modes:
     - memory: In-memory storage (default, no dependencies)
@@ -115,6 +131,20 @@ class SessionRegistry(SessionBackend):
 
     In distributed mode (redis/database), session existence is tracked in the shared
     backend while transports themselves remain local to each worker process.
+
+    Doctest (memory backend only):
+    >>> from mcpgateway.cache.session_registry import SessionRegistry
+    >>> class DummyTransport:
+    ...     pass
+    >>> reg = SessionRegistry(backend='memory')
+    >>> import asyncio
+    >>> asyncio.run(reg.add_session('sid', DummyTransport()))
+    >>> t = asyncio.run(reg.get_session('sid'))
+    >>> isinstance(t, DummyTransport)
+    True
+    >>> asyncio.run(reg.remove_session('sid'))
+    >>> asyncio.run(reg.get_session('sid')) is None
+    True
     """
 
     def __init__(

mcpgateway/utils/create_jwt_token.py

@@ -16,12 +16,24 @@
     $ python3 jwt_cli.py
 
 Library:
-```python
-from mcpgateway.utils.create_jwt_token import create_jwt_token, get_jwt_token
-
-# inside async context
-jwt = await create_jwt_token({"username": "alice"})
-```
+    from mcpgateway.utils.create_jwt_token import create_jwt_token, get_jwt_token
+
+    # inside async context
+    jwt = await create_jwt_token({"username": "alice"})
+
+Doctest examples
+----------------
+>>> from mcpgateway.utils import create_jwt_token as jwt_util
+>>> jwt_util.settings.jwt_secret_key = 'secret'
+>>> jwt_util.settings.jwt_algorithm = 'HS256'
+>>> token = jwt_util._create_jwt_token({'sub': 'alice'}, expires_in_minutes=1, secret='secret', algorithm='HS256')
+>>> import jwt
+>>> jwt.decode(token, 'secret', algorithms=['HS256'])['sub'] == 'alice'
+True
+>>> import asyncio
+>>> t = asyncio.run(jwt_util.create_jwt_token({'sub': 'bob'}, expires_in_minutes=1, secret='secret', algorithm='HS256'))
+>>> jwt.decode(t, 'secret', algorithms=['HS256'])['sub'] == 'bob'
+True
 """
 
 # Future
@@ -67,7 +79,8 @@ def _create_jwt_token(
     secret: str = DEFAULT_SECRET,
     algorithm: str = DEFAULT_ALGO,
 ) -> str:
-    """Return a signed JWT string (synchronous, timezone-aware).
+    """
+    Return a signed JWT string (synchronous, timezone-aware).
 
     Args:
         data: Dictionary containing payload data to encode in the token.
@@ -78,6 +91,15 @@ def _create_jwt_token(
 
     Returns:
         The JWT token string.
+
+    Doctest:
+    >>> from mcpgateway.utils import create_jwt_token as jwt_util
+    >>> jwt_util.settings.jwt_secret_key = 'secret'
+    >>> jwt_util.settings.jwt_algorithm = 'HS256'
+    >>> token = jwt_util._create_jwt_token({'sub': 'alice'}, expires_in_minutes=1, secret='secret', algorithm='HS256')
+    >>> import jwt
+    >>> jwt.decode(token, 'secret', algorithms=['HS256'])['sub'] == 'alice'
+    True
     """
     payload = data.copy()
     if expires_in_minutes > 0:
@@ -98,7 +120,8 @@ async def create_jwt_token(
     secret: str = DEFAULT_SECRET,
     algorithm: str = DEFAULT_ALGO,
 ) -> str:
-    """Async facade for historic code. Internally synchronous-almost instant.
+    """
+    Async facade for historic code. Internally synchronous-almost instant.
 
     Args:
         data: Dictionary containing payload data to encode in the token.
@@ -109,6 +132,16 @@ async def create_jwt_token(
 
     Returns:
         The JWT token string.
+
+    Doctest:
+    >>> from mcpgateway.utils import create_jwt_token as jwt_util
+    >>> jwt_util.settings.jwt_secret_key = 'secret'
+    >>> jwt_util.settings.jwt_algorithm = 'HS256'
+    >>> import asyncio
+    >>> t = asyncio.run(jwt_util.create_jwt_token({'sub': 'bob'}, expires_in_minutes=1, secret='secret', algorithm='HS256'))
+    >>> import jwt
+    >>> jwt.decode(t, 'secret', algorithms=['HS256'])['sub'] == 'bob'
+    True
     """
     return _create_jwt_token(data, expires_in_minutes, secret, algorithm)
 

mcpgateway/utils/db_isready.py

@@ -45,15 +45,43 @@
 Shell ::
 
     python3 db_isready.py
-    python3 db_isready.py --database-url "postgresql://user:pw@db:5432/mcp" \
-                         --max-tries 20 --interval 1 --timeout 1
+    python3 db_isready.py --database-url "sqlite:///./mcp.db" --max-tries 2 --interval 1 --timeout 1
 
 Python ::
 
-    from db_isready import wait_for_db_ready
-
-    await wait_for_db_ready()          # asynchronous
-    wait_for_db_ready(sync=True)       # synchronous / blocking
+    from mcpgateway.utils.db_isready import wait_for_db_ready
+
+    # Synchronous/blocking
+    wait_for_db_ready(sync=True)
+
+    # Asynchronous
+    import asyncio
+    asyncio.run(wait_for_db_ready())
+
+Doctest examples
+----------------
+>>> from mcpgateway.utils.db_isready import wait_for_db_ready
+>>> import logging
+>>> class DummyLogger:
+...     def __init__(self): self.infos = []
+...     def info(self, msg): self.infos.append(msg)
+...     def debug(self, msg): pass
+...     def error(self, msg): pass
+...     @property
+...     def handlers(self): return [True]
+>>> import sys
+>>> sys.modules['sqlalchemy'] = type('sqlalchemy', (), {
+...     'create_engine': lambda *a, **k: type('E', (), {'connect': lambda self: type('C', (), {'execute': lambda self, q: 1, '__enter__': lambda self: self, '__exit__': lambda self, exc_type, exc_val, exc_tb: None})()})(),
+...     'text': lambda q: q,
+...     'engine': type('engine', (), {'Engine': object, 'URL': object, 'url': type('url', (), {'make_url': lambda u: type('U', (), {'get_backend_name': lambda self: "sqlite"})()}),}),
+...     'exc': type('exc', (), {'OperationalError': Exception})
+... })
+>>> wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=1, interval=1, timeout=1, logger=DummyLogger(), sync=True)
+>>> try:
+...     wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=0, interval=1, timeout=1, logger=DummyLogger(), sync=True)
+... except RuntimeError as e:
+...     print('error')
+error
 """
 
 # Future
@@ -173,7 +201,8 @@ def wait_for_db_ready(
     logger: Optional[logging.Logger] = None,
     sync: bool = False,
 ) -> None:
-    """Block until the database replies to ``SELECT 1``.
+    """
+    Block until the database replies to ``SELECT 1``.
 
     The helper can be awaited **asynchronously** *or* called in *blocking*
     mode by passing ``sync=True``.
@@ -194,6 +223,30 @@ def wait_for_db_ready(
     Raises:
         RuntimeError: If *invalid* parameters are supplied or the database is
             still unavailable after the configured number of attempts.
+
+    Doctest:
+    >>> from mcpgateway.utils.db_isready import wait_for_db_ready
+    >>> import logging
+    >>> class DummyLogger:
+    ...     def __init__(self): self.infos = []
+    ...     def info(self, msg): self.infos.append(msg)
+    ...     def debug(self, msg): pass
+    ...     def error(self, msg): pass
+    ...     @property
+    ...     def handlers(self): return [True]
+    >>> import sys
+    >>> sys.modules['sqlalchemy'] = type('sqlalchemy', (), {
+    ...     'create_engine': lambda *a, **k: type('E', (), {'connect': lambda self: type('C', (), {'execute': lambda self, q: 1, '__enter__': lambda self: self, '__exit__': lambda self, exc_type, exc_val, exc_tb: None})()})(),
+    ...     'text': lambda q: q,
+    ...     'engine': type('engine', (), {'Engine': object, 'URL': object, 'url': type('url', (), {'make_url': lambda u: type('U', (), {'get_backend_name': lambda self: "sqlite"})()}),}),
+    ...     'exc': type('exc', (), {'OperationalError': Exception})
+    ... })
+    >>> wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=1, interval=1, timeout=1, logger=DummyLogger(), sync=True)
+    >>> try:
+    ...     wait_for_db_ready(database_url='sqlite:///./mcp.db', max_tries=0, interval=1, timeout=1, logger=DummyLogger(), sync=True)
+    ... except RuntimeError as e:
+    ...     print('error')
+    error
     """
 
     log = logger or logging.getLogger("db_isready")