Merge pull request #144 from livechat/add-custom-header-to-ws-handshake

Allow passing custom header for websocket handshake. Change pre-commi…
livechat · Dec 30, 2024 · 22dbad9 · 22dbad9
2 parents 4d8c3a5 + fbec948
commit 22dbad9
Show file tree

Hide file tree

Showing 12 changed files with 85 additions and 35 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -34,7 +34,7 @@ repos:
              --disable=too-many-locals,
              --disable=duplicate-code,
              --disable=logging-fstring-interpolation]
-- repo: https://gitlab.com/pycqa/flake8
+- repo: https://github.com/pycqa/flake8
   rev: '3.8.4'
   hooks:
     - id: flake8

diff --git a/changelog.md b/changelog.md
@@ -8,10 +8,12 @@ All notable changes to this project will be documented in this file.
 - Added `response_timeout` parameter in `open_connection` methods.
 - New `get_license_info` method in agent-api v3.5.
 - New `update_session` method in agent-api v3.6 (rtm).
+- Allow passing custom header for websocket handshake.
 
 ### Changed
 - Added missing top-level arguments to `update_auto_access` method in configuration-api.
 - Updated outdated packages.
+- Changed pre-commit flake8 URL
 
 ## [0.3.9] - 2024-04-22
 

diff --git a/livechat/agent/rtm/api/v33.py b/livechat/agent/rtm/api/v33.py
@@ -1,6 +1,6 @@
 ''' Module containing Agent RTM API client implementation for v3.3. '''
 
-from typing import Any, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,8 +11,13 @@
 
 class AgentRtmV33:
     ''' Agent RTM API Class containing methods in version 3.3. '''
-    def __init__(self, url: str):
-        self.ws = WebsocketClient(url=f'wss://{url}/v3.3/agent/rtm/ws')
+    def __init__(
+        self,
+        url: str,
+        header: Union[list, dict, Callable, None],
+    ):
+        self.ws = WebsocketClient(url=f'wss://{url}/v3.3/agent/rtm/ws',
+                                  header=header)
 
     def open_connection(self,
                         origin: dict = None,

diff --git a/livechat/agent/rtm/api/v34.py b/livechat/agent/rtm/api/v34.py
@@ -1,6 +1,6 @@
 ''' Module containing Agent RTM API client implementation for v3.4. '''
 
-from typing import Any, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,8 +11,13 @@
 
 class AgentRtmV34:
     ''' Agent RTM API Class containing methods in version 3.4. '''
-    def __init__(self, url: str):
-        self.ws = WebsocketClient(url=f'wss://{url}/v3.4/agent/rtm/ws')
+    def __init__(
+        self,
+        url: str,
+        header: Union[list, dict, Callable, None],
+    ):
+        self.ws = WebsocketClient(url=f'wss://{url}/v3.4/agent/rtm/ws',
+                                  header=header)
 
     def open_connection(self,
                         origin: dict = None,

diff --git a/livechat/agent/rtm/api/v35.py b/livechat/agent/rtm/api/v35.py
@@ -1,6 +1,6 @@
 ''' Module containing Agent RTM API client implementation for v3.5. '''
 
-from typing import Any, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,8 +11,13 @@
 
 class AgentRtmV35:
     ''' Agent RTM API Class containing methods in version 3.5. '''
-    def __init__(self, url: str):
-        self.ws = WebsocketClient(url=f'wss://{url}/v3.5/agent/rtm/ws')
+    def __init__(
+        self,
+        url: str,
+        header: Union[list, dict, Callable, None],
+    ):
+        self.ws = WebsocketClient(url=f'wss://{url}/v3.5/agent/rtm/ws',
+                                  header=header)
 
     def open_connection(self,
                         origin: dict = None,

diff --git a/livechat/agent/rtm/api/v36.py b/livechat/agent/rtm/api/v36.py
@@ -1,6 +1,6 @@
 ''' Module containing Agent RTM API client implementation for v3.6. '''
 
-from typing import Any, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,8 +11,15 @@
 
 class AgentRtmV36:
     ''' Agent RTM API Class containing methods in version 3.6. '''
-    def __init__(self, url: str):
-        self.ws = WebsocketClient(url=f'wss://{url}/v3.6/agent/rtm/ws')
+    def __init__(
+        self,
+        url: str,
+        header: Union[list, dict, Callable, None],
+    ):
+        self.ws = WebsocketClient(
+            url=f'wss://{url}/v3.6/agent/rtm/ws',
+            header=header,
+        )
 
     def open_connection(self,
                         origin: dict = None,

diff --git a/livechat/agent/rtm/base.py b/livechat/agent/rtm/base.py
@@ -3,7 +3,7 @@
 # pylint: disable=W0613,W0622,C0103,R0913,R0903,W0107,W0221
 from __future__ import annotations
 
-from typing import Union
+from typing import Callable, Union
 
 from livechat.agent.rtm.api.v33 import AgentRtmV33
 from livechat.agent.rtm.api.v34 import AgentRtmV34
@@ -20,13 +20,16 @@ class AgentRTM:
     @staticmethod
     def get_client(
         version: str = stable_version,
-        base_url: str = api_url
+        base_url: str = api_url,
+        header: Union[list, dict, Callable, None] = None,
     ) -> Union[AgentRtmV33, AgentRtmV34, AgentRtmV35, AgentRtmV36]:
         ''' Returns client for specific Agent RTM version.
 
             Args:
                 version (str): API's version. Defaults to the stable version of API.
                 base_url (str): API's base url. Defaults to API's production URL.
+                header (Union[list, dict, Callable, None]): Custom header for websocket handshake.
+                        If the parameter is a callable object, it is called just before the connection attempt.
 
             Returns:
                 API client object for specified version.
@@ -42,4 +45,4 @@ def get_client(
         }.get(version)
         if not client:
             raise ValueError('Provided version does not exist.')
-        return client(base_url)
+        return client(base_url, header)
diff --git a/livechat/customer/rtm/api/v33.py b/livechat/customer/rtm/api/v33.py
@@ -2,7 +2,7 @@
 
 # pylint: disable=C0103,R0903,R0913,W0107,W0231,W0613,W0622
 
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,12 +11,17 @@
 
 class CustomerRtmV33:
     ''' Customer RTM API client class in version 3.3. '''
-    def __init__(self, license_id: str, base_url: str):
+    def __init__(
+        self,
+        license_id: str,
+        base_url: str,
+        header: Union[list, dict, Callable, None],
+    ):
         if isinstance(license_id, (int, str)):
             self.ws = WebsocketClient(
                 url=
-                f'wss://{base_url}/v3.3/customer/rtm/ws?license_id={license_id}'
-            )
+                f'wss://{base_url}/v3.3/customer/rtm/ws?license_id={license_id}',
+                header=header)
         else:
             raise ValueError(
                 f'Provided `license_id` (`{license_id}`) seems invalid. Websocket connection may not open.'

diff --git a/livechat/customer/rtm/api/v34.py b/livechat/customer/rtm/api/v34.py
@@ -2,7 +2,7 @@
 
 # pylint: disable=C0103,R0903,R0913,W0107,W0231,W0613,W0622
 
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,12 +11,17 @@
 
 class CustomerRtmV34:
     ''' Customer RTM API client class in version 3.4. '''
-    def __init__(self, organization_id: str, base_url: str):
+    def __init__(
+        self,
+        organization_id: str,
+        base_url: str,
+        header: Union[list, dict, Callable, None],
+    ):
         if isinstance(organization_id, str):
             self.ws = WebsocketClient(
                 url=
-                f'wss://{base_url}/v3.4/customer/rtm/ws?organization_id={organization_id}'
-            )
+                f'wss://{base_url}/v3.5/customer/rtm/ws?organization_id={organization_id}',
+                header=header)
         else:
             raise ValueError(
                 f'Provided `organization_id` (`{organization_id}`) seems invalid. Websocket connection may not open.'

diff --git a/livechat/customer/rtm/api/v35.py b/livechat/customer/rtm/api/v35.py
@@ -2,7 +2,7 @@
 
 # pylint: disable=C0103,R0903,R0913,W0107,W0231,W0613,W0622
 
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,12 +11,17 @@
 
 class CustomerRtmV35:
     ''' Customer RTM API client class in version 3.5. '''
-    def __init__(self, organization_id: str, base_url: str):
+    def __init__(
+        self,
+        organization_id: str,
+        base_url: str,
+        header: Union[list, dict, Callable, None],
+    ):
         if isinstance(organization_id, str):
             self.ws = WebsocketClient(
                 url=
-                f'wss://{base_url}/v3.5/customer/rtm/ws?organization_id={organization_id}'
-            )
+                f'wss://{base_url}/v3.5/customer/rtm/ws?organization_id={organization_id}',
+                header=header)
         else:
             raise ValueError(
                 f'Provided `organization_id` (`{organization_id}`) seems invalid. Websocket connection may not open.'

diff --git a/livechat/customer/rtm/api/v36.py b/livechat/customer/rtm/api/v36.py
@@ -2,7 +2,7 @@
 
 # pylint: disable=C0103,R0903,R0913,W0107,W0231,W0613,W0622
 
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 from livechat.utils.helpers import prepare_payload
 from livechat.utils.structures import AccessToken, RtmResponse
@@ -11,12 +11,17 @@
 
 class CustomerRtmV36:
     ''' Customer RTM API client class in version 3.6. '''
-    def __init__(self, organization_id: str, base_url: str):
+    def __init__(
+        self,
+        organization_id: str,
+        base_url: str,
+        header: Union[list, dict, Callable, None],
+    ):
         if isinstance(organization_id, str):
             self.ws = WebsocketClient(
                 url=
-                f'wss://{base_url}/v3.6/customer/rtm/ws?organization_id={organization_id}'
-            )
+                f'wss://{base_url}/v3.5/customer/rtm/ws?organization_id={organization_id}',
+                header=header)
         else:
             raise ValueError(
                 f'Provided `organization_id` (`{organization_id}`) seems invalid. Websocket connection may not open.'

diff --git a/livechat/customer/rtm/base.py b/livechat/customer/rtm/base.py
@@ -2,7 +2,7 @@
 
 # pylint: disable=C0103,R0903,R0913,W0107,W0231,W0613,W0622
 
-from typing import Union
+from typing import Callable, Union
 
 from livechat.config import CONFIG
 from livechat.customer.rtm.api.v33 import CustomerRtmV33
@@ -21,7 +21,8 @@ def get_client(
         version: str = stable_version,
         base_url: str = api_url,
         license_id: int = None,
-        organization_id: str = None
+        organization_id: str = None,
+        header: Union[list, dict, Callable, None] = None,
     ) -> Union[CustomerRtmV33, CustomerRtmV34, CustomerRtmV35, CustomerRtmV36]:
         ''' Returns client for specific Customer RTM version.
 
@@ -30,6 +31,8 @@ def get_client(
                 base_url (str): API's base url. Defaults to API's production URL.
                 license_id (int): License ID. Required to use for API version <= 3.3.
                 organization_id (str): Organization ID, replaced license ID in v3.4.
+                header (Union[list, dict, Callable, None]): Custom header for websocket handshake.
+                        If the parameter is a callable object, it is called just before the connection attempt.
 
             Returns:
                 API client object for specified version.
@@ -62,5 +65,5 @@ def get_client(
             },
         }.get(version)
         if client:
-            return client(**client_kwargs)
+            return client(**client_kwargs, header=header)
         raise ValueError('Provided version does not exist.')