Clean up logic surrounding session auto-creation

Author: puddly

nanokvm/client.py

@@ -152,7 +152,7 @@ def __init__(
         """
         self.url = yarl.URL(url)
         self._session: ClientSession | None = session
-        self._auto_close_session = session is None
+        self._external_session_provided = session is not None
         self._token = token
         self._request_timeout = request_timeout
         self._ws: aiohttp.ClientWebSocketResponse | None = None
@@ -203,7 +203,7 @@ def token(self) -> str | None:
 
     async def __aenter__(self) -> NanoKVMClient:
         """Async context manager entry."""
-        if self._session is None:
+        if self._session is None and not self._external_session_provided:
             self._session = ClientSession()
 
         self._ssl_config = await asyncio.to_thread(self._create_ssl_context)
@@ -217,10 +217,8 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
             self._ws = None
 
         # Close HTTP session
-        if self._session is not None:
-            if self._auto_close_session:
-                await self._session.close()
-
+        if self._session is not None and not self._external_session_provided:
+            await self._session.close()
             self._session = None
 
     @contextlib.asynccontextmanager

tests/test_client.py

@@ -85,18 +85,23 @@ async def test_client_context_manager() -> None:
 async def test_client_context_manager_external_session() -> None:
     """Test that client properly deals with an external session."""
     async with ClientSession() as session:
-        # Both clients connect with the same external session
+        client3 = NanoKVMClient("http://localhost:8888/api/", session=session)
+
+        # All clients connect with the same external session
         async with (
             NanoKVMClient("http://localhost:8888/api/", session=session) as client1,
             NanoKVMClient("http://localhost:8888/api/", session=session) as client2,
+            client3,
         ):
             # Verify session is created
             assert client1._session is session
             assert client2._session is session
+            assert client3._session is session
+
+        # Reusing a client with an external session should not close the session
+        async with client3:
+            assert client3._session is session
 
-        # Both exit but the session is still open
         assert not session.closed
-        assert client1._session is None
-        assert client2._session is None
 
     assert session.closed