Close session on disconnect and add context manager (noahhusby#151)

Author: Solmath

aiostreammagic/stream_magic.py

@@ -58,6 +58,7 @@ def __init__(self, host: str, session: ClientSession | None = None) -> None:
         self._attempt_reconnection = False
         self._reconnect_task: Optional[Task[Any]] = None
         self.position_last_updated: datetime = datetime.now()
+        self._subscription_tasks: dict[str, asyncio.Task[Any]] = {}
 
     async def register_state_update_callbacks(self, callback: Any) -> None:
         """Register state update callback."""
@@ -109,6 +110,14 @@ async def disconnect(self) -> None:
                 except asyncio.CancelledError:
                     pass
             await self.do_state_update_callbacks(CallbackType.CONNECTION)
+        # Cancel all subscription handler tasks
+        for task in self._subscription_tasks.values():
+            task.cancel()
+        await asyncio.gather(*self._subscription_tasks.values(), return_exceptions=True)
+        self._subscription_tasks.clear()
+        # Properly close the aiohttp session if it was created by this client
+        if self.session is not None and not self.session.closed:
+            await self.session.close()
 
     def is_connected(self) -> bool:
         """Return True if device is connected."""
@@ -230,7 +239,6 @@ async def consumer_handler(
     ) -> None:
         """Callback consumer handler."""
         subscription_queues = {}
-        subscription_tasks = {}
         try:
             async for raw_msg in ws:
                 if futures or subscriptions:
@@ -244,14 +252,13 @@ async def consumer_handler(
                             if not future.done():
                                 future.set_result(msg)
                     if subscription:
-                        if path not in subscription_tasks:
+                        if path not in self._subscription_tasks:
                             queue: Queue[dict[str, Any]] = asyncio.Queue()
                             subscription_queues[path] = queue
-                            subscription_tasks[path] = asyncio.create_task(
+                            self._subscription_tasks[path] = asyncio.create_task(
                                 self.subscription_handler(queue, subscription)
                             )
                         subscription_queues[path].put_nowait(msg)
-
         except (asyncio.CancelledError,):
             pass
 
@@ -645,3 +652,10 @@ async def set_auto_power_down(self, auto_power_down_time_seconds: int) -> None:
         await self.request(
             ep.POWER, params={"auto_power_down": auto_power_down_time_seconds}
         )
+
+    async def __aenter__(self):
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.disconnect()