Resolve "Alternative Kraken instances not working with websockets (Spot)" (#402)

Author: btschwertfeger

src/kraken/spot/websocket/__init__.py

@@ -58,8 +58,8 @@ class SpotWSClientBase(SpotAsyncClient):
     """
 
     LOG: logging.Logger = logging.getLogger(__name__)
-    PROD_ENV_URL: str = "ws.kraken.com"
-    AUTH_PROD_ENV_URL: str = "ws-auth.kraken.com"
+    WS_URL: str = "wss://ws.kraken.com"
+    AUTH_WS_URL: str = "wss://ws-auth.kraken.com"
     # Changing this can cause errors, as this class is designed for v2.
     API_V: str = "/v2"
 
@@ -70,8 +70,14 @@ def __init__(  # nosec: B107
         callback: Callable | None = None,
         *,
         no_public: bool = False,
+        rest_url: str | None = None,
+        ws_url: str | None = None,
+        auth_ws_url: str | None = None,
     ) -> None:
-        super().__init__(key=key, secret=secret)
+        super().__init__(key=key, secret=secret, url=rest_url)
+        self.WS_URL = ws_url or self.WS_URL
+        self.AUTH_WS_URL = auth_ws_url or self.AUTH_WS_URL
+
         self.state: WSState = WSState.INIT
         self._is_auth: bool = bool(key and secret)
         self.__callback: Callable | None = callback
@@ -92,13 +98,13 @@ def __prepare_connect(self: SpotWSClientBase, *, no_public: bool) -> None:
         """Set up functions and attributes based on the API version."""
 
         # pylint: disable=invalid-name
-        self.PROD_ENV_URL += self.API_V
-        self.AUTH_PROD_ENV_URL += self.API_V
+        self.WS_URL += self.API_V
+        self.AUTH_WS_URL += self.API_V
 
         self._pub_conn = (
             ConnectSpotWebsocket(
                 client=self,
-                endpoint=self.PROD_ENV_URL,
+                endpoint=self.WS_URL,
                 is_auth=False,
                 callback=self.on_message,
             )
@@ -109,7 +115,7 @@ def __prepare_connect(self: SpotWSClientBase, *, no_public: bool) -> None:
         self._priv_conn = (
             ConnectSpotWebsocket(
                 client=self,
-                endpoint=self.AUTH_PROD_ENV_URL,
+                endpoint=self.AUTH_WS_URL,
                 is_auth=True,
                 callback=self.on_message,
             )

src/kraken/spot/websocket/connectors.py

@@ -147,7 +147,7 @@ async def __run(self: ConnectSpotWebsocketBase, event: asyncio.Event) -> None:
         LOG.debug("Websocket token: %s", self.ws_conn_details)
 
         async with connect(  # pylint: disable=no-member
-            f"wss://{self.__ws_endpoint}",
+            self.__ws_endpoint,
             additional_headers={"User-Agent": "btschwertfeger/python-kraken-sdk"},
             ping_interval=30,
             max_queue=None,  # FIXME: This is not recommended by the docs https://websockets.readthedocs.io/en/stable/reference/asyncio/client.html#module-websockets.asyncio.client

src/kraken/spot/ws_client.py

@@ -61,9 +61,14 @@ class SpotWSClient(SpotWSClientBase):
         set or set to ``False``, the client will create a public and a private
         connection per default. If only a private connection is required, this
         parameter should be set to ``True``.
-    :param beta: Use the beta websocket channels (maybe not supported anymore,
-        default: ``False``)
-    :type beta: bool
+    :param rest_url: Set a specific REST URL to access the Kraken Spot API
+        (default: ``None``).
+    :type rest_url: str, optional
+    :param ws_url: Set a specific Websocket URL to access the Kraken Spot API
+        (default: ``None``).
+    :type ws_url: str, optional
+    :param auth_ws_url: Set a specific Authenticated Websocket URL to access
+        the Kraken Spot API (default: ``None``).
 
     .. code-block:: python
         :linenos:
@@ -163,6 +168,18 @@ async def main():
                 asyncio.run(main())
             except KeyboardInterrupt:
                 pass
+
+    .. code-block:: python
+        :linenos:
+        :caption: HowTo: How to connect to other Kraken instances
+
+        client_auth = SpotWSClient(
+            key="api-key",
+            secret="secret-key",
+            rest_url="https://api.vip.uat.lobster.kraken.com",
+            ws_url="wss://ws.vip.uat.lobster.kraken.com",
+            auth_ws_url="wss://ws-auth.vip.uat.lobster.kraken.com",
+        )
     """
 
     def __init__(  # nosec: B107
@@ -172,12 +189,18 @@ def __init__(  # nosec: B107
         callback: Callable | None = None,
         *,
         no_public: bool = False,
+        rest_url: str | None = None,
+        ws_url: str | None = None,
+        auth_ws_url: str | None = None,
     ) -> None:
         super().__init__(
             key=key,
             secret=secret,
             callback=callback,
             no_public=no_public,
+            rest_url=rest_url,
+            ws_url=ws_url,
+            auth_ws_url=auth_ws_url,
         )
 
     async def send_message(  # noqa: C901 # pylint: disable=arguments-differ