From 1b7a1e1e30a035ae6e3723ef7ac0d9e2ad926c72 Mon Sep 17 00:00:00 2001
From: noahhusby <32528627+noahhusby@users.noreply.github.com>
Date: Fri, 22 Nov 2024 19:39:45 -0500
Subject: [PATCH] Use new async websockets API

---
 aiostreammagic/stream_magic.py | 22 ++++++++++++----------
 pyproject.toml                 |  2 +-
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/aiostreammagic/stream_magic.py b/aiostreammagic/stream_magic.py
index c66b4ab..3e01a51 100644
--- a/aiostreammagic/stream_magic.py
+++ b/aiostreammagic/stream_magic.py
@@ -6,9 +6,8 @@
 from datetime import datetime, UTC
 from typing import Any, Optional
 
-import websockets
-from websockets import WebSocketClientProtocol
-from websockets.client import connect as ws_connect
+from websockets import ConnectionClosedError, ConnectionClosedOK
+from websockets.asyncio.client import connect, ClientConnection
 
 from aiostreammagic.exceptions import StreamMagicError
 from aiostreammagic.models import (
@@ -37,7 +36,7 @@ class StreamMagicClient:
 
     def __init__(self, host: str) -> None:
         self.host = host
-        self.connection: WebSocketClientProtocol | None = None
+        self.connection: ClientConnection | None = None
         self.futures: dict[str, list[Future[Any]]] = {}
         self._subscriptions: dict[str, Any] = {}
         self._loop: AbstractEventLoop = asyncio.get_running_loop()
@@ -110,11 +109,14 @@ def is_connected(self) -> bool:
         """Return True if device is connected."""
         return self.connect_task is not None and not self.connect_task.done()
 
-    async def _ws_connect(self, uri: str) -> WebSocketClientProtocol:
+    async def _ws_connect(self, uri: str) -> ClientConnection:
         """Establish a connection with a WebSocket."""
-        return await ws_connect(
+        return await connect(
             uri,
-            extra_headers={"Origin": f"ws://{self.host}", "Host": f"{self.host}:80"},
+            additional_headers={
+                "Origin": f"ws://{self.host}",
+                "Host": f"{self.host}:80",
+            },
         )
 
     async def _reconnect_handler(self, res: Future[bool]) -> None:
@@ -210,7 +212,7 @@ async def subscription_handler(queue: Queue, callback) -> None:
 
     async def consumer_handler(
         self,
-        ws: WebSocketClientProtocol,
+        ws: ClientConnection,
         subscriptions: dict[str, list[Any]],
         futures: dict[str, list[asyncio.Future]],
     ) -> None:
@@ -240,8 +242,8 @@ async def consumer_handler(
 
         except (
             asyncio.CancelledError,
-            websockets.exceptions.ConnectionClosedError,
-            websockets.exceptions.ConnectionClosedOK,
+            ConnectionClosedError,
+            ConnectionClosedOK,
         ):
             pass
 
diff --git a/pyproject.toml b/pyproject.toml
index 283202a..0878b2b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aiostreammagic"
-version = "2.8.6"
+version = "2.9.0"
 description = "An async python package for interfacing with Cambridge Audio / Stream Magic compatible streamers."
 authors = ["Noah Husby <32528627+noahhusby@users.noreply.github.com>"]
 maintainers = ["Noah Husby <32528627+noahhusby@users.noreply.github.com>"]