unflake and type

Author: Skn0tt

tests/async/test_extension.py

@@ -14,7 +14,7 @@
 
 import asyncio
 from pathlib import Path
-from typing import Any, Awaitable, Callable, Dict, Generator, Optional
+from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional
 
 import pytest
 
@@ -24,20 +24,22 @@
 
 
 @pytest.fixture()
-def launch_persistent_context(
+async def launch_persistent_context(
     browser_type: BrowserType,
     browser_channel: Optional[str],
     tmp_path: Path,
     launch_arguments: Dict[str, Any],
     is_headless_shell: bool,
-) -> Generator[Callable[[str], Awaitable[BrowserContext]], None, None]:
+) -> AsyncGenerator[Callable[..., Awaitable[BrowserContext]], None]:
     if browser_channel and browser_channel.startswith("chrome"):
         pytest.skip(
             "--load-extension is not supported in Chrome anymore. https://groups.google.com/a/chromium.org/g/chromium-extensions/c/1-g8EFx2BBY/m/S0ET5wPjCAAJ"
         )
     if is_headless_shell:
         pytest.skip("Headless Shell has no support for extensions")
 
+    contexts: List[BrowserContext] = []
+
     async def launch(extension_path: str, **kwargs: Any) -> BrowserContext:
         context = await browser_type.launch_persistent_context(
             str(tmp_path),
@@ -48,23 +50,27 @@ async def launch(extension_path: str, **kwargs: Any) -> BrowserContext:
                 f"--load-extension={extension_path}",
             ],
         )
+        contexts.append(context)
         return context
 
     yield launch
 
+    for context in contexts:
+        await context.close()
+
 
 @pytest.mark.only_browser("chromium")
 async def test_should_give_access_to_the_service_worker(
-    launch_persistent_context: Any,
+    launch_persistent_context: Callable[..., Awaitable[BrowserContext]],
     assetdir: Path,
 ) -> None:
     extension_path = str(assetdir / "extension-mv3-simple")
-    context: BrowserContext = await launch_persistent_context(extension_path)
+    context = await launch_persistent_context(extension_path)
     service_workers = context.service_workers
     service_worker = (
         service_workers[0]
         if len(service_workers)
-        else await context.wait_for_event("backgroundpage")
+        else await context.wait_for_event("serviceworker")
     )
     assert service_worker
     assert service_worker in context.service_workers
@@ -76,19 +82,19 @@ async def test_should_give_access_to_the_service_worker(
 
 @pytest.mark.only_browser("chromium")
 async def test_should_give_access_to_the_service_worker_when_recording_video(
-    launch_persistent_context: Any,
+    launch_persistent_context: Callable[..., Awaitable[BrowserContext]],
     tmp_path: Path,
     assetdir: Path,
 ) -> None:
     extension_path = str(assetdir / "extension-mv3-simple")
-    context: BrowserContext = await launch_persistent_context(
+    context = await launch_persistent_context(
         extension_path, record_video_dir=(tmp_path / "videos")
     )
     service_workers = context.service_workers
     service_worker = (
         service_workers[0]
         if len(service_workers)
-        else await context.wait_for_event("backgroundpage")
+        else await context.wait_for_event("serviceworker")
     )
     assert service_worker
     assert service_worker in context.service_workers
@@ -101,12 +107,12 @@ async def test_should_give_access_to_the_service_worker_when_recording_video(
 # https://github.com/microsoft/playwright/issues/32762
 @pytest.mark.only_browser("chromium")
 async def test_should_report_console_messages_from_content_script(
-    launch_persistent_context: Any,
+    launch_persistent_context: Callable[..., Awaitable[BrowserContext]],
     assetdir: Path,
     server: Server,
 ) -> None:
     extension_path = str(assetdir / "extension-mv3-with-logging")
-    context: BrowserContext = await launch_persistent_context(extension_path)
+    context = await launch_persistent_context(extension_path)
     page = await context.new_page()
     [message, _] = await asyncio.gather(
         page.context.wait_for_event(

tests/sync/test_extension.py

@@ -14,7 +14,7 @@
 
 from pathlib import Path
 from time import sleep
-from typing import Any, Callable, Dict, Generator, Optional
+from typing import Any, Callable, Dict, Generator, List, Optional
 
 import pytest
 
@@ -28,14 +28,16 @@ def launch_persistent_context(
     tmp_path: Path,
     launch_arguments: Dict[str, Any],
     is_headless_shell: bool,
-) -> Generator[Callable[[str], BrowserContext], None, None]:
+) -> Generator[Callable[..., BrowserContext], None, None]:
     if browser_channel and browser_channel.startswith("chrome"):
         pytest.skip(
             "--load-extension is not supported in Chrome anymore. https://groups.google.com/a/chromium.org/g/chromium-extensions/c/1-g8EFx2BBY/m/S0ET5wPjCAAJ"
         )
     if is_headless_shell:
         pytest.skip("Headless Shell has no support for extensions")
 
+    contexts: List[BrowserContext] = []
+
     def launch(extension_path: str, **kwargs: Any) -> BrowserContext:
         context = browser_type.launch_persistent_context(
             str(tmp_path),
@@ -46,23 +48,27 @@ def launch(extension_path: str, **kwargs: Any) -> BrowserContext:
                 f"--load-extension={extension_path}",
             ],
         )
+        contexts.append(context)
         return context
 
     yield launch
 
+    for context in contexts:
+        context.close()
+
 
 @pytest.mark.only_browser("chromium")
 def test_should_give_access_to_the_service_worker(
-    launch_persistent_context: Any,
+    launch_persistent_context: Callable[..., BrowserContext],
     assetdir: Path,
 ) -> None:
     extension_path = str(assetdir / "extension-mv3-simple")
-    context: BrowserContext = launch_persistent_context(extension_path)
+    context = launch_persistent_context(extension_path)
     service_workers = context.service_workers
     service_worker = (
         service_workers[0]
         if len(service_workers)
-        else context.wait_for_event("backgroundpage")
+        else context.wait_for_event("serviceworker")
     )
     assert service_worker
     assert service_worker in context.service_workers
@@ -74,19 +80,19 @@ def test_should_give_access_to_the_service_worker(
 
 @pytest.mark.only_browser("chromium")
 def test_should_give_access_to_the_service_worker_when_recording_video(
-    launch_persistent_context: Any,
+    launch_persistent_context: Callable[..., BrowserContext],
     tmp_path: Path,
     assetdir: Path,
 ) -> None:
     extension_path = str(assetdir / "extension-mv3-simple")
-    context: BrowserContext = launch_persistent_context(
+    context = launch_persistent_context(
         extension_path, record_video_dir=(tmp_path / "videos")
     )
     service_workers = context.service_workers
     service_worker = (
         service_workers[0]
         if len(service_workers)
-        else context.wait_for_event("backgroundpage")
+        else context.wait_for_event("serviceworker")
     )
     assert service_worker
     assert service_worker in context.service_workers