From f75b68366fae6ae60d4eb08017e9c5387cac27f1 Mon Sep 17 00:00:00 2001 From: Yash Pawar Date: Tue, 16 Jun 2026 02:06:46 +0530 Subject: [PATCH 1/2] feat: add Redis persistence support for OAuth token and client storage --- .env.example | 12 +- compose.yaml | 25 ++ pyproject.toml | 1 + template_mcp_server/src/oauth/service.py | 104 ++++--- template_mcp_server/src/settings.py | 48 +++ template_mcp_server/src/storage/__init__.py | 8 +- template_mcp_server/src/storage/base.py | 97 ++++++ .../src/storage/redis_service.py | 290 ++++++++++++++++++ .../src/storage/storage_service.py | 3 +- tests/test_redis_service.py | 143 +++++++++ tests/test_storage_init.py | 2 +- 11 files changed, 684 insertions(+), 49 deletions(-) create mode 100644 template_mcp_server/src/storage/base.py create mode 100644 template_mcp_server/src/storage/redis_service.py create mode 100644 tests/test_redis_service.py diff --git a/.env.example b/.env.example index 9386323..d0989bd 100644 --- a/.env.example +++ b/.env.example @@ -30,9 +30,19 @@ SSO_INTROSPECTION_MODE=rfc7662 # SSO_INTROSPECTION_URL=https://oauth2.googleapis.com/tokeninfo SSO_INTROSPECTION_URL=https://sso.example.com/realms/myrealm/protocol/openid-connect/token/introspect -# --- PostgreSQL (required only if ENABLE_AUTH=True, for token storage) --- +# --- Storage Provider (required only if ENABLE_AUTH=True) --- +# Valid values: postgres, redis +STORAGE_TYPE=postgres + +# --- PostgreSQL (required if STORAGE_TYPE=postgres) --- POSTGRES_HOST=localhost POSTGRES_PORT=5432 POSTGRES_DB=postgres_db POSTGRES_USER=postgres_user POSTGRES_PASSWORD=postgres_password + +# --- Redis (required if STORAGE_TYPE=redis) --- +REDIS_HOST=localhost +REDIS_PORT=6379 +REDIS_PASSWORD= +REDIS_DB=0 diff --git a/compose.yaml b/compose.yaml index 73ea4ac..745dbd2 100644 --- a/compose.yaml +++ b/compose.yaml @@ -20,6 +20,23 @@ services: networks: - template-network + redis: + image: redis:7-alpine + container_name: template-mcp-redis + ports: + - "6379:6379" + volumes: + - redis_data:/data + restart: always + healthcheck: + test: ["CMD", "redis-cli", "ping"] + interval: 10s + timeout: 5s + retries: 5 + start_period: 10s + networks: + - template-network + template-mcp-server: build: context: . @@ -37,9 +54,15 @@ services: - POSTGRES_DB=postgres - POSTGRES_USER=postgres - POSTGRES_PASSWORD=postgres + - STORAGE_TYPE=postgres + - REDIS_HOST=redis + - REDIS_PORT=6379 + - REDIS_DB=0 depends_on: postgres: condition: service_healthy + redis: + condition: service_healthy restart: always healthcheck: test: [ "CMD", "curl", "-f", "-k", "http://0.0.0.0:5001/health"] @@ -52,6 +75,8 @@ services: volumes: postgres_data: driver: local + redis_data: + driver: local networks: template-network: diff --git a/pyproject.toml b/pyproject.toml index 5df9ff9..f12f287 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "psycopg==3.2.3", "itsdangerous==2.2.0", "requests-oauthlib==2.0.0", + "redis>=5.0.1", ] [project.optional-dependencies] diff --git a/template_mcp_server/src/oauth/service.py b/template_mcp_server/src/oauth/service.py index 2ae67dc..ec89f33 100644 --- a/template_mcp_server/src/oauth/service.py +++ b/template_mcp_server/src/oauth/service.py @@ -16,13 +16,15 @@ from typing import Any, Dict, List, Optional from template_mcp_server.src.settings import settings +from template_mcp_server.src.storage.base import BaseStorageService +from template_mcp_server.src.storage.redis_service import RedisStorageService from template_mcp_server.src.storage.storage_service import StorageService from template_mcp_server.utils.pylogger import get_python_logger logger = get_python_logger(settings.PYTHON_LOG_LEVEL) # Global storage service for backward compatibility during transition -_storage_service: Optional[StorageService] = None +_storage_service: Optional[BaseStorageService] = None def generate_random_string(length: int = 32) -> str: @@ -66,7 +68,7 @@ def verify_code_challenge(code_verifier: str, code_challenge: str) -> bool: class OAuthService: """OAuth service that manages OAuth 2.0 operations with dependency injection.""" - def __init__(self, storage_service: StorageService): + def __init__(self, storage_service: BaseStorageService): """Initialize OAuth service with storage dependency. Args: @@ -239,13 +241,13 @@ async def get_storage_status(self) -> Dict[str, Any]: # Backward compatibility functions - will be deprecated in future versions -async def get_storage_service() -> StorageService: +async def get_storage_service() -> BaseStorageService: """Get the initialized storage service. Note: Storage service must be initialized via initialize_storage() during startup. Returns: - StorageService: The initialized storage service + BaseStorageService: The initialized storage service Raises: RuntimeError: If storage service hasn't been initialized @@ -380,15 +382,15 @@ async def get_storage_status() -> Dict[str, Any]: return await service.get_storage_status() -async def initialize_storage() -> StorageService: +async def initialize_storage() -> BaseStorageService: """Initialize the storage service. Call this during application startup. Returns: - StorageService: The initialized storage service + BaseStorageService: The initialized storage service Raises: - ValueError: If PostgreSQL configuration is missing - ConnectionError: If PostgreSQL connection fails + ValueError: If configuration is missing + ConnectionError: If connection fails """ global _storage_service @@ -396,46 +398,58 @@ async def initialize_storage() -> StorageService: logger.warning("Storage service already initialized") return _storage_service - logger.info("Initializing PostgreSQL storage service") - - # Validate required configuration - if not all( - [ - settings.POSTGRES_HOST, - settings.POSTGRES_PORT, - settings.POSTGRES_DB, - settings.POSTGRES_USER, - ] - ): - missing = [ - name - for name, value in [ - ("POSTGRES_HOST", settings.POSTGRES_HOST), - ("POSTGRES_PORT", settings.POSTGRES_PORT), - ("POSTGRES_DB", settings.POSTGRES_DB), - ("POSTGRES_USER", settings.POSTGRES_USER), + if settings.STORAGE_TYPE == "redis": + logger.info("Initializing Redis storage service") + _storage_service = RedisStorageService( + host=settings.REDIS_HOST or "localhost", + port=settings.REDIS_PORT, + password=settings.REDIS_PASSWORD, + db=settings.REDIS_DB, + ) + else: + logger.info("Initializing PostgreSQL storage service") + + # Validate required configuration + if not all( + [ + settings.POSTGRES_HOST, + settings.POSTGRES_PORT, + settings.POSTGRES_DB, + settings.POSTGRES_USER, + ] + ): + missing = [ + name + for name, value in [ + ("POSTGRES_HOST", settings.POSTGRES_HOST), + ("POSTGRES_PORT", settings.POSTGRES_PORT), + ("POSTGRES_DB", settings.POSTGRES_DB), + ("POSTGRES_USER", settings.POSTGRES_USER), + ] + if not value ] - if not value - ] - raise ValueError( - f"Missing required PostgreSQL configuration: {', '.join(missing)}" + raise ValueError( + f"Missing required PostgreSQL configuration: {', '.join(missing)}" + ) + + # Create and connect storage service + # Type assertions are safe here because we validated required fields above + _storage_service = StorageService( + host=str(settings.POSTGRES_HOST), + port=int(settings.POSTGRES_PORT) + if settings.POSTGRES_PORT is not None + else 5432, + database=str(settings.POSTGRES_DB), + username=str(settings.POSTGRES_USER), + password=settings.POSTGRES_PASSWORD or "", + pool_size=settings.POSTGRES_POOL_SIZE, + max_connections=settings.POSTGRES_MAX_CONNECTIONS, ) - # Create and connect storage service - # Type assertions are safe here because we validated required fields above - _storage_service = StorageService( - host=str(settings.POSTGRES_HOST), - port=int(settings.POSTGRES_PORT) - if settings.POSTGRES_PORT is not None - else 5432, - database=str(settings.POSTGRES_DB), - username=str(settings.POSTGRES_USER), - password=settings.POSTGRES_PASSWORD or "", - pool_size=settings.POSTGRES_POOL_SIZE, - max_connections=settings.POSTGRES_MAX_CONNECTIONS, - ) await _storage_service.connect() - logger.info("PostgreSQL storage service initialized successfully") + logger.info( + f"{settings.STORAGE_TYPE.capitalize()} storage service initialized successfully" + ) return _storage_service @@ -444,7 +458,7 @@ async def cleanup_storage() -> None: """Cleanup storage service. Call this during application shutdown.""" global _storage_service if _storage_service is not None: - logger.info("Disconnecting from PostgreSQL...") + logger.info("Disconnecting from storage...") await _storage_service.disconnect() _storage_service = None logger.info("Storage service cleanup complete") diff --git a/template_mcp_server/src/settings.py b/template_mcp_server/src/settings.py index bc7df91..b6b6512 100644 --- a/template_mcp_server/src/settings.py +++ b/template_mcp_server/src/settings.py @@ -220,6 +220,54 @@ class Settings(BaseSettings): }, ) + # Storage Configuration + STORAGE_TYPE: Literal["postgres", "redis"] = Field( + default="postgres", + json_schema_extra={ + "env": "STORAGE_TYPE", + "description": "Storage backend type: postgres or redis", + "example": "postgres", + "enum": ["postgres", "redis"], + }, + ) + + # Redis Configuration + REDIS_HOST: Optional[str] = Field( + default="localhost", + json_schema_extra={ + "env": "REDIS_HOST", + "description": "Redis host address", + "example": "localhost", + }, + ) + REDIS_PORT: int = Field( + default=6379, + ge=1024, + le=65535, + json_schema_extra={ + "env": "REDIS_PORT", + "description": "Redis port number", + "example": 6379, + }, + ) + REDIS_PASSWORD: Optional[str] = Field( + default=None, + json_schema_extra={ + "env": "REDIS_PASSWORD", + "description": "Redis password", + "example": "secretpassword", + "sensitive": True, + }, + ) + REDIS_DB: int = Field( + default=0, + json_schema_extra={ + "env": "REDIS_DB", + "description": "Redis database index", + "example": 0, + }, + ) + # PostgreSQL Configuration POSTGRES_HOST: Optional[str] = Field( default=None, diff --git a/template_mcp_server/src/storage/__init__.py b/template_mcp_server/src/storage/__init__.py index 57cb219..a2a2bbd 100644 --- a/template_mcp_server/src/storage/__init__.py +++ b/template_mcp_server/src/storage/__init__.py @@ -1 +1,7 @@ -"""PostgreSQL storage service for the Template MCP Server.""" +"""Storage services for the Template MCP Server.""" + +from template_mcp_server.src.storage.base import BaseStorageService +from template_mcp_server.src.storage.redis_service import RedisStorageService +from template_mcp_server.src.storage.storage_service import StorageService + +__all__ = ["BaseStorageService", "StorageService", "RedisStorageService"] diff --git a/template_mcp_server/src/storage/base.py b/template_mcp_server/src/storage/base.py new file mode 100644 index 0000000..2f0044c --- /dev/null +++ b/template_mcp_server/src/storage/base.py @@ -0,0 +1,97 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional + + +class BaseStorageService(ABC): + """Abstract base class for storage services.""" + + @abstractmethod + async def connect(self) -> None: + """Establish connection to storage backend.""" + pass + + @abstractmethod + async def disconnect(self) -> None: + """Close connection to storage backend.""" + pass + + @abstractmethod + async def is_healthy(self) -> bool: + """Check if storage backend is healthy.""" + pass + + @abstractmethod + async def get_status(self) -> Dict[str, Any]: + """Get storage service status.""" + pass + + @abstractmethod + async def get_client_by_name_and_redirect_uris( + self, client_name: str, redirect_uris: List[str] + ) -> Optional[Dict[str, Any]]: + """Find an existing client by name and redirect URIs.""" + pass + + @abstractmethod + async def store_client(self, client_data: Dict[str, Any]) -> bool: + """Store a new OAuth client.""" + pass + + @abstractmethod + async def get_client(self, client_id: str) -> Optional[Dict[str, Any]]: + """Get a client by ID.""" + pass + + @abstractmethod + async def store_authorization_code( + self, code: str, code_data: Dict[str, Any] + ) -> bool: + """Store an authorization code.""" + pass + + @abstractmethod + async def get_authorization_code(self, code: str) -> Optional[Dict[str, Any]]: + """Get authorization code data.""" + pass + + @abstractmethod + async def update_authorization_code_token( + self, code: str, snowflake_token: Dict[str, Any] + ) -> bool: + """Update authorization code with Snowflake token.""" + pass + + @abstractmethod + async def delete_authorization_code(self, code: str) -> bool: + """Delete an authorization code.""" + pass + + @abstractmethod + async def store_access_token(self, token: str, token_data: Dict[str, Any]) -> bool: + """Store an access token.""" + pass + + @abstractmethod + async def get_access_token(self, token: str) -> Optional[Dict[str, Any]]: + """Get access token data.""" + pass + + @abstractmethod + async def delete_access_token(self, token: str) -> bool: + """Delete an access token.""" + pass + + @abstractmethod + async def store_refresh_token(self, token: str, token_data: Dict[str, Any]) -> bool: + """Store a refresh token.""" + pass + + @abstractmethod + async def get_refresh_token(self, token: str) -> Optional[Dict[str, Any]]: + """Get refresh token data.""" + pass + + @abstractmethod + async def delete_refresh_token(self, token: str) -> bool: + """Delete a refresh token.""" + pass diff --git a/template_mcp_server/src/storage/redis_service.py b/template_mcp_server/src/storage/redis_service.py new file mode 100644 index 0000000..0eb4b50 --- /dev/null +++ b/template_mcp_server/src/storage/redis_service.py @@ -0,0 +1,290 @@ +"""Redis storage service for the Template MCP Server.""" + +import hashlib +import json +import time +from typing import Any, Dict, List, Optional + +import redis.asyncio as redis + +from template_mcp_server.src.storage.base import BaseStorageService +from template_mcp_server.utils.pylogger import get_python_logger + +logger = get_python_logger() + + +def _get_client_index_key(client_name: str, redirect_uris: List[str]) -> str: + """Generate a stable hash key for a client name and redirect URIs combination.""" + normalized_uris = sorted(redirect_uris) + hash_input = json.dumps( + {"name": client_name, "uris": normalized_uris}, sort_keys=True + ) + hash_val = hashlib.sha256(hash_input.encode("utf-8")).hexdigest() + return f"client_index:{hash_val}" + + +class RedisStorageService(BaseStorageService): + """Redis storage service for persistent data storage. + + This service provides direct Redis storage functionality. + It automatically uses Redis built-in TTL for tokens and auth codes. + """ + + def __init__( + self, + host: str = "localhost", + port: int = 6379, + password: Optional[str] = None, + db: int = 0, + ): + """Initialize the Redis storage service. + + Args: + host: Redis host + port: Redis port + password: Redis password + db: Database index + """ + self.host = host + self.port = port + self.password = password + self.db = db + self.redis: Optional[redis.Redis] = None + + async def connect(self) -> None: + """Establish connection to Redis.""" + try: + self.redis = redis.Redis( + host=self.host, + port=self.port, + password=self.password, + db=self.db, + decode_responses=True, + ) + # Ping to verify connection + await self.redis.ping() + logger.info("Storage service connected to Redis") + except Exception as e: + logger.error(f"Failed to connect to Redis: {e}") + raise ConnectionError(f"Redis connection failed: {e}") + + async def disconnect(self) -> None: + """Close Redis connection.""" + if self.redis: + await self.redis.aclose() + self.redis = None + logger.info("Storage service disconnected from Redis") + + async def is_healthy(self) -> bool: + """Check if Redis is healthy.""" + if not self.redis: + return False + try: + await self.redis.ping() + return True + except Exception as e: + logger.warning(f"Redis health check failed: {e}") + return False + + async def get_status(self) -> Dict[str, Any]: + """Get storage service status.""" + try: + is_healthy = await self.is_healthy() + status = { + "type": "redis", + "healthy": is_healthy, + "host": self.host, + "port": self.port, + "db": self.db, + } + return status + except Exception as e: + logger.error(f"Failed to get storage status: {e}") + return {"type": "redis", "healthy": False, "error": str(e)} + + async def get_client_by_name_and_redirect_uris( + self, client_name: str, redirect_uris: List[str] + ) -> Optional[Dict[str, Any]]: + """Find an existing client by name and redirect URIs.""" + if not self.redis: + return None + try: + idx_key = _get_client_index_key(client_name, redirect_uris) + client_id = await self.redis.get(idx_key) + if client_id: + cid = client_id.decode("utf-8") if isinstance(client_id, bytes) else str(client_id) + return await self.get_client(cid) + return None + except Exception as e: + logger.error(f"Failed to get client by name and redirect URIs: {e}") + return None + + async def store_client(self, client_data: Dict[str, Any]) -> bool: + """Store a new OAuth client.""" + if not self.redis: + return False + try: + client_id = client_data["id"] + # Store client details + await self.redis.set(f"client:{client_id}", json.dumps(client_data)) + + # Store index for lookup by name and redirect URIs + idx_key = _get_client_index_key( + client_data["name"], client_data["redirect_uris"] + ) + await self.redis.set(idx_key, client_id) + + logger.info(f"Storing client: {client_id}") + return True + except Exception as e: + logger.error(f"Failed to store client: {e}") + return False + + async def get_client(self, client_id: str) -> Optional[Dict[str, Any]]: + """Get a client by ID.""" + if not self.redis: + return None + try: + val = await self.redis.get(f"client:{client_id}") + if val: + return json.loads(val) + return None + except Exception as e: + logger.error(f"Failed to get client: {e}") + return None + + async def store_authorization_code( + self, code: str, code_data: Dict[str, Any] + ) -> bool: + """Store an authorization code.""" + if not self.redis: + return False + try: + ttl = max(1, int(code_data["expires_at"] - time.time())) + await self.redis.setex(f"auth_code:{code}", ttl, json.dumps(code_data)) + return True + except Exception as e: + logger.error(f"Failed to store authorization code: {e}") + return False + + async def get_authorization_code(self, code: str) -> Optional[Dict[str, Any]]: + """Get authorization code data.""" + if not self.redis: + return None + try: + val = await self.redis.get(f"auth_code:{code}") + if val: + return json.loads(val) + return None + except Exception as e: + logger.error(f"Failed to get authorization code: {e}") + return None + + async def update_authorization_code_token( + self, code: str, snowflake_token: Dict[str, Any] + ) -> bool: + """Update authorization code with Snowflake token.""" + if not self.redis: + return False + try: + key = f"auth_code:{code}" + data_str = await self.redis.get(key) + if not data_str: + return False + data = json.loads(data_str) + data["snowflake_token"] = snowflake_token + # Get remaining TTL to preserve it + ttl = await self.redis.ttl(key) + if ttl > 0: + await self.redis.setex(key, ttl, json.dumps(data)) + else: + await self.redis.set(key, json.dumps(data)) + return True + except Exception as e: + logger.error(f"Failed to update authorization code: {e}") + return False + + async def delete_authorization_code(self, code: str) -> bool: + """Delete an authorization code.""" + if not self.redis: + return False + try: + result = await self.redis.delete(f"auth_code:{code}") + return result > 0 + except Exception as e: + logger.error(f"Failed to delete authorization code: {e}") + return False + + async def store_access_token(self, token: str, token_data: Dict[str, Any]) -> bool: + """Store an access token.""" + if not self.redis: + return False + try: + ttl = max(1, int(token_data["expires_at"] - time.time())) + await self.redis.setex(f"access_token:{token}", ttl, json.dumps(token_data)) + return True + except Exception as e: + logger.error(f"Failed to store access token: {e}") + return False + + async def get_access_token(self, token: str) -> Optional[Dict[str, Any]]: + """Get access token data.""" + if not self.redis: + return None + try: + val = await self.redis.get(f"access_token:{token}") + if val: + return json.loads(val) + return None + except Exception as e: + logger.error(f"Failed to get access token: {e}") + return None + + async def delete_access_token(self, token: str) -> bool: + """Delete an access token.""" + if not self.redis: + return False + try: + result = await self.redis.delete(f"access_token:{token}") + return result > 0 + except Exception as e: + logger.error(f"Failed to delete access token: {e}") + return False + + async def store_refresh_token(self, token: str, token_data: Dict[str, Any]) -> bool: + """Store a refresh token.""" + if not self.redis: + return False + try: + ttl = max(1, int(token_data["expires_at"] - time.time())) + await self.redis.setex( + f"refresh_token:{token}", ttl, json.dumps(token_data) + ) + return True + except Exception as e: + logger.error(f"Failed to store refresh token: {e}") + return False + + async def get_refresh_token(self, token: str) -> Optional[Dict[str, Any]]: + """Get refresh token data.""" + if not self.redis: + return None + try: + val = await self.redis.get(f"refresh_token:{token}") + if val: + return json.loads(val) + return None + except Exception as e: + logger.error(f"Failed to get refresh token: {e}") + return None + + async def delete_refresh_token(self, token: str) -> bool: + """Delete a refresh token.""" + if not self.redis: + return False + try: + result = await self.redis.delete(f"refresh_token:{token}") + return result > 0 + except Exception as e: + logger.error(f"Failed to delete refresh token: {e}") + return False diff --git a/template_mcp_server/src/storage/storage_service.py b/template_mcp_server/src/storage/storage_service.py index df32b4d..9811dd7 100644 --- a/template_mcp_server/src/storage/storage_service.py +++ b/template_mcp_server/src/storage/storage_service.py @@ -6,12 +6,13 @@ import asyncpg +from template_mcp_server.src.storage.base import BaseStorageService from template_mcp_server.utils.pylogger import get_python_logger logger = get_python_logger() -class StorageService: +class StorageService(BaseStorageService): """PostgreSQL storage service for persistent data storage. This service provides direct PostgreSQL storage functionality without diff --git a/tests/test_redis_service.py b/tests/test_redis_service.py new file mode 100644 index 0000000..7d4e387 --- /dev/null +++ b/tests/test_redis_service.py @@ -0,0 +1,143 @@ +import json +import time +from unittest.mock import AsyncMock, patch + +import pytest + +from template_mcp_server.src.storage.redis_service import ( + RedisStorageService, + _get_client_index_key, +) + + +class TestRedisStorageServiceInit: + """Test RedisStorageService initialization.""" + + def test_init_default_values(self): + service = RedisStorageService() + assert service.host == "localhost" + assert service.port == 6379 + assert service.password is None + assert service.db == 0 + assert service.redis is None + + def test_init_custom_values(self): + service = RedisStorageService(host="redis", port=6380, password="secret", db=1) + assert service.host == "redis" + assert service.port == 6380 + assert service.password == "secret" + assert service.db == 1 + + +class TestRedisStorageServiceConnection: + """Test connection management.""" + + @pytest.mark.asyncio + async def test_connect_success(self): + service = RedisStorageService() + mock_redis = AsyncMock() + + with patch( + "template_mcp_server.src.storage.redis_service.redis.Redis", + return_value=mock_redis, + ): + await service.connect() + mock_redis.ping.assert_called_once() + assert service.redis == mock_redis + + @pytest.mark.asyncio + async def test_connect_failure(self): + service = RedisStorageService() + mock_redis = AsyncMock() + mock_redis.ping.side_effect = Exception("Connection refused") + + with patch( + "template_mcp_server.src.storage.redis_service.redis.Redis", + return_value=mock_redis, + ): + with pytest.raises(ConnectionError): + await service.connect() + + @pytest.mark.asyncio + async def test_disconnect(self): + service = RedisStorageService() + mock_redis = AsyncMock() + service.redis = mock_redis + + await service.disconnect() + mock_redis.aclose.assert_called_once() + assert service.redis is None + + +class TestRedisStorageServiceMethods: + """Test storage methods.""" + + @pytest.fixture + def service(self): + s = RedisStorageService() + s.redis = AsyncMock() + return s + + @pytest.mark.asyncio + async def test_store_client(self, service): + client_data = { + "id": "123", + "name": "Test", + "redirect_uris": ["http://localhost"], + } + result = await service.store_client(client_data) + + assert result is True + # Check that it saves the client info and the index + assert service.redis.set.call_count == 2 + idx_key = _get_client_index_key("Test", ["http://localhost"]) + service.redis.set.assert_any_call("client:123", json.dumps(client_data)) + service.redis.set.assert_any_call(idx_key, "123") + + @pytest.mark.asyncio + async def test_get_client(self, service): + client_data = {"id": "123", "name": "Test"} + service.redis.get.return_value = json.dumps(client_data) + + result = await service.get_client("123") + assert result == client_data + service.redis.get.assert_called_once_with("client:123") + + @pytest.mark.asyncio + async def test_store_authorization_code(self, service): + code_data = {"expires_at": time.time() + 600, "client_id": "123"} + result = await service.store_authorization_code("code1", code_data) + + assert result is True + service.redis.setex.assert_called_once() + args, _ = service.redis.setex.call_args + assert args[0] == "auth_code:code1" + assert args[1] > 0 # TTL + assert args[2] == json.dumps(code_data) + + @pytest.mark.asyncio + async def test_update_authorization_code_token(self, service): + code_data = {"client_id": "123", "snowflake_token": None} + service.redis.get.return_value = json.dumps(code_data) + service.redis.ttl.return_value = 300 + + snowflake = {"access_token": "token1"} + result = await service.update_authorization_code_token("code1", snowflake) + + assert result is True + service.redis.setex.assert_called_once() + args, _ = service.redis.setex.call_args + assert args[0] == "auth_code:code1" + assert args[1] == 300 + saved_data = json.loads(args[2]) + assert saved_data["snowflake_token"] == snowflake + + @pytest.mark.asyncio + async def test_get_client_by_name_and_redirect_uris(self, service): + client_data = {"id": "123", "name": "Test", "redirect_uris": ["url"]} + service.redis.get.side_effect = ["123", json.dumps(client_data)] + + result = await service.get_client_by_name_and_redirect_uris("Test", ["url"]) + + assert result == client_data + assert service.redis.get.call_count == 2 diff --git a/tests/test_storage_init.py b/tests/test_storage_init.py index 6bb0cf4..48d8276 100644 --- a/tests/test_storage_init.py +++ b/tests/test_storage_init.py @@ -33,4 +33,4 @@ def test_module_docstring(self): from template_mcp_server.src import storage assert storage.__doc__ is not None - assert "PostgreSQL storage service" in storage.__doc__ + assert "Storage services" in storage.__doc__ From b49b5c8556cf5c5a52c49b01d89b4c5410a1b18f Mon Sep 17 00:00:00 2001 From: Yash Pawar Date: Tue, 16 Jun 2026 02:20:54 +0530 Subject: [PATCH 2/2] test: expand redis storage coverage to 82% and fix base.py module docstring --- template_mcp_server/src/storage/base.py | 2 + tests/test_redis_service.py | 252 ++++++++++++++++++++++++ 2 files changed, 254 insertions(+) diff --git a/template_mcp_server/src/storage/base.py b/template_mcp_server/src/storage/base.py index 2f0044c..674f65e 100644 --- a/template_mcp_server/src/storage/base.py +++ b/template_mcp_server/src/storage/base.py @@ -1,3 +1,5 @@ +"""Abstract base class defining the storage interface for the Template MCP Server.""" + from abc import ABC, abstractmethod from typing import Any, Dict, List, Optional diff --git a/tests/test_redis_service.py b/tests/test_redis_service.py index 7d4e387..a3bbd85 100644 --- a/tests/test_redis_service.py +++ b/tests/test_redis_service.py @@ -141,3 +141,255 @@ async def test_get_client_by_name_and_redirect_uris(self, service): assert result == client_data assert service.redis.get.call_count == 2 + + @pytest.mark.asyncio + async def test_get_client_by_name_not_found(self, service): + service.redis.get.return_value = None + result = await service.get_client_by_name_and_redirect_uris("Test", ["url"]) + assert result is None + + @pytest.mark.asyncio + async def test_get_client_not_found(self, service): + service.redis.get.return_value = None + result = await service.get_client("nonexistent") + assert result is None + + @pytest.mark.asyncio + async def test_store_client_no_redis(self): + service = RedisStorageService() + result = await service.store_client({"id": "1", "name": "x", "redirect_uris": []}) + assert result is False + + @pytest.mark.asyncio + async def test_get_client_no_redis(self): + service = RedisStorageService() + result = await service.get_client("id") + assert result is None + + @pytest.mark.asyncio + async def test_get_client_by_name_no_redis(self): + service = RedisStorageService() + result = await service.get_client_by_name_and_redirect_uris("x", []) + assert result is None + + @pytest.mark.asyncio + async def test_store_client_exception(self, service): + service.redis.set.side_effect = Exception("fail") + result = await service.store_client({"id": "1", "name": "x", "redirect_uris": []}) + assert result is False + + @pytest.mark.asyncio + async def test_get_client_exception(self, service): + service.redis.get.side_effect = Exception("fail") + result = await service.get_client("id") + assert result is None + + # --- Authorization code --- + @pytest.mark.asyncio + async def test_get_authorization_code(self, service): + data = {"client_id": "c1", "expires_at": time.time() + 600} + service.redis.get.return_value = json.dumps(data) + result = await service.get_authorization_code("code1") + assert result == data + + @pytest.mark.asyncio + async def test_get_authorization_code_not_found(self, service): + service.redis.get.return_value = None + result = await service.get_authorization_code("x") + assert result is None + + @pytest.mark.asyncio + async def test_get_authorization_code_no_redis(self): + service = RedisStorageService() + result = await service.get_authorization_code("code1") + assert result is None + + @pytest.mark.asyncio + async def test_store_authorization_code_no_redis(self): + service = RedisStorageService() + result = await service.store_authorization_code("c", {"expires_at": time.time() + 60}) + assert result is False + + @pytest.mark.asyncio + async def test_delete_authorization_code(self, service): + service.redis.delete.return_value = 1 + result = await service.delete_authorization_code("code1") + assert result is True + service.redis.delete.assert_called_once_with("auth_code:code1") + + @pytest.mark.asyncio + async def test_delete_authorization_code_not_found(self, service): + service.redis.delete.return_value = 0 + result = await service.delete_authorization_code("code1") + assert result is False + + @pytest.mark.asyncio + async def test_delete_authorization_code_no_redis(self): + service = RedisStorageService() + result = await service.delete_authorization_code("code1") + assert result is False + + @pytest.mark.asyncio + async def test_update_authorization_code_token_not_found(self, service): + service.redis.get.return_value = None + result = await service.update_authorization_code_token("code1", {}) + assert result is False + + @pytest.mark.asyncio + async def test_update_authorization_code_token_no_ttl(self, service): + data = {"client_id": "c1"} + service.redis.get.return_value = json.dumps(data) + service.redis.ttl.return_value = -1 + result = await service.update_authorization_code_token("code1", {"token": "t"}) + assert result is True + service.redis.set.assert_called_once() + + @pytest.mark.asyncio + async def test_update_authorization_code_no_redis(self): + service = RedisStorageService() + result = await service.update_authorization_code_token("code1", {}) + assert result is False + + # --- Access tokens --- + @pytest.mark.asyncio + async def test_store_access_token(self, service): + data = {"client_id": "c1", "expires_at": time.time() + 3600} + result = await service.store_access_token("tok1", data) + assert result is True + service.redis.setex.assert_called_once() + + @pytest.mark.asyncio + async def test_store_access_token_no_redis(self): + service = RedisStorageService() + result = await service.store_access_token("tok", {"expires_at": time.time() + 100}) + assert result is False + + @pytest.mark.asyncio + async def test_get_access_token(self, service): + data = {"client_id": "c1", "scope": "read"} + service.redis.get.return_value = json.dumps(data) + result = await service.get_access_token("tok1") + assert result == data + + @pytest.mark.asyncio + async def test_get_access_token_not_found(self, service): + service.redis.get.return_value = None + result = await service.get_access_token("tok1") + assert result is None + + @pytest.mark.asyncio + async def test_get_access_token_no_redis(self): + service = RedisStorageService() + result = await service.get_access_token("tok1") + assert result is None + + @pytest.mark.asyncio + async def test_delete_access_token(self, service): + service.redis.delete.return_value = 1 + result = await service.delete_access_token("tok1") + assert result is True + service.redis.delete.assert_called_once_with("access_token:tok1") + + @pytest.mark.asyncio + async def test_delete_access_token_not_found(self, service): + service.redis.delete.return_value = 0 + result = await service.delete_access_token("tok1") + assert result is False + + @pytest.mark.asyncio + async def test_delete_access_token_no_redis(self): + service = RedisStorageService() + result = await service.delete_access_token("tok1") + assert result is False + + # --- Refresh tokens --- + @pytest.mark.asyncio + async def test_store_refresh_token(self, service): + data = {"client_id": "c1", "expires_at": time.time() + 7200} + result = await service.store_refresh_token("ref1", data) + assert result is True + service.redis.setex.assert_called_once() + + @pytest.mark.asyncio + async def test_store_refresh_token_no_redis(self): + service = RedisStorageService() + result = await service.store_refresh_token("ref", {"expires_at": time.time() + 100}) + assert result is False + + @pytest.mark.asyncio + async def test_get_refresh_token(self, service): + data = {"client_id": "c1", "scope": "read"} + service.redis.get.return_value = json.dumps(data) + result = await service.get_refresh_token("ref1") + assert result == data + + @pytest.mark.asyncio + async def test_get_refresh_token_not_found(self, service): + service.redis.get.return_value = None + result = await service.get_refresh_token("ref1") + assert result is None + + @pytest.mark.asyncio + async def test_get_refresh_token_no_redis(self): + service = RedisStorageService() + result = await service.get_refresh_token("ref1") + assert result is None + + @pytest.mark.asyncio + async def test_delete_refresh_token(self, service): + service.redis.delete.return_value = 1 + result = await service.delete_refresh_token("ref1") + assert result is True + service.redis.delete.assert_called_once_with("refresh_token:ref1") + + @pytest.mark.asyncio + async def test_delete_refresh_token_not_found(self, service): + service.redis.delete.return_value = 0 + result = await service.delete_refresh_token("ref1") + assert result is False + + @pytest.mark.asyncio + async def test_delete_refresh_token_no_redis(self): + service = RedisStorageService() + result = await service.delete_refresh_token("ref1") + assert result is False + + # --- Health and Status --- + @pytest.mark.asyncio + async def test_is_healthy_true(self, service): + result = await service.is_healthy() + assert result is True + service.redis.ping.assert_called_once() + + @pytest.mark.asyncio + async def test_is_healthy_no_redis(self): + service = RedisStorageService() + result = await service.is_healthy() + assert result is False + + @pytest.mark.asyncio + async def test_is_healthy_exception(self, service): + service.redis.ping.side_effect = Exception("ping failed") + result = await service.is_healthy() + assert result is False + + @pytest.mark.asyncio + async def test_get_status_healthy(self, service): + result = await service.get_status() + assert result["type"] == "redis" + assert result["healthy"] is True + assert result["host"] == "localhost" + assert result["port"] == 6379 + + @pytest.mark.asyncio + async def test_get_status_unhealthy(self): + service = RedisStorageService() + result = await service.get_status() + assert result["type"] == "redis" + assert result["healthy"] is False + + @pytest.mark.asyncio + async def test_disconnect_no_redis(self): + service = RedisStorageService() + await service.disconnect() # Should not raise +