diff --git a/tests/test_customization.py b/tests/test_customization.py new file mode 100644 index 0000000..55eedfc --- /dev/null +++ b/tests/test_customization.py @@ -0,0 +1,64 @@ +# tests/test_customization.py +"""Tests for unraid_customization tool.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, patch + +import pytest +from conftest import make_tool_fn + + +@pytest.fixture +def _mock_graphql(): + with patch("unraid_mcp.tools.customization.make_graphql_request", new_callable=AsyncMock) as m: + yield m + + +def _make_tool(): + return make_tool_fn( + "unraid_mcp.tools.customization", + "register_customization_tool", + "unraid_customization", + ) + + +@pytest.mark.asyncio +async def test_theme_returns_customization(_mock_graphql): + _mock_graphql.return_value = { + "customization": {"theme": {"name": "azure"}, "partnerInfo": None, "activationCode": None} + } + result = await _make_tool()(action="theme") + assert result["success"] is True + + +@pytest.mark.asyncio +async def test_public_theme(_mock_graphql): + _mock_graphql.return_value = {"publicTheme": {"name": "black"}} + result = await _make_tool()(action="public_theme") + assert result["success"] is True + + +@pytest.mark.asyncio +async def test_is_initial_setup(_mock_graphql): + _mock_graphql.return_value = {"isInitialSetup": False} + result = await _make_tool()(action="is_initial_setup") + assert result["success"] is True + assert result["data"]["isInitialSetup"] is False + + +@pytest.mark.asyncio +async def test_set_theme_requires_theme(_mock_graphql): + from unraid_mcp.core.exceptions import ToolError + + with pytest.raises(ToolError, match="theme_name"): + await _make_tool()(action="set_theme") + + +@pytest.mark.asyncio +async def test_set_theme_success(_mock_graphql): + _mock_graphql.return_value = { + "customization": {"setTheme": {"name": "azure", "showBannerImage": True}} + } + result = await _make_tool()(action="set_theme", theme_name="azure") + assert result["success"] is True diff --git a/unraid_mcp/server.py b/unraid_mcp/server.py index 13e5bc7..a8ef698 100644 --- a/unraid_mcp/server.py +++ b/unraid_mcp/server.py @@ -20,6 +20,7 @@ from .config.settings import ( from .subscriptions.diagnostics import register_diagnostic_tools from .subscriptions.resources import register_subscription_resources from .tools.array import register_array_tool +from .tools.customization import register_customization_tool from .tools.docker import register_docker_tool from .tools.health import register_health_tool from .tools.info import register_info_tool @@ -66,6 +67,7 @@ def register_all_modules() -> None: register_health_tool, register_settings_tool, register_live_tool, + register_customization_tool, ] for registrar in registrars: registrar(mcp) diff --git a/unraid_mcp/tools/customization.py b/unraid_mcp/tools/customization.py new file mode 100644 index 0000000..08bbb47 --- /dev/null +++ b/unraid_mcp/tools/customization.py @@ -0,0 +1,119 @@ +"""UI customization and system state queries. + +Provides the `unraid_customization` tool with 5 actions covering +theme/customization data, public UI config, initial setup state, and +theme mutation. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal, get_args + + +if TYPE_CHECKING: + from fastmcp import FastMCP + +from ..config.logging import logger +from ..core.client import make_graphql_request +from ..core.exceptions import ToolError, tool_error_handler + + +QUERIES: dict[str, str] = { + "theme": """ + query GetCustomization { + customization { + theme { name showBannerImage showBannerGradient showHeaderDescription + headerBackgroundColor headerPrimaryTextColor headerSecondaryTextColor } + partnerInfo { partnerName hasPartnerLogo partnerUrl partnerLogoUrl } + activationCode { code partnerName serverName sysModel comment header theme } + } + } + """, + "public_theme": """ + query GetPublicTheme { + publicTheme { name showBannerImage showBannerGradient showHeaderDescription + headerBackgroundColor headerPrimaryTextColor headerSecondaryTextColor } + publicPartnerInfo { partnerName hasPartnerLogo partnerUrl partnerLogoUrl } + } + """, + "is_initial_setup": """ + query IsInitialSetup { + isInitialSetup + } + """, + "sso_enabled": """ + query IsSSOEnabled { + isSSOEnabled + } + """, +} + +MUTATIONS: dict[str, str] = { + "set_theme": """ + mutation SetTheme($theme: ThemeName!) { + customization { setTheme(theme: $theme) { + name showBannerImage showBannerGradient showHeaderDescription + }} + } + """, +} + +ALL_ACTIONS = set(QUERIES) | set(MUTATIONS) + +CUSTOMIZATION_ACTIONS = Literal[ + "is_initial_setup", + "public_theme", + "set_theme", + "sso_enabled", + "theme", +] + +if set(get_args(CUSTOMIZATION_ACTIONS)) != ALL_ACTIONS: + _missing = ALL_ACTIONS - set(get_args(CUSTOMIZATION_ACTIONS)) + _extra = set(get_args(CUSTOMIZATION_ACTIONS)) - ALL_ACTIONS + raise RuntimeError( + f"CUSTOMIZATION_ACTIONS and ALL_ACTIONS are out of sync. " + f"Missing: {_missing or 'none'}. Extra: {_extra or 'none'}" + ) + + +def register_customization_tool(mcp: FastMCP) -> None: + """Register the unraid_customization tool with the FastMCP instance.""" + + @mcp.tool() + async def unraid_customization( + action: CUSTOMIZATION_ACTIONS, + theme_name: str | None = None, + ) -> dict[str, Any]: + """Manage Unraid UI customization and system state. + + Actions: + theme - Get full customization (theme, partner info, activation code) + public_theme - Get public theme and partner info (no auth required) + is_initial_setup - Check if server is in initial setup mode + sso_enabled - Check if SSO is enabled + set_theme - Change the UI theme (requires theme_name: azure/black/gray/white) + """ + if action not in ALL_ACTIONS: + raise ToolError(f"Invalid action '{action}'. Must be one of: {sorted(ALL_ACTIONS)}") + + if action == "set_theme" and not theme_name: + raise ToolError( + "theme_name is required for 'set_theme' action " + "(valid values: azure, black, gray, white)" + ) + + with tool_error_handler("customization", action, logger): + logger.info(f"Executing unraid_customization action={action}") + + if action in QUERIES: + data = await make_graphql_request(QUERIES[action]) + return {"success": True, "action": action, "data": data} + + if action == "set_theme": + data = await make_graphql_request(MUTATIONS[action], {"theme": theme_name}) + return {"success": True, "action": action, "data": data} + + raise ToolError(f"Unhandled action '{action}' — this is a bug") + + logger.info("Customization tool registered successfully")