Move from trio to threading

Author: jkim2492

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

@@ -201,13 +201,13 @@ class Network:
     def __init__(self, conn):
         self.conn = conn
 
-    async def add_intercept(self, params: AddInterceptParameters):
-        result = await self.conn.execute(AddIntercept(params).cmd())
+    def add_intercept(self, params: AddInterceptParameters):
+        result = self.conn.execute(AddIntercept(params).cmd())
         return result
 
-    async def continue_request(self, params: ContinueRequestParameters):
-        result = await self.conn.execute(ContinueRequest(params).cmd())
+    def continue_request(self, params: ContinueRequestParameters):
+        result = self.conn.execute(ContinueRequest(params).cmd())
         return result
 
-    async def remove_intercept(self, params: RemoveInterceptParameters):
-        await self.conn.execute(RemoveIntercept(params).cmd())
+    def remove_intercept(self, params: RemoveInterceptParameters):
+        self.conn.execute(RemoveIntercept(params).cmd())

py/selenium/webdriver/remote/network.py

@@ -15,14 +15,10 @@
 # specific language governing permissions and limitations
 # under the License.
 from collections import defaultdict
-from contextlib import asynccontextmanager
-
-import trio
 
 from selenium.webdriver.common.bidi import network
 from selenium.webdriver.common.bidi.browsing_context import Navigate
 from selenium.webdriver.common.bidi.browsing_context import NavigateParameters
-from selenium.webdriver.common.bidi.cdp import open_cdp
 from selenium.webdriver.common.bidi.network import AddInterceptParameters
 from selenium.webdriver.common.bidi.network import BeforeRequestSent
 from selenium.webdriver.common.bidi.network import BeforeRequestSentParameters
@@ -32,68 +28,46 @@
 
 
 class Network:
-    def __init__(self, driver):
+    def __init__(self, conn, driver):
+        self.intercepts = defaultdict(lambda: {"event": None, "handlers": []})
+        self.callback_ids = {}
         self.driver = driver
-        self.listeners = {}
-        self.intercepts = defaultdict(lambda: {"event_name": None, "handlers": []})
-        self.bidi_network = None
-        self.conn = None
-        self.nursery = None
+        self.conn = conn
+        self.bidi_network = network.Network(self.conn)
 
         self.remove_request_handler = self.remove_intercept
         self.clear_request_handlers = self.clear_intercepts
 
-    @asynccontextmanager
-    async def set_context(self):
-        ws_url = self.driver.caps.get("webSocketUrl")
-        async with open_cdp(ws_url) as conn:
-            self.conn = conn
-            self.bidi_network = network.Network(self.conn)
-            async with trio.open_nursery() as nursery:
-                self.nursery = nursery
-                yield
-
-    async def get(self, url, wait="complete"):
+    def get(self, url, wait="none"):
         params = NavigateParameters(context=self.driver.current_window_handle, url=url, wait=wait)
-        await self.conn.execute(Navigate(params).cmd())
+        self.conn.execute(Navigate(params).cmd())
 
-    async def add_listener(self, event, callback):
-        event_name = event.event_class
-        if event_name in self.listeners:
-            return
-        self.listeners[event_name] = self.conn.listen(event)
-        try:
-            async for event in self.listeners[event_name]:
-                request_data = event.params
-                if request_data.isBlocked:
-                    await callback(request_data)
-        except trio.ClosedResourceError:
-            pass
-
-    async def add_handler(self, event, handler, urlPatterns=None):
+    def add_handler(self, event, handler, urlPatterns=None):
         event_name = event.event_class
         phase_name = event_name.split(".")[-1]
 
-        await self.conn.execute(session_subscribe(event_name))
+        self.conn.execute(session_subscribe(event_name))
 
         params = AddInterceptParameters(phases=[phase_name], urlPatterns=urlPatterns)
-        result = await self.bidi_network.add_intercept(params)
+        result = self.bidi_network.add_intercept(params)
         intercept = result["intercept"]
 
-        self.intercepts[intercept]["event_name"] = event_name
+        self.intercepts[intercept]["event"] = event
         self.intercepts[intercept]["handlers"].append(handler)
-        self.nursery.start_soon(self.add_listener, event, self.handle_events)
+        if not self.callback_ids.get(event_name):
+            self.callback_ids[event_name] = self.conn.add_callback(event, self.handle_events)
         return intercept
 
-    async def add_request_handler(self, handler, urlPatterns=None):
-        intercept = await self.add_handler(BeforeRequestSent, handler, urlPatterns)
+    def add_request_handler(self, handler, urlPatterns=None):
+        intercept = self.add_handler(BeforeRequestSent, handler, urlPatterns)
         return intercept
 
-    async def handle_events(self, event_params):
-        if isinstance(event_params, BeforeRequestSentParameters):
+    def handle_events(self, event):
+        event_params = event.params
+        if isinstance(event_params, BeforeRequestSentParameters) and event_params.isBlocked:
             json = self.handle_requests(event_params)
             params = ContinueRequestParameters(**json)
-            await self.bidi_network.continue_request(params)
+            self.bidi_network.continue_request(params)
 
     def handle_requests(self, params):
         request = params.request
@@ -102,24 +76,19 @@ def handle_requests(self, params):
                 request = handler(request)
         return request
 
-    async def remove_listener(self, event_name):
-        listener = self.listeners.pop(event_name)
-        listener.close()
-
-    async def remove_intercept(self, intercept):
-        await self.bidi_network.remove_intercept(
+    def remove_intercept(self, intercept):
+        self.bidi_network.remove_intercept(
             params=network.RemoveInterceptParameters(intercept),
         )
-        event_name = self.intercepts.pop(intercept)["event_name"]
-        remaining = [i for i in self.intercepts.values() if i["event_name"] == event_name]
+
+        event = self.intercepts.pop(intercept)["event"]
+        event_name = event.event_class
+
+        remaining = [i for i in self.intercepts.values() if i["event"].event_class == event_name]
         if len(remaining) == 0:
-            await self.remove_listener(event_name)
-            await self.conn.execute(session_unsubscribe(event_name))
+            self.conn.execute(session_unsubscribe(event_name))
+            self.conn.remove_callback(event, self.callback_ids[event_name])
 
-    async def clear_intercepts(self):
+    def clear_intercepts(self):
         for intercept in self.intercepts:
-            await self.remove_intercept(intercept)
-
-    async def disable_cache(self):
-        # Bidi 'network.setCacheBehavior' is not implemented in v130
-        self.driver.execute_cdp_cmd("Network.setCacheDisabled", {"cacheDisabled": True})
+            self.remove_intercept(intercept)

py/selenium/webdriver/remote/webdriver.py

@@ -1253,8 +1253,13 @@ def script(self):
 
     @property
     def network(self):
+        if not self._websocket_connection:
+            self._start_bidi()
+
         if not self._network:
-            self._network = Network(self)
+            # Bidi 'network.setCacheBehavior' is not implemented in v130
+            self.execute_cdp_cmd("Network.setCacheDisabled", {"cacheDisabled": True})
+            self._network = Network(self._websocket_connection, self)
 
         return self._network
 

py/test/selenium/webdriver/common/bidi_network_tests.py

@@ -21,7 +21,7 @@
 
 @pytest.mark.xfail_firefox
 @pytest.mark.xfail_safari
-async def test_request_handler(driver, pages):
+def test_request_handler(driver, pages):
 
     url1 = pages.url("simpleTest.html")
     url2 = pages.url("clicks.html")
@@ -35,22 +35,21 @@ def request_handler(params):
         json = {"request": request, "url": url3}
         return json
 
-    async with driver.network.set_context():
-        # Multiple intercepts
-        intercept1 = await driver.network.add_request_handler(request_handler, pattern1)
-        intercept2 = await driver.network.add_request_handler(request_handler, pattern2)
-        await driver.network.get(url1)
-        assert driver.title == "We Leave From Here"
-        await driver.network.get(url2)
-        assert driver.title == "We Leave From Here"
-
-        # Removal of a single intercept
-        await driver.network.remove_intercept(intercept2)
-        await driver.network.get(url2)
-        assert driver.title == "clicks"
-        await driver.network.get(url1)
-        assert driver.title == "We Leave From Here"
-
-        await driver.network.remove_intercept(intercept1)
-        await driver.network.get(url1)
-        assert driver.title == "Hello WebDriver"
+    # Multiple intercepts
+    intercept1 = driver.network.add_request_handler(request_handler, pattern1)
+    intercept2 = driver.network.add_request_handler(request_handler, pattern2)
+    driver.network.get(url1)
+    assert driver.title == "We Leave From Here"
+    driver.network.get(url2)
+    assert driver.title == "We Leave From Here"
+
+    # Removal of a single intercept
+    driver.network.remove_intercept(intercept2)
+    driver.network.get(url2)
+    assert driver.title == "clicks"
+    driver.network.get(url1)
+    assert driver.title == "We Leave From Here"
+
+    driver.network.remove_intercept(intercept1)
+    driver.network.get(url1)
+    assert driver.title == "Hello WebDriver"