make int tests better

Author: akash329d

src/mcp/client/sse.py

@@ -98,9 +98,7 @@ async def sse_reader(
                                             await read_stream_writer.send(exc)
                                             continue
 
-                                        session_message = SessionMessage(
-                                            message=message
-                                        )
+                                        session_message = SessionMessage(message=message)
                                         await read_stream_writer.send(session_message)
                                     case _:
                                         logger.warning(
@@ -150,3 +148,6 @@ async def post_writer(endpoint_url: str):
         finally:
             await read_stream_writer.aclose()
             await write_stream.aclose()
+            await read_stream.aclose()
+            await write_stream_reader.aclose()
+            

src/mcp/server/message_queue/base.py

@@ -28,7 +28,7 @@ async def publish_message(
 
         Args:
             session_id: The UUID of the session this message is for
-            message: The message to publish (JSONRPCMessage or str for invalid JSON)
+            message: The message to publish (SessionMessage or str for invalid JSON)
 
         Returns:
             bool: True if message was published, False if session not found

src/mcp/server/sse.py

@@ -152,6 +152,8 @@ async def response_wrapper(scope: Scope, receive: Receive, send: Send):
                 )(scope, receive, send)
                 await read_stream_writer.aclose()
                 await write_stream_reader.aclose()
+                await sse_stream_writer.aclose()
+                await sse_stream_reader.aclose()
                 logging.debug(f"Client session disconnected {session_id}")
 
             logger.debug("Starting SSE response task")

tests/server/message_dispatch/test_redis.py

@@ -10,33 +10,6 @@
 from mcp.shared.message import SessionMessage
 
 
-@pytest.mark.anyio
-async def test_session_exists(message_dispatch):
-    """Test session existence check."""
-    session_id = uuid4()
-
-    # Initially session should not exist
-    assert not await message_dispatch.session_exists(session_id)
-
-    # After subscribing, session should exist
-    async with message_dispatch.subscribe(session_id, AsyncMock()):
-        assert await message_dispatch.session_exists(session_id)
-
-    # After unsubscribing, session should not exist
-    assert not await message_dispatch.session_exists(session_id)
-
-
-@pytest.mark.anyio
-async def test_session_ttl(message_dispatch):
-    """Test that session has proper TTL set."""
-    session_id = uuid4()
-
-    async with message_dispatch.subscribe(session_id, AsyncMock()):
-        session_key = message_dispatch._session_key(session_id)
-        ttl = await message_dispatch._redis.ttl(session_key)  # type: ignore
-        assert ttl > 0
-        assert ttl <= message_dispatch._session_ttl
-
 
 @pytest.mark.anyio
 async def test_session_heartbeat(message_dispatch):
@@ -129,12 +102,12 @@ async def test_publish_message_invalid_json(message_dispatch):
 
 
 @pytest.mark.anyio
-async def test_publish_to_nonexistent_session(message_dispatch):
+async def test_publish_to_nonexistent_session(message_dispatch: RedisMessageDispatch):
     """Test publishing to a session that doesn't exist."""
     session_id = uuid4()
-    message = types.JSONRPCMessage.model_validate(
+    message = SessionMessage(message=types.JSONRPCMessage.model_validate(
         {"jsonrpc": "2.0", "method": "test", "params": {}, "id": 1}
-    )
+    ))
 
     published = await message_dispatch.publish_message(session_id, message)
     assert not published