Exit nursery on listener removal

Author: jman-sketch

py/selenium/webdriver/remote/network.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 from dataclasses import fields
-
+import trio
 from selenium.webdriver.common.bidi import network
 from selenium.webdriver.common.bidi.browsing_context import Navigate, NavigateParameters
 from selenium.webdriver.common.bidi.network import (
@@ -35,23 +35,27 @@ def __init__(self, driver):
         self.network = None
         self.driver = driver
         self.intercept = None
+        self.scope = 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
+        with trio.CancelScope() as scope:
+            self.scope = scope
+            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
 
     async def get(self, url, conn):
         params = NavigateParameters(context=self.driver.current_window_handle, url=url)
         await conn.execute(Navigate(params).cmd())
 
     async def remove_request_handler(self):
         await self.network.remove_intercept(event=BeforeRequestSent, params=network.RemoveInterceptParameters(self.intercept))
+        self.scope.cancel()
 
     def _callback(self, request_filter, handler):
         async def callback(request):

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

@@ -20,30 +20,26 @@
 from selenium.webdriver.common.bidi.cdp import open_cdp
 import trio
 
-@pytest.mark.xfail_safari
-@pytest.mark.xfail_firefox
-@pytest.mark.xfail_opera
+
+@pytest.mark.no_driver_after_test
 async def test_add_request_handler(driver):
 
     def request_filter(params: BeforeRequestSentParameters):
         return params.request["url"] == "https://www.example.com/"
 
     def request_handler(params: BeforeRequestSentParameters):
         request = params.request["request"]
-        json = {
-            "request": request,
-            "url" : "https://www.selenium.dev/about/"
-        }
+        json = {"request": request, "url": "https://www.selenium.dev/about/"}
         return ContinueRequestParameters(**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("https://www.example.com",conn)
+            await driver.network.get("https://www.example.com", conn)
             assert "Selenium" in driver.title
             await trio.sleep(1)
             await driver.network.remove_request_handler()
-            await driver.network.get("https://www.example.com",conn)
-            assert "Example" in driver.title
+            await driver.network.get("https://www.example.com", conn)
+            assert "Example" in driver.title