context manager: provide a context manager for the whole client

in this way we don't need to create a ClientSession context, and we
keep control of the websocket closing.

this is backwards compatible.

Author: mangelajo

README.md

@@ -4,26 +4,54 @@ Async Python client for [NanoKVM](https://github.com/sipeed/NanoKVM).
 
 ## Usage
 
-```python
+### Simple Usage (Recommended)
 
-from aiohttp import ClientSession
-from nanokvm.models import GpioType
-from nanokvm.client import NanoKVMClient
+The client automatically manages the HTTP session:
 
+```python
+from nanokvm.client import NanoKVMClient
+from nanokvm.models import GpioType, MouseButton
 
-async with ClientSession() as session:
-    client = NanoKVMClient("http://kvm-8b76.local/api/", session)
+async with NanoKVMClient("http://kvm-8b76.local/api/") as client:
     await client.authenticate("username", "password")
 
+    # Get device information
     dev = await client.get_info()
     hw = await client.get_hardware()
     gpio = await client.get_gpio()
 
+    # Keyboard input
     await client.paste_text("Hello\nworld!")
 
+    # Mouse control
+    await client.mouse_click(MouseButton.LEFT, 0.5, 0.5)
+    await client.mouse_move_abs(0.25, 0.75)
+    await client.mouse_scroll(0, -3)
+
+    # Stream video
     async for frame in client.mjpeg_stream():
         print(frame)
 
+    # Control GPIO
+    await client.push_button(GpioType.POWER, duration_ms=1000)
+```
+
+### Backwards compatible custom Session)
+
+For more control, you can provide your own `ClientSession`:
+
+```python
+from aiohttp import ClientSession
+from nanokvm.client import NanoKVMClient
+from nanokvm.models import GpioType
+
+async with ClientSession() as session:
+    client = NanoKVMClient("http://kvm-8b76.local/api/", session)
+    await client.authenticate("username", "password")
+
+    dev = await client.get_info()
+    hw = await client.get_hardware()
+
     await client.push_button(GpioType.POWER, duration_ms=1000)
 ```
 

nanokvm/client.py

@@ -109,23 +109,51 @@ class NanoKVMClient:
     def __init__(
         self,
         url: str,
-        session: ClientSession,
+        session: ClientSession | None = None,
         *,
         token: str | None = None,
         request_timeout: int = 10,
     ) -> None:
-        """Initialize the NanoKVM client."""
+        """
+        Initialize the NanoKVM client.
+
+        Args:
+            url: Base URL of the NanoKVM API (e.g., "http://192.168.1.1/api/")
+            session: Optional aiohttp ClientSession. If not provided, one will
+                be created and managed automatically when using the client as
+                an async context manager.
+            token: Optional pre-existing authentication token
+            request_timeout: Request timeout in seconds (default: 10)
+        """
         self.url = yarl.URL(url)
         self.session = session
         self._token = token
         self._request_timeout = request_timeout
         self._ws: aiohttp.ClientWebSocketResponse | None = None
+        self._owns_session = session is None
 
     @property
     def token(self) -> str | None:
         """Return the current auth token."""
         return self._token
 
+    async def __aenter__(self) -> NanoKVMClient:
+        """Async context manager entry."""
+        if self._owns_session:
+            self.session = ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Async context manager exit - cleanup resources."""
+        await self.close()
+
+    async def close(self) -> None:
+        """Close all connections and cleanup resources."""
+        await self.close_ws()
+        if self._owns_session and self.session is not None:
+            await self.session.close()
+            self.session = None
+
     @contextlib.asynccontextmanager
     async def _request(
         self,
@@ -136,6 +164,12 @@ async def _request(
         **kwargs: Any,
     ) -> AsyncIterator[ClientResponse]:
         """Make an API request."""
+        if self.session is None:
+            raise RuntimeError(
+                "Client session not initialized. Use the client as an async "
+                "context manager: 'async with NanoKVMClient(url) as client:'"
+            )
+
         cookies = {}
         if authenticate:
             if not self._token:
@@ -660,6 +694,12 @@ async def set_mouse_jiggler_state(
     async def _get_ws(self) -> aiohttp.ClientWebSocketResponse:
         """Get or create WebSocket connection for mouse events."""
         if self._ws is None or self._ws.closed:
+            if self.session is None:
+                raise RuntimeError(
+                    "Client session not initialized. Use the client as an async "
+                    "context manager: 'async with NanoKVMClient(url) as client:'"
+                )
+
             if not self._token:
                 raise NanoKVMNotAuthenticatedError("Client is not authenticated")
 

tests/test_client.py

@@ -1,3 +1,5 @@
+from unittest.mock import AsyncMock
+
 from aiohttp import ClientSession
 
 from nanokvm.client import NanoKVMClient
@@ -8,3 +10,61 @@ async def test_client() -> None:
     async with ClientSession() as session:
         client = NanoKVMClient("http://localhost:8888/api/", session)
         assert client is not None
+
+
+async def test_client_context_manager_with_session() -> None:
+    """Test the NanoKVMClient as async context manager with provided session."""
+    async with ClientSession() as session:
+        async with NanoKVMClient(
+            "http://localhost:8888/api/", session, token="test-token"
+        ) as client:
+            assert client is not None
+            assert client.session is session
+            # Mock WebSocket to test cleanup
+            mock_ws = AsyncMock()
+            mock_ws.closed = False
+            client._ws = mock_ws
+
+        # After exiting context, WebSocket should be closed
+        mock_ws.close.assert_called_once()
+        # Session should not be closed (we don't own it)
+        assert not session.closed
+
+
+async def test_client_context_manager_without_session() -> None:
+    """Test the NanoKVMClient creates and manages session automatically."""
+    async with NanoKVMClient(
+        "http://localhost:8888/api/", token="test-token"
+    ) as client:
+        assert client is not None
+        assert client.session is not None
+        assert isinstance(client.session, ClientSession)
+        session = client.session
+        # Mock WebSocket to test cleanup
+        mock_ws = AsyncMock()
+        mock_ws.closed = False
+        client._ws = mock_ws
+
+    # After exiting context, both WebSocket and session should be closed
+    mock_ws.close.assert_called_once()
+    assert session.closed
+
+
+async def test_client_without_context_manager_requires_session() -> None:
+    """Test that using client without context manager requires explicit session."""
+    # This should work with explicit session (backward compatible)
+    async with ClientSession() as session:
+        client = NanoKVMClient("http://localhost:8888/api/", session)
+        assert client is not None
+        assert client.session is session
+
+    # Using without session and without context manager should fail on first use
+    client_no_session = NanoKVMClient("http://localhost:8888/api/")
+    assert client_no_session.session is None
+
+    # Should raise RuntimeError when trying to make a request
+    try:
+        await client_no_session.get_info()
+        raise AssertionError("Expected RuntimeError")
+    except RuntimeError as e:
+        assert "context manager" in str(e)