Merge pull request #110 from cooperlees/master

Add http proxy support

CHANGELOG.rst

@@ -1,3 +1,10 @@
+Unreleased
+'''''''''''
+
+Added
+- HTTP proxy support for http + websocket calls
+  - Websocket calls are now using aiohttp
+
 7.3.2
 ''''''
 Fixed

requirements.txt

@@ -1,3 +1,3 @@
+aiohttp>=3.8.1<4.0.0
 httpx>=0.20.0<1.0.0
-websockets==9.1
 sphinx-rtd-theme==0.4.1

setup.py

@@ -42,7 +42,7 @@
     packages=find_packages('src'),
     python_requires=">=3.6",
     install_requires=[
-        'websockets>=8',
+        'aiohttp>=3.8.1<4.0.0'
         'httpx>=0.20.0<1.0.0',
     ],
 )

src/mattermostdriver/client.py

@@ -35,6 +35,9 @@ def __init__(self, options):
         self._cookies = None
         self._userid = ""
         self._username = ""
+        self._proxies = None
+        if options["proxy"]:
+            self._proxies = {"all://": options["proxy"]}
 
     @staticmethod
     def _make_url(scheme, url, port, basepath):
@@ -186,7 +189,11 @@ def _get_request_method(method, client):
 class Client(BaseClient):
     def __init__(self, options):
         super().__init__(options)
-        self.client = httpx.Client(http2=options.get("http2", False), verify=options.get("verify", True))
+        self.client = httpx.Client(
+            http2=options.get("http2", False),
+            proxies=self._proxies,
+            verify=options.get("verify", True),
+        )
 
     def make_request(self, method, endpoint, options=None, params=None, data=None, files=None, basepath=None):
         request, url, request_params = self._build_request(method, options, params, data, files, basepath)
@@ -228,7 +235,11 @@ def delete(self, endpoint, options=None, params=None, data=None):
 class AsyncClient(BaseClient):
     def __init__(self, options):
         super().__init__(options)
-        self.client = httpx.AsyncClient(http2=options.get("http2", False), verify=options.get("verify", True))
+        self.client = httpx.AsyncClient(
+            http2=options.get("http2", False),
+            proxies=self._proxies,
+            verify=options.get("verify", True),
+        )
 
     async def __aenter__(self):
         await self.client.__aenter__()

src/mattermostdriver/driver.py

@@ -57,6 +57,7 @@ class BaseDriver:
         "websocket_kw_args": None,
         "debug": False,
         "http2": False,
+        "proxy": None,
     }
     """
 	Required options

src/mattermostdriver/websocket.py

@@ -4,7 +4,7 @@
 import logging
 import time
 
-import websockets
+import aiohttp
 
 log = logging.getLogger("mattermostdriver.websocket")
 log.setLevel(logging.INFO)
@@ -49,21 +49,23 @@ async def connect(self, event_handler):
                 kw_args = {}
                 if self.options["websocket_kw_args"] is not None:
                     kw_args = self.options["websocket_kw_args"]
-                websocket = await websockets.connect(
-                    url,
-                    ssl=context,
-                    **kw_args,
-                )
-                await self._authenticate_websocket(websocket, event_handler)
-                while self._alive:
-                    try:
-                        await self._start_loop(websocket, event_handler)
-                    except websockets.ConnectionClosedError:
-                        break
-                if (not self.options["keepalive"]) or (not self._alive):
-                    break
+                async with aiohttp.ClientSession() as session:
+                    async with session.ws_connect(
+                        url,
+                        ssl=context,
+                        proxy=self.options["proxy"],
+                        **kw_args,
+                    ) as websocket:
+                        await self._authenticate_websocket(websocket, event_handler)
+                        while self._alive:
+                            try:
+                                await self._start_loop(websocket, event_handler)
+                            except aiohttp.ClientError:
+                                break
+                        if (not self.options["keepalive"]) or (not self._alive):
+                            break
             except Exception as e:
-                log.warning(f"Failed to establish websocket connection: {e}")
+                log.exception(f"Failed to establish websocket connection: {type(e)} thrown")
                 await asyncio.sleep(self.options["keepalive_delay"])
 
     async def _start_loop(self, websocket, event_handler):
@@ -77,7 +79,7 @@ async def _start_loop(self, websocket, event_handler):
         keep_alive = asyncio.ensure_future(self._do_heartbeats(websocket))
         log.debug("Waiting for messages on websocket")
         while self._alive:
-            message = await websocket.recv()
+            message = await websocket.receive_str()
             self._last_msg = time.time()
             await event_handler(message)
         log.debug("cancelling heartbeat task")
@@ -118,9 +120,9 @@ async def _authenticate_websocket(self, websocket, event_handler):
         """
         log.debug("Authenticating websocket")
         json_data = json.dumps({"seq": 1, "action": "authentication_challenge", "data": {"token": self._token}})
-        await websocket.send(json_data)
+        await websocket.send_bytes(json_data)
         while True:
-            message = await websocket.recv()
+            message = await websocket.receive_str()
             status = json.loads(message)
             log.debug(status)
             # We want to pass the events to the event_handler already