Merge pull request #4 from promptmesh/feat/testing

Setup pytest

Author: SecretiveShell

.gitignore

@@ -10,4 +10,7 @@ wheels/
 .venv
 
 # repomix
-repomix-output.xml
+repomix-output.xml
+
+# pytest
+.coverage

.vscode/settings.json

@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "tests"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}

pyproject.toml

@@ -25,4 +25,12 @@ build-backend = "hatchling.build"
 [dependency-groups]
 dev = [
     "mypy>=1.15.0",
+    "pytest>=8.3.5",
+    "pytest-asyncio>=0.26.0",
+    "pytest-cov>=6.1.0",
+    "pytest-rerunfailures>=15.0",
 ]
+
+[tool.pytest.ini_options]
+asyncio_mode = "strict"
+asyncio_default_fixture_loop_scope = "function"

src/easymcp/client/iobuffers.py

@@ -1,4 +1,4 @@
-from asyncio import Queue, Task
+from asyncio import Queue, create_task
 
 from loguru import logger
 import pydantic
@@ -24,7 +24,7 @@ async def _reader():
 
             queue.put_nowait(parsed)
 
-    task = Task(_reader())
+    task = create_task(_reader())
     return task
 
 
@@ -36,5 +36,5 @@ async def _writer():
             data = await queue.get()
             await transport.send(data.model_dump_json())
 
-    task = Task(_writer())
+    task = create_task(_writer())
     return task

src/easymcp/client/sessions/mcp/session.py

@@ -1,4 +1,5 @@
-from asyncio import Queue, Task
+from asyncio import Queue, Task, create_task
+import asyncio
 from inspect import iscoroutinefunction
 import json
 from typing import Awaitable, Callable
@@ -23,6 +24,8 @@ class MCPClientSession(BaseSessionProtocol):
     reader_task: Task[None]
     writer_task: Task[None]
 
+    _start_reading_messages_task: Task[None]
+
     request_map: RequestMap
 
     roots_callback: Callable[[types.ListRootsRequest], Awaitable[types.ListRootsResult]] | None = None
@@ -136,7 +139,7 @@ async def __start_reading_messages():
                 else:
                     logger.error(f"Unknown message type: {message.root}")
 
-        Task(__start_reading_messages())
+        self._start_reading_messages_task = create_task(__start_reading_messages())
 
     async def start(self) -> types.InitializeResult:
         """start the client session"""
@@ -188,7 +191,26 @@ async def start(self) -> types.InitializeResult:
 
     async def stop(self):
         """stop the client session"""
+        self.reader_task.cancel()
+        self.writer_task.cancel()
+        try:
+            await self.reader_task
+        except asyncio.CancelledError:
+            pass
+
+        try:
+            self.writer_task
+        except asyncio.CancelledError:
+            pass
+
+        self._start_reading_messages_task.cancel()
+        try:
+            await self._start_reading_messages_task
+        except asyncio.CancelledError:
+            pass
+
         await self.transport.stop()
+        await asyncio.sleep(0)
 
     async def list_tools(self, force: bool = False):
         """list available tools"""

src/easymcp/client/transports/docker.py

@@ -1,4 +1,5 @@
-from typing import Literal
+import asyncio
+from typing import Any, Literal
 
 import anyio
 import anyio.abc
@@ -7,6 +8,8 @@
 from loguru import logger
 from pydantic import BaseModel, Field
 
+from easymcp.client.transports.generic import TransportProtocol
+
 
 class DockerServerParameters(BaseModel):
     """Configuration for Docker transport."""
@@ -22,14 +25,14 @@ class DockerServerParameters(BaseModel):
     )
 
 
-class DockerTransport:
+class DockerTransport(TransportProtocol):
     state: Literal["constructed", "initialized", "started", "stopped"] = "constructed"
 
     def __init__(self, config: DockerServerParameters) -> None:
         self.config = config
         self.docker: Docker | None = None
         self.container: containers.DockerContainer | None = None
-        self.attach_result: any = None
+        self.attach_result: Any = None
 
         self._reader_send: MemoryObjectSendStream[str]
         self._reader_recv: MemoryObjectReceiveStream[str]
@@ -116,6 +119,7 @@ async def write_stdin():
         self._task_group.start_soon(write_stdin)
 
         self.state = "started"
+        await asyncio.sleep(1)
 
     async def stop(self) -> None:
         logger.debug("Stopping DockerTransport...")

src/easymcp/client/transports/stdio.py

@@ -3,6 +3,7 @@
 import os
 import shutil
 import sys
+import gc
 from typing import Literal
 
 from loguru import logger
@@ -58,6 +59,7 @@ class StdioTransport(TransportProtocol):
     arguments: StdioServerParameters
     subprocess: asyncio.subprocess.Process | None
     read_buffer: ReadBuffer
+    stderr_task: asyncio.Task | None
 
     def __init__(self, arguments: StdioServerParameters):
         self.state = "constructed"
@@ -125,18 +127,52 @@ async def read_stderr(self) -> None:
             print(line.decode(), file=sys.stderr, end="")
 
     async def stop(self) -> None:
-        try:
-            if self.subprocess:
-                self.subprocess.terminate()
+        if self.stderr_task:
+            self.stderr_task.cancel()
+            try:
+                await self.stderr_task
+            except asyncio.CancelledError:
+                pass
+            self.stderr_task = None
+
+        if self.subprocess:
+            if self.subprocess.returncode is None:
                 try:
+                    self.subprocess.terminate()
                     await asyncio.wait_for(self.subprocess.wait(), timeout=5)
                 except asyncio.TimeoutError:
                     self.subprocess.kill()
                     await self.subprocess.wait()
 
-            logger.info("Transport stopped successfully.")
-        except Exception as e:
-            logger.error(f"Error stopping transport: {e}")
-        finally:
+            if self.subprocess.stdin and not self.subprocess.stdin.is_closing():
+                logger.debug("Closing stdin")
+                self.subprocess.stdin.close()
+                await self.subprocess.stdin.wait_closed()
+            self.subprocess.stdin = None
+
+            if self.subprocess.stdout:
+                logger.debug("Closing stdout")
+                self.subprocess.stdout.feed_eof()
+                try:
+                    await self.subprocess.stdout.read()
+                except Exception:
+                    pass
+                self.subprocess.stdout = None
+
+            if self.subprocess.stderr:
+                logger.debug("Closing stderr")
+                self.subprocess.stderr.feed_eof()
+                try:
+                    await self.subprocess.stderr.read()
+                except Exception:
+                    pass
+                self.subprocess.stderr = None
+
+            await asyncio.sleep(0.1)
+            gc.collect()
+
             self.subprocess = None
-            self.state = "stopped"
+
+
+        self.state = "stopped"
+        logger.info("Transport stopped successfully.")