Skip to content

Commit 0ad4f8b

Browse files
committed
fix: add comprehensive type annotations to test fixture
- Add type aliases for stream types to improve clarity - Add return type annotation to fixture - Add type hints to internal _create function - Fixes all pyright type checking errors
1 parent d7b04ee commit 0ad4f8b

File tree

1 file changed

+13
-3
lines changed

1 file changed

+13
-3
lines changed

tests/shared/test_proxy.py

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,29 @@
11
"""Tests for the MCP proxy pattern."""
22

3+
from collections.abc import Callable
4+
35
import anyio
46
import pytest
7+
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
58

69
from mcp.shared.message import SessionMessage
710
from mcp.shared.proxy import mcp_proxy
811
from mcp.types import JSONRPCMessage, JSONRPCRequest
912

13+
# Type aliases for clarity
14+
ReadStream = MemoryObjectReceiveStream[SessionMessage | Exception]
15+
WriteStream = MemoryObjectSendStream[SessionMessage]
16+
StreamPair = tuple[ReadStream, WriteStream]
17+
WriterReaderPair = tuple[MemoryObjectSendStream[SessionMessage | Exception], MemoryObjectReceiveStream[SessionMessage]]
18+
StreamsFixtureReturn = tuple[StreamPair, StreamPair, WriterReaderPair, WriterReaderPair]
19+
1020

1121
@pytest.fixture
12-
async def create_streams():
22+
async def create_streams() -> Callable[[], StreamsFixtureReturn]:
1323
"""Helper fixture to create memory streams for testing with proper cleanup."""
14-
streams_to_cleanup = []
24+
streams_to_cleanup: list[MemoryObjectSendStream[SessionMessage | Exception] | MemoryObjectReceiveStream[SessionMessage | Exception] | MemoryObjectSendStream[SessionMessage] | MemoryObjectReceiveStream[SessionMessage]] = []
1525

16-
def _create():
26+
def _create() -> StreamsFixtureReturn:
1727
client_read_writer, client_read = anyio.create_memory_object_stream[SessionMessage | Exception](10)
1828
client_write, client_write_reader = anyio.create_memory_object_stream[SessionMessage](10)
1929

0 commit comments

Comments
 (0)