feat(customization): add unraid_customization tool with theme, public_theme, is_initial_setup, sso_enabled, set_theme

This commit is contained in:
Jacob Magar
2026-03-15 19:19:06 -04:00
parent 76391b4d2b
commit d26467a4d0
3 changed files with 185 additions and 0 deletions

View File

@@ -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

View File

@@ -20,6 +20,7 @@ from .config.settings import (
from .subscriptions.diagnostics import register_diagnostic_tools from .subscriptions.diagnostics import register_diagnostic_tools
from .subscriptions.resources import register_subscription_resources from .subscriptions.resources import register_subscription_resources
from .tools.array import register_array_tool from .tools.array import register_array_tool
from .tools.customization import register_customization_tool
from .tools.docker import register_docker_tool from .tools.docker import register_docker_tool
from .tools.health import register_health_tool from .tools.health import register_health_tool
from .tools.info import register_info_tool from .tools.info import register_info_tool
@@ -66,6 +67,7 @@ def register_all_modules() -> None:
register_health_tool, register_health_tool,
register_settings_tool, register_settings_tool,
register_live_tool, register_live_tool,
register_customization_tool,
] ]
for registrar in registrars: for registrar in registrars:
registrar(mcp) registrar(mcp)

View File

@@ -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")