test: browser context tests, part 2

pavelfeldman · pavelfeldman · commit 850db1600976 · 2020-07-14T14:57:23.000-07:00
diff --git a/playwright/browser_context.py b/playwright/browser_context.py
@@ -170,7 +170,9 @@ async def route(self, match: URLMatch, handler: RouteHandler) -> None:
                 "setNetworkInterceptionEnabled", dict(enabled=True)
             )
 
-    async def unroute(self, match: URLMatch, handler: Optional[RouteHandler]) -> None:
+    async def unroute(
+        self, match: URLMatch, handler: Optional[RouteHandler] = None
+    ) -> None:
         self._routes = list(
             filter(
                 lambda r: r.matcher.match != match
diff --git a/tests/assets/playground.html b/tests/assets/playground.html
@@ -0,0 +1,15 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <title>Playground</title>
+  </head>
+  <body>
+    <button>A button</button>
+    <textarea>A text area</textarea>
+    <div id="first">First div</div>
+    <div id="second">
+        Second div
+        <span class="inner">Inner span</span>
+    </div>
+  </body>
+</html>
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -69,7 +69,10 @@ async def context(browser):
     context = await browser.newContext()
     yield context
     await context.close()
-    assert len(browser.contexts) == 0
+    if len(browser.contexts):
+        for context in browser.contexts:
+            await context.close()
+        assert False
 
 
 @pytest.fixture
diff --git a/tests/server.py b/tests/server.py
@@ -51,10 +51,12 @@ def __repr__(self) -> str:
     def start(self):
         request_subscribers = {}
         auth = {}
+        csp = {}
         routes = {}
         gzip_routes = set()
         self.request_subscribers = request_subscribers
         self.auth = auth
+        self.csp = csp
         self.routes = routes
         self.gzip_routes = gzip_routes
         static_path = os.path.join(os.path.dirname(__file__), "assets")
@@ -77,13 +79,15 @@ def process(self):
                             request.getUser(),
                             request.getPassword(),
                         )
-                    if not (authorization_header or creds_correct):
+                    if not creds_correct:
                         request.setHeader(
                             b"www-authenticate", 'Basic realm="Secure Area"'
                         )
                         request.setResponseCode(HTTPStatus.UNAUTHORIZED)
                         request.finish()
                         return
+                if csp.get(uri_path):
+                    request.setHeader(b"Content-Security-Policy", csp[uri_path])
                 if routes.get(uri_path):
                     routes[uri_path](request)
                     return
@@ -123,15 +127,18 @@ def stop(self):
     async def wait_for_request(self, path):
         future = asyncio.Future()
         self.request_subscribers[path] = future
-        future.add_done_callback(lambda f: self.request_subscribers.pop(path, None))
         return await future
 
     def set_auth(self, path: str, username: str, password: str):
         self.auth[path] = (username, password)
 
+    def set_csp(self, path: str, value: str):
+        self.csp[path] = value
+
     def reset(self):
         self.request_subscribers.clear()
         self.auth.clear()
+        self.csp.clear()
         self.gzip_routes.clear()
         self.routes.clear()
 
diff --git a/tests/test_browsercontext.py b/tests/test_browsercontext.py
