mirror of
https://github.com/jmagar/unraid-mcp.git
synced 2026-03-23 04:29:17 -07:00
feat(live): add unraid_live tool with 11 subscription snapshot actions
Creates unraid_mcp/tools/live.py with SNAPSHOT_ACTIONS (9 one-shot reads) and COLLECT_ACTIONS (2 streaming collectors), plus tests/test_live.py with 6 passing tests. Registers register_live_tool in server.py, bringing the total to 12 tools.
This commit is contained in:
95
tests/test_live.py
Normal file
95
tests/test_live.py
Normal file
@@ -0,0 +1,95 @@
|
||||
# tests/test_live.py
|
||||
"""Tests for unraid_live subscription snapshot tool."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastmcp import FastMCP
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mcp():
|
||||
return FastMCP("test")
|
||||
|
||||
|
||||
def _make_live_tool(mcp):
|
||||
from unraid_mcp.tools.live import register_live_tool
|
||||
|
||||
register_live_tool(mcp)
|
||||
local_provider = mcp.providers[0]
|
||||
tool = local_provider._components["tool:unraid_live@"]
|
||||
return tool.fn
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _mock_subscribe_once():
|
||||
with patch("unraid_mcp.tools.live.subscribe_once") as m:
|
||||
yield m
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _mock_subscribe_collect():
|
||||
with patch("unraid_mcp.tools.live.subscribe_collect") as m:
|
||||
yield m
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cpu_returns_snapshot(mcp, _mock_subscribe_once):
|
||||
_mock_subscribe_once.return_value = {"systemMetricsCpu": {"percentTotal": 23.5, "cpus": []}}
|
||||
tool_fn = _make_live_tool(mcp)
|
||||
result = await tool_fn(action="cpu")
|
||||
assert result["success"] is True
|
||||
assert result["data"]["systemMetricsCpu"]["percentTotal"] == 23.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_returns_snapshot(mcp, _mock_subscribe_once):
|
||||
_mock_subscribe_once.return_value = {
|
||||
"systemMetricsMemory": {"total": 32000000000, "used": 10000000000, "percentTotal": 31.2}
|
||||
}
|
||||
tool_fn = _make_live_tool(mcp)
|
||||
result = await tool_fn(action="memory")
|
||||
assert result["success"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_tail_requires_path(mcp, _mock_subscribe_collect):
|
||||
_mock_subscribe_collect.return_value = []
|
||||
tool_fn = _make_live_tool(mcp)
|
||||
from unraid_mcp.core.exceptions import ToolError
|
||||
|
||||
with pytest.raises(ToolError, match="path"):
|
||||
await tool_fn(action="log_tail")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_log_tail_with_path(mcp, _mock_subscribe_collect):
|
||||
_mock_subscribe_collect.return_value = [
|
||||
{"logFile": {"path": "/var/log/syslog", "content": "line1\nline2", "totalLines": 2}}
|
||||
]
|
||||
tool_fn = _make_live_tool(mcp)
|
||||
result = await tool_fn(action="log_tail", path="/var/log/syslog", collect_for=1.0)
|
||||
assert result["success"] is True
|
||||
assert result["event_count"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_notification_feed_collects_events(mcp, _mock_subscribe_collect):
|
||||
_mock_subscribe_collect.return_value = [
|
||||
{"notificationAdded": {"id": "1", "title": "Alert"}},
|
||||
{"notificationAdded": {"id": "2", "title": "Info"}},
|
||||
]
|
||||
tool_fn = _make_live_tool(mcp)
|
||||
result = await tool_fn(action="notification_feed", collect_for=2.0)
|
||||
assert result["event_count"] == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_action_raises(mcp):
|
||||
from unraid_mcp.core.exceptions import ToolError
|
||||
|
||||
tool_fn = _make_live_tool(mcp)
|
||||
with pytest.raises(ToolError, match="Invalid action"):
|
||||
await tool_fn(action="nonexistent") # type: ignore[arg-type]
|
||||
@@ -24,6 +24,7 @@ from .tools.docker import register_docker_tool
|
||||
from .tools.health import register_health_tool
|
||||
from .tools.info import register_info_tool
|
||||
from .tools.keys import register_keys_tool
|
||||
from .tools.live import register_live_tool
|
||||
from .tools.notifications import register_notifications_tool
|
||||
from .tools.rclone import register_rclone_tool
|
||||
from .tools.settings import register_settings_tool
|
||||
@@ -64,6 +65,7 @@ def register_all_modules() -> None:
|
||||
register_keys_tool,
|
||||
register_health_tool,
|
||||
register_settings_tool,
|
||||
register_live_tool,
|
||||
]
|
||||
for registrar in registrars:
|
||||
registrar(mcp)
|
||||
|
||||
167
unraid_mcp/tools/live.py
Normal file
167
unraid_mcp/tools/live.py
Normal file
@@ -0,0 +1,167 @@
|
||||
"""Real-time subscription snapshot tool.
|
||||
|
||||
Provides the `unraid_live` tool with 11 actions — one per GraphQL
|
||||
subscription. Each action opens a transient WebSocket, receives one event
|
||||
(or collects events for `collect_for` seconds), then closes.
|
||||
|
||||
Use `subscribe_once` actions for current-state reads (cpu, memory, array_state).
|
||||
Use `subscribe_collect` actions for event streams (notification_feed, log_tail).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal, get_args
|
||||
|
||||
from fastmcp import FastMCP
|
||||
|
||||
from ..config.logging import logger
|
||||
from ..core.exceptions import ToolError, tool_error_handler
|
||||
from ..subscriptions.snapshot import subscribe_collect, subscribe_once
|
||||
|
||||
|
||||
SNAPSHOT_ACTIONS = {
|
||||
"cpu": """
|
||||
subscription { systemMetricsCpu { id percentTotal cpus { percentTotal percentUser percentSystem percentIdle } } }
|
||||
""",
|
||||
"memory": """
|
||||
subscription { systemMetricsMemory { id total used free available active buffcache percentTotal swapTotal swapUsed swapFree percentSwapTotal } }
|
||||
""",
|
||||
"cpu_telemetry": """
|
||||
subscription { systemMetricsCpuTelemetry { id totalPower power temp } }
|
||||
""",
|
||||
"array_state": """
|
||||
subscription { arraySubscription { id state capacity { kilobytes { free used total } } parityCheckStatus { status progress speed errors } } }
|
||||
""",
|
||||
"parity_progress": """
|
||||
subscription { parityHistorySubscription { date status progress speed errors correcting paused running } }
|
||||
""",
|
||||
"ups_status": """
|
||||
subscription { upsUpdates { id name model status battery { chargeLevel estimatedRuntime health } power { inputVoltage outputVoltage loadPercentage } } }
|
||||
""",
|
||||
"notifications_overview": """
|
||||
subscription { notificationsOverview { unread { info warning alert total } archive { info warning alert total } } }
|
||||
""",
|
||||
"owner": """
|
||||
subscription { ownerSubscription { username url avatar } }
|
||||
""",
|
||||
"server_status": """
|
||||
subscription { serversSubscription { id name status guid wanip lanip localurl remoteurl } }
|
||||
""",
|
||||
}
|
||||
|
||||
COLLECT_ACTIONS = {
|
||||
"notification_feed": """
|
||||
subscription { notificationAdded { id title subject description importance type timestamp } }
|
||||
""",
|
||||
"log_tail": """
|
||||
subscription LogTail($path: String!) { logFile(path: $path) { path content totalLines startLine } }
|
||||
""",
|
||||
}
|
||||
|
||||
ALL_LIVE_ACTIONS = set(SNAPSHOT_ACTIONS) | set(COLLECT_ACTIONS)
|
||||
|
||||
LIVE_ACTIONS = Literal[
|
||||
"array_state",
|
||||
"cpu",
|
||||
"cpu_telemetry",
|
||||
"log_tail",
|
||||
"memory",
|
||||
"notification_feed",
|
||||
"notifications_overview",
|
||||
"owner",
|
||||
"parity_progress",
|
||||
"server_status",
|
||||
"ups_status",
|
||||
]
|
||||
|
||||
if set(get_args(LIVE_ACTIONS)) != ALL_LIVE_ACTIONS:
|
||||
_missing = ALL_LIVE_ACTIONS - set(get_args(LIVE_ACTIONS))
|
||||
_extra = set(get_args(LIVE_ACTIONS)) - ALL_LIVE_ACTIONS
|
||||
raise RuntimeError(
|
||||
f"LIVE_ACTIONS and ALL_LIVE_ACTIONS are out of sync. "
|
||||
f"Missing: {_missing or 'none'}. Extra: {_extra or 'none'}"
|
||||
)
|
||||
|
||||
|
||||
def register_live_tool(mcp: FastMCP) -> None:
|
||||
"""Register the unraid_live tool with the FastMCP instance."""
|
||||
|
||||
@mcp.tool()
|
||||
async def unraid_live(
|
||||
action: LIVE_ACTIONS,
|
||||
path: str | None = None,
|
||||
collect_for: float = 5.0,
|
||||
timeout: float = 10.0,
|
||||
) -> dict[str, Any]:
|
||||
"""Get real-time data from Unraid via WebSocket subscriptions.
|
||||
|
||||
Each action opens a transient WebSocket, receives data, then closes.
|
||||
|
||||
Snapshot actions (return current state):
|
||||
cpu - Real-time CPU utilization (all cores)
|
||||
memory - Real-time memory and swap utilization
|
||||
cpu_telemetry - CPU power draw and temperature per package
|
||||
array_state - Live array state and parity status
|
||||
parity_progress - Live parity check progress
|
||||
ups_status - Real-time UPS battery and power state
|
||||
notifications_overview - Live notification counts by severity
|
||||
owner - Live owner info
|
||||
server_status - Live server connection state
|
||||
|
||||
Collection actions (collect events for `collect_for` seconds):
|
||||
notification_feed - Collect new notification events (default: 5s window)
|
||||
log_tail - Tail a log file (requires path; default: 5s window)
|
||||
|
||||
Parameters:
|
||||
path - Log file path for log_tail action (required)
|
||||
collect_for - Seconds to collect events for collect actions (default: 5.0)
|
||||
timeout - WebSocket connection/handshake timeout in seconds (default: 10.0)
|
||||
"""
|
||||
if action not in ALL_LIVE_ACTIONS:
|
||||
raise ToolError(
|
||||
f"Invalid action '{action}'. Must be one of: {sorted(ALL_LIVE_ACTIONS)}"
|
||||
)
|
||||
|
||||
with tool_error_handler("live", action, logger):
|
||||
logger.info(f"Executing unraid_live action={action} timeout={timeout}")
|
||||
|
||||
if action in SNAPSHOT_ACTIONS:
|
||||
data = await subscribe_once(SNAPSHOT_ACTIONS[action], timeout=timeout)
|
||||
return {"success": True, "action": action, "data": data}
|
||||
|
||||
# Collect actions
|
||||
if action == "log_tail":
|
||||
if not path:
|
||||
raise ToolError("path is required for 'log_tail' action")
|
||||
events = await subscribe_collect(
|
||||
COLLECT_ACTIONS["log_tail"],
|
||||
variables={"path": path},
|
||||
collect_for=collect_for,
|
||||
timeout=timeout,
|
||||
)
|
||||
return {
|
||||
"success": True,
|
||||
"action": action,
|
||||
"path": path,
|
||||
"collect_for": collect_for,
|
||||
"event_count": len(events),
|
||||
"events": events,
|
||||
}
|
||||
|
||||
if action == "notification_feed":
|
||||
events = await subscribe_collect(
|
||||
COLLECT_ACTIONS["notification_feed"],
|
||||
collect_for=collect_for,
|
||||
timeout=timeout,
|
||||
)
|
||||
return {
|
||||
"success": True,
|
||||
"action": action,
|
||||
"collect_for": collect_for,
|
||||
"event_count": len(events),
|
||||
"events": events,
|
||||
}
|
||||
|
||||
raise ToolError(f"Unhandled action '{action}' — this is a bug")
|
||||
|
||||
logger.info("Live tool registered successfully")
|
||||
Reference in New Issue
Block a user