test(client): cover session group cleanup branches

Author: go165

src/mcp/client/session_group.py

@@ -147,7 +147,7 @@ def __init__(
         self._session_exit_stacks = {}
         self._component_name_hook = component_name_hook
 
-    async def __aenter__(self) -> Self:  # pragma: no cover
+    async def __aenter__(self) -> Self:
         # Enter the exit stack only if we created it ourselves
         if self._owns_exit_stack:
             await self._exit_stack.__aenter__()
@@ -158,7 +158,7 @@ async def __aexit__(
         _exc_type: type[BaseException] | None,
         _exc_val: BaseException | None,
         _exc_tb: TracebackType | None,
-    ) -> bool | None:  # pragma: no cover
+    ) -> bool | None:
         """Closes session exit stacks and main exit stack upon completion."""
 
         # Only close the main exit stack if we created it
@@ -323,7 +323,7 @@ async def _establish_session(
             await self._exit_stack.enter_async_context(session_stack)
 
             return result.server_info, session
-        except Exception:  # pragma: no cover
+        except Exception:
             # If anything during this setup fails, ensure the session-specific
             # stack is closed.
             await session_stack.aclose()

tests/client/test_session_group.py

@@ -278,6 +278,24 @@ async def test_client_session_group_disconnect_non_existent_server():
         await group.disconnect_from_server(session)
 
 
+@pytest.mark.anyio
+async def test_client_session_group_context_manager_with_provided_exit_stack():
+    """Provided exit stacks are not entered or closed by the session group."""
+    provided_stack = mock.AsyncMock(spec=contextlib.AsyncExitStack)
+    session_stack = mock.AsyncMock(spec=contextlib.AsyncExitStack)
+    session = mock.Mock(spec=mcp.ClientSession)
+
+    group = ClientSessionGroup(exit_stack=provided_stack)
+    group._session_exit_stacks[session] = session_stack
+
+    assert await group.__aenter__() is group
+    await group.__aexit__(None, None, None)
+
+    provided_stack.__aenter__.assert_not_awaited()
+    provided_stack.aclose.assert_not_awaited()
+    session_stack.aclose.assert_awaited_once()
+
+
 # TODO(Marcelo): This is horrible. We should drop this test.
 @pytest.mark.anyio
 @pytest.mark.parametrize(