test: Add load testing framework

tests/load/workloads.py

@@ -0,0 +1,282 @@
+"""
+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()