Add support for multiple intercepts

Author: jman-sketch

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

@@ -201,20 +201,29 @@ class RemoveIntercept(BidiCommand):
     method: typing.Literal["network.removeIntercept"] = "network.removeIntercept"
 
 
+@dataclass
+class SetCacheBehaviorParameters(BidiObject):
+    cacheBehavior: typing.Literal["default", "bypass"]
+    contexts: typing.Optional[typing.List[browsing_context.BrowsingContext]] = None
+
+
+@dataclass
+class SetCacheBehavior(BidiCommand):
+    params: SetCacheBehaviorParameters
+    method: typing.Literal["network.setCacheBehavior"] = "network.setCacheBehavior"
+
+
 class Network:
     def __init__(self, conn):
         self.conn = conn
-        self.callbacks = {}
 
-    async def add_intercept(self, event, params: AddInterceptParameters):
-        await self.conn.execute(session_subscribe(event.event_class))
+    async def add_intercept(self, params: AddInterceptParameters):
         result = await self.conn.execute(AddIntercept(params).cmd())
         return result
 
     async def continue_request(self, params: ContinueRequestParameters):
         result = await self.conn.execute(ContinueRequest(params).cmd())
         return result
 
-    async def remove_intercept(self, event, params: RemoveInterceptParameters):
-        await self.conn.execute(session_unsubscribe(event.event_class))
+    async def remove_intercept(self, params: RemoveInterceptParameters):
         await self.conn.execute(RemoveIntercept(params).cmd())

py/selenium/webdriver/remote/network.py

@@ -14,6 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from collections import defaultdict
+
 import trio
 
 from selenium.webdriver.common.bidi import network
@@ -23,65 +25,93 @@
 from selenium.webdriver.common.bidi.network import BeforeRequestSent
 from selenium.webdriver.common.bidi.network import BeforeRequestSentParameters
 from selenium.webdriver.common.bidi.network import ContinueRequestParameters
-
-
-def default_request_handler(params: BeforeRequestSentParameters):
-    return ContinueRequestParameters(request=params.request["request"])
+from selenium.webdriver.common.bidi.session import session_subscribe
+from selenium.webdriver.common.bidi.session import session_unsubscribe
 
 
 class Network:
     def __init__(self, driver):
-        self.network = None
         self.driver = driver
-        self.scopes = {}
+        self.listeners = {}
+        self.intercepts = defaultdict(lambda: {"event_name": None, "handlers": []})
+        self.bidi_network = None
         self.conn = None
 
-    async def get(self, url, conn):
-        params = NavigateParameters(context=self.driver.current_window_handle, url=url, wait="complete")
+        self.remove_request_handler = self.remove_intercept
+        self.clear_request_handlers = self.clear_intercepts
+
+    async def get(self, url, conn, wait="complete"):
+        params = NavigateParameters(context=self.driver.current_window_handle, url=url, wait=wait)
         await conn.execute(Navigate(params).cmd())
 
-    def create_callback(self, request_filter, handler):
-        async def callback(request):
-            if request_filter(request):
-                request = handler(request)
-            else:
-                request = default_request_handler(request)
-            await self.network.continue_request(request)
+    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
 
-        return callback
 
-    async def add_listener(self, event, callback):
-        listener = self.conn.listen(event)
-
-        async for event in listener:
-            request_data = BeforeRequestSentParameters.from_json(event.to_json()["params"])
-            if request_data.isBlocked:
-                await callback(request_data)
-
-    async def add_request_handler(
-        self,
-        request_filter=lambda _: True,
-        handler=default_request_handler,
-        conn=None,
-        task_status=trio.TASK_STATUS_IGNORED,
-    ):
+    async def add_handler(self, event, handler, urlPatterns=None, conn=None, task_status=trio.TASK_STATUS_IGNORED):
         if not self.conn:
             self.conn = conn
-        with trio.CancelScope() as scope:
-            self.network = network.Network(conn)
-            params = AddInterceptParameters(["beforeRequestSent"])
-            result = await self.network.add_intercept(event=BeforeRequestSent, params=params)
-            intercept = result["intercept"]
-            self.scopes[intercept] = scope
-            task_status.started(intercept)
-            callback = self.create_callback(request,filter,handler)
-            await self.add_listener(event=BeforeRequestSent, callback=callback)
-            return intercept
-
-    async def remove_request_handler(self, intercept):
-        await self.network.remove_intercept(
-            event=BeforeRequestSent,
-            params=network.RemoveInterceptParameters(self.intercept),
+            self.bidi_network = network.Network(conn)
+
+        event_name = event.event_class
+        phase_name = event_name.split(".")[-1]
+
+        await self.conn.execute(session_subscribe(event_name))
+
+        params = AddInterceptParameters(phases=[phase_name], urlPatterns=urlPatterns)
+        result = await self.bidi_network.add_intercept(params)
+        intercept = result["intercept"]
+
+        self.intercepts[intercept]["event_name"] = event_name
+        self.intercepts[intercept]["handlers"].append(handler)
+        task_status.started(intercept)
+        await self.add_listener(event=event, callback=self.handle_events)
+
+    async def add_request_handler(self, handler, urlPatterns=None, conn=None, task_status=trio.TASK_STATUS_IGNORED):
+        intercept = await self.add_handler(BeforeRequestSent, handler, urlPatterns, conn, task_status)
+        return intercept
+
+    async def handle_events(self, event_params):
+        if isinstance(event_params, BeforeRequestSentParameters):
+            json = self.handle_requests(event_params)
+            params = ContinueRequestParameters(**json)
+            await self.bidi_network.continue_request(params)
+
+    def handle_requests(self, params):
+        request = params.request
+        for intercept in params.intercepts:
+            for handler in self.intercepts[intercept]["handlers"]:
+                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(
+            params=network.RemoveInterceptParameters(intercept),
         )
-        self.scopes[intercept].cancel()
-        self.scopes.pop(intercept)
+        event_name = self.intercepts.pop(intercept)["event_name"]
+        remaining = [i for i in self.intercepts.values() if i["event_name"] == event_name]
+        if len(remaining) == 0:
+            await self.remove_listener(event_name)
+            await self.conn.execute(session_unsubscribe(event_name))
+
+    async 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})

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

@@ -20,36 +20,44 @@
 from selenium.webdriver.common.bidi.cdp import open_cdp
 from selenium.webdriver.common.bidi.network import BeforeRequestSentParameters
 from selenium.webdriver.common.bidi.network import ContinueRequestParameters
+from selenium.webdriver.common.bidi.network import UrlPatternString
+from selenium.webdriver.common.bidi.browsing_context import BrowsingContext
 
 
 @pytest.mark.xfail_firefox
 @pytest.mark.xfail_safari
 @pytest.mark.xfail_edge
-async def test_add_request_handler(driver, pages):
+async def test_request_handler(driver, pages):
 
-    target = pages.url("simpleTest.html")
+    url1 = pages.url("simpleTest.html")
+    url2 = pages.url("clicks.html")
+    url3 = pages.url("formPage.html")
 
-    def request_filter(params: BeforeRequestSentParameters):
-        return params.request["url"] == target
+    pattern1 = [UrlPatternString(url1)]
+    pattern2 = [UrlPatternString(url2)]
 
-    def request_handler(params: BeforeRequestSentParameters):
-        request = params.request["request"]
-        json = {"request": request, "url": pages.url("formPage.html")}
-        return ContinueRequestParameters(**json)
+    def request_handler(params):
+        request = params["request"]
+        json = {"request": request, "url": url3}
+        return json
 
     ws_url = driver.caps.get("webSocketUrl")
     async with open_cdp(ws_url) as conn:
         async with trio.open_nursery() as nursery:
-            nursery.start_soon(
-                driver.network.add_request_handler,
-                request_filter,
-                request_handler,
-                conn,
-            )
-            await trio.sleep(1)
-            await driver.network.get(target, conn)
-            assert "We Leave From Here" == driver.title
-            await trio.sleep(1)
-            await driver.network.remove_request_handler()
-            await driver.network.get(target, conn)
-            assert "Hello WebDriver" == driver.title
+            # Multiple intercepts
+            intercept1 = await nursery.start(driver.network.add_request_handler, request_handler, pattern1, conn)
+            intercept2 = await nursery.start(driver.network.add_request_handler, request_handler, pattern2, conn)
+            await driver.network.get(url1, conn)
+            assert driver.title ==  "We Leave From Here"
+            await driver.network.get(url2, conn)
+            assert  driver.title == "We Leave From Here"
+            
+            # Removal of a single intercept
+            await driver.network.remove_intercept(intercept2)
+            await driver.network.get(url2, conn)
+            assert  driver.title == "clicks"
+            await driver.network.get(url1, conn)
+            assert driver.title == "We Leave From Here"
+            
+            await driver.network.remove_intercept(intercept1)
+            assert driver.title == "We Leave From Here"