"""
Workload definitions for load testing the MCP server.
Defines realistic operation mixes and individual operation functions.
"""
import logging
import random
import time
import uuid
from mcp import ClientSession
logger = logging.getLogger(__name__)
class OperationResult:
"""Result of a single operation execution."""
def __init__(
self,
operation: str,
success: bool,
duration: float,
error: str | None = None,
):
self.operation = operation
self.success = success
self.duration = duration
self.error = error
self.timestamp = time.time()
class WorkloadOperations:
"""Collection of MCP operations for load testing."""
def __init__(self, session: ClientSession):
self.session = session
self._created_notes: list[int] = []
self._created_boards: list[int] = []
async def get_capabilities(self) -> OperationResult:
"""Fetch server capabilities (lightweight operation)."""
start = time.time()
try:
await self.session.read_resource("nc://capabilities")
duration = time.time() - start
return OperationResult("get_capabilities", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("get_capabilities", False, duration, str(e))
async def list_notes(self) -> OperationResult:
"""List all notes (read operation)."""
start = time.time()
try:
await self.session.call_tool("nc_notes_search_notes", {"query": ""})
duration = time.time() - start
return OperationResult("list_notes", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("list_notes", False, duration, str(e))
async def search_notes(self, query: str = "test") -> OperationResult:
"""Search notes by query (read operation with filtering)."""
start = time.time()
try:
await self.session.call_tool("nc_notes_search_notes", {"query": query})
duration = time.time() - start
return OperationResult("search_notes", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("search_notes", False, duration, str(e))
async def create_note(self) -> OperationResult:
"""Create a new note (write operation)."""
start = time.time()
unique_id = uuid.uuid4().hex[:8]
try:
result = await self.session.call_tool(
"nc_notes_create_note",
{
"title": f"Load Test Note {unique_id}",
"content": f"Content for load test note {unique_id}",
"category": "LoadTesting",
},
)
duration = time.time() - start
# Track created note ID for cleanup
if result and len(result.content) > 0:
content = result.content[0]
if hasattr(content, "text"):
import json
note_data = json.loads(content.text)
note_id = note_data.get("id")
if note_id:
self._created_notes.append(note_id)
return OperationResult("create_note", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("create_note", False, duration, str(e))
async def get_note(self, note_id: int) -> OperationResult:
"""Get a specific note by ID (read operation)."""
start = time.time()
try:
await self.session.call_tool("nc_notes_get_note", {"note_id": note_id})
duration = time.time() - start
return OperationResult("get_note", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("get_note", False, duration, str(e))
async def update_note(self, note_id: int, etag: str) -> OperationResult:
"""Update an existing note (write operation)."""
start = time.time()
try:
await self.session.call_tool(
"nc_notes_update_note",
{
"note_id": note_id,
"etag": etag,
"title": f"Updated Note {note_id}",
"content": f"Updated content at {time.time()}",
"category": "LoadTesting",
},
)
duration = time.time() - start
return OperationResult("update_note", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("update_note", False, duration, str(e))
async def delete_note(self, note_id: int) -> OperationResult:
"""Delete a note (write operation)."""
start = time.time()
try:
await self.session.call_tool("nc_notes_delete_note", {"note_id": note_id})
duration = time.time() - start
# Remove from tracking
if note_id in self._created_notes:
self._created_notes.remove(note_id)
return OperationResult("delete_note", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("delete_note", False, duration, str(e))
async def list_webdav_files(self, path: str = "/") -> OperationResult:
"""List files via WebDAV (read operation)."""
start = time.time()
try:
await self.session.call_tool("nc_webdav_list", {"path": path})
duration = time.time() - start
return OperationResult("list_webdav_files", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("list_webdav_files", False, duration, str(e))
async def list_calendars(self) -> OperationResult:
"""List calendars (read operation)."""
start = time.time()
try:
await self.session.call_tool("nc_calendar_list_calendars", {})
duration = time.time() - start
return OperationResult("list_calendars", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("list_calendars", False, duration, str(e))
async def list_deck_boards(self) -> OperationResult:
"""List deck boards (read operation)."""
start = time.time()
try:
await self.session.call_tool("nc_deck_list_boards", {})
duration = time.time() - start
return OperationResult("list_deck_boards", True, duration)
except Exception as e:
duration = time.time() - start
return OperationResult("list_deck_boards", False, duration, str(e))
async def cleanup(self):
"""Clean up any resources created during testing."""
logger.info(f"Cleaning up {len(self._created_notes)} test notes...")
for note_id in self._created_notes[:]:
try:
await self.delete_note(note_id)
except Exception as e:
logger.warning(f"Failed to delete note {note_id}: {e}")
class MixedWorkload:
"""
Realistic mixed workload simulating typical MCP server usage.
Operation distribution:
- 40% Notes read (list/get/search)
- 20% Notes write (create/update/delete)
- 15% Notes search
- 10% WebDAV operations
- 10% Calendar operations
- 5% Other (capabilities, deck)
"""
def __init__(self, operations: WorkloadOperations):
self.ops = operations
# Pre-create some notes for read/update operations
self._warmup_note_ids: list[tuple[int, str]] = []
async def warmup(self, count: int = 10):
"""Create initial notes for read/update operations."""
logger.info(f"Warming up with {count} test notes...")
for _ in range(count):
result = await self.ops.create_note()
if result.success and self.ops._created_notes:
note_id = self.ops._created_notes[-1]
# Get the note to fetch its etag
try:
get_result = await self.ops.session.call_tool(
"nc_notes_get_note", {"note_id": note_id}
)
if get_result and len(get_result.content) > 0:
import json
note_data = json.loads(get_result.content[0].text)
etag = note_data.get("etag", "")
self._warmup_note_ids.append((note_id, etag))
except Exception as e:
logger.warning(f"Failed to get etag for note {note_id}: {e}")
async def run_operation(self) -> OperationResult:
"""Execute one random operation based on the workload distribution."""
rand = random.random()
# 40% reads (list/get/search)
if rand < 0.40:
op_rand = random.random()
if op_rand < 0.5:
return await self.ops.list_notes()
elif op_rand < 0.8 and self._warmup_note_ids:
note_id, _ = random.choice(self._warmup_note_ids)
return await self.ops.get_note(note_id)
else:
return await self.ops.search_notes()
# 20% writes (create/update/delete)
elif rand < 0.60:
op_rand = random.random()
if op_rand < 0.5:
return await self.ops.create_note()
elif op_rand < 0.8 and self._warmup_note_ids:
note_id, etag = random.choice(self._warmup_note_ids)
return await self.ops.update_note(note_id, etag)
elif self.ops._created_notes and len(self.ops._created_notes) > 5:
# Only delete if we have enough notes
note_id = random.choice(self.ops._created_notes)
return await self.ops.delete_note(note_id)
else:
return await self.ops.create_note()
# 15% search
elif rand < 0.75:
queries = ["test", "load", "note", "content", ""]
return await self.ops.search_notes(random.choice(queries))
# 10% WebDAV
elif rand < 0.85:
return await self.ops.list_webdav_files()
# 10% Calendar
elif rand < 0.95:
return await self.ops.list_calendars()
# 5% Other
else:
op_rand = random.random()
if op_rand < 0.5:
return await self.ops.get_capabilities()
else:
return await self.ops.list_deck_boards()