Take connection as parameter

Author: jman-sketch

py/selenium/webdriver/common/bidi/network.py

@@ -498,12 +498,12 @@ def __init__(self, conn):
         self.callbacks = {}
 
     async def add_intercept(self, event, params: AddInterceptParameters):
-        await self.bidi.execute(session_subscribe(event.event_class))
-        result = await self.bidi.execute(AddIntercept(params).cmd())
+        await self.conn.execute(session_subscribe(event.event_class))
+        result = await self.conn.execute(AddIntercept(params).cmd())
         return result
 
     async def add_listener(self, event, callback):
-        listener = self.bidi.listen(event)
+        listener = self.conn.listen(event)
 
         async for event in listener:
             request_data = BeforeRequestSentParameters.from_json(
@@ -512,5 +512,5 @@ async def add_listener(self, event, callback):
             await callback(request_data)
 
     async def continue_request(self, params: ContinueRequestParameters):
-        result = await self.bidi.execute(ContinueRequest(params).cmd())
+        result = await self.conn.execute(ContinueRequest(params).cmd())
         return result

py/selenium/webdriver/remote/network.py

@@ -1,7 +1,6 @@
 from dataclasses import fields
 
 from selenium.webdriver.common.bidi import network
-from selenium.webdriver.common.bidi.cdp import open_cdp
 from selenium.webdriver.common.bidi.network import (
     AddInterceptParameters,
     BeforeRequestSent,
@@ -15,20 +14,20 @@ def default_request_handler(params: BeforeRequestSentParameters):
 
 
 class Network:
-    def __init__(self, ws_url):
-        self.ws_url = ws_url
+    def __init__(self, driver):
         self.network = None
-
-    async def add_request_handler(self, request_filter=lambda _: True, handler=default_request_handler):
-        async with open_cdp(self.ws_url) as conn:
-            if not self.network:
-                self.network = network.Network(conn)
-            params = AddInterceptParameters(["beforeRequestSent"])
-            result = await self.network.add_intercept(event=BeforeRequestSent, params=params)
-            intercept = result["intercept"]
-            callback = self._callback(request_filter, handler)
-            await self.network.add_listener(BeforeRequestSent, callback)
-            return intercept
+        self.driver = driver
+        self.intercept = None
+
+    async def add_request_handler(self, request_filter=lambda _: True, handler=default_request_handler, conn=None):
+        self.network = network.Network(conn)
+        params = AddInterceptParameters(["beforeRequestSent"])
+        callback = self._callback(request_filter, handler)
+        result = await self.network.add_intercept(event=BeforeRequestSent, params=params)
+        intercept = result["intercept"]
+        self.intercept = intercept
+        await self.network.add_listener(event=BeforeRequestSent, callback=callback)
+        return intercept
 
     def _callback(self, request_filter, handler):
         async def callback(request):

py/selenium/webdriver/remote/webdriver.py

@@ -1085,7 +1085,7 @@ def script(self):
     @property
     def network(self):
         if not self._network:
-            self._network = Network(self.caps.get("webSocketUrl"))
+            self._network = Network(self)
 
         return self._network