feat(firefox): add expect_download to the BiDi tab (portable download API)

thalissonvs · thalissonvs · commit ae13bab716a9 · 2026-05-23T15:42:13.000-03:00
Brings Firefox/BiDi to parity with CDP's Tab.expect_download and makes the
download API portable across both backends.

- BiDiTab.expect_download(keep_file_at, timeout): async context manager that
  routes downloads to a managed directory (temp by default, cleaned up on exit)
  and yields a handle. BiDi reports the saved path on
  browsingContext.downloadEnd, so completion needs no progress polling — it
  filters events to this browsing context and resolves on status 'complete'.
- _BiDiDownloadHandle mirrors CDP's _DownloadHandle: file_path, wait_started,
  wait_finished, read_bytes, read_base64.
- New DownloadHandleProtocol + expect_download on TabProtocol, so callers get
  the same surface regardless of protocol; _conformance.py asserts both concrete
  handles satisfy it.

Verified e2e on real Firefox (temp-dir read, base64, keep_file_at persist);
CDP expect_download still green. mypy/ruff clean.
diff --git a/pydoll/browser/_conformance.py b/pydoll/browser/_conformance.py
@@ -13,10 +13,10 @@
 
 if TYPE_CHECKING:
     from pydoll.browser.chromium.base import Browser
-    from pydoll.browser.chromium.tab import Tab
+    from pydoll.browser.chromium.tab import Tab, _DownloadHandle
     from pydoll.browser.firefox.base import FirefoxBrowser
-    from pydoll.browser.firefox.tab import BiDiTab
-    from pydoll.browser.protocols import BrowserProtocol, TabProtocol
+    from pydoll.browser.firefox.tab import BiDiTab, _BiDiDownloadHandle
+    from pydoll.browser.protocols import BrowserProtocol, DownloadHandleProtocol, TabProtocol
     from pydoll.elements.bidi.shadow_root import BiDiShadowRoot
     from pydoll.elements.bidi.web_element import BiDiWebElement
     from pydoll.elements.cdp.shadow_root import ShadowRoot
@@ -29,6 +29,12 @@ def _cdp_tab_conforms(tab: Tab) -> TabProtocol:
     def _bidi_tab_conforms(tab: BiDiTab) -> TabProtocol:
         return tab
 
+    def _cdp_download_handle_conforms(handle: _DownloadHandle) -> DownloadHandleProtocol:
+        return handle
+
+    def _bidi_download_handle_conforms(handle: _BiDiDownloadHandle) -> DownloadHandleProtocol:
+        return handle
+
     def _cdp_browser_conforms(browser: Browser) -> BrowserProtocol:
         return browser
 
diff --git a/pydoll/browser/firefox/tab.py b/pydoll/browser/firefox/tab.py
@@ -3,13 +3,16 @@
 import asyncio
 import base64 as _b64
 import logging
+import shutil
 from collections.abc import Mapping
 from contextlib import asynccontextmanager
 from pathlib import Path
+from tempfile import mkdtemp
 from typing import TYPE_CHECKING, Optional, Union, overload
 
 import aiofiles
 
+from pydoll.commands.bidi.browser_commands import BrowserCommands
 from pydoll.commands.bidi.browsing_context_commands import BrowsingContextCommands
 from pydoll.commands.bidi.emulation_commands import EmulationCommands
 from pydoll.commands.bidi.input_commands import InputCommands as BiDiInputCommands
@@ -18,11 +21,12 @@
 from pydoll.commands.bidi.storage_commands import StorageCommands
 from pydoll.connection.bidi_connection_handler import BiDiConnectionHandler
 from pydoll.elements.mixins.bidi_find_elements_mixin import BidiFindElementsMixin
-from pydoll.exceptions import ScriptExecutionError
+from pydoll.exceptions import DownloadTimeout, ScriptExecutionError
 from pydoll.interactions.keyboard import BiDiKeyboard
 from pydoll.interactions.mouse import BiDiMouse
 from pydoll.interactions.scroll import BiDiScroll
 from pydoll.protocol.bidi.base import Command, T_CommandParams, T_CommandResult
+from pydoll.protocol.bidi.browser.types import DownloadBehaviorAllowed
 from pydoll.protocol.bidi.browsing_context.types import ImageFormat, ReadinessState
 from pydoll.protocol.bidi.network.types import Cookie as BiDiCookie
 from pydoll.protocol.bidi.network.types import SameSite, StringValue
@@ -413,6 +417,118 @@ async def event_handler(event: dict) -> None:
         finally:
             await self.remove_callback(callback_id)
 
+    @asynccontextmanager
+    async def expect_download(
+        self,
+        keep_file_at: Optional[Union[str, Path]] = None,
+        timeout: Optional[float] = None,
+    ) -> AsyncGenerator[_BiDiDownloadHandle, None]:
+        """Capture a file download triggered inside the block (mirrors the CDP API).
+
+        Routes downloads to a directory and yields a handle to await completion
+        and read the file. BiDi reports the saved path directly on
+        browsingContext.downloadEnd, so no progress polling is needed.
+
+        Args:
+            keep_file_at: Directory to persist the file. If None, a temporary
+                directory is used and removed when the block exits.
+            timeout: Max seconds to wait for the download to finish (default 60).
+
+        Yields:
+            _BiDiDownloadHandle: reads the downloaded file (bytes/base64) and its path.
+        """
+        download_timeout = 60.0 if timeout is None else float(timeout)
+
+        cleanup_dir = False
+        if keep_file_at is None:
+            download_dir = mkdtemp(prefix='pydoll-download-')
+            cleanup_dir = True
+        else:
+            download_dir = str(Path(keep_file_at))
+            Path(download_dir).mkdir(parents=True, exist_ok=True)
+
+        logger.info(f'Expecting download (dir={download_dir}, timeout={download_timeout}s)')
+        await self._execute_command(
+            BrowserCommands.set_download_behavior(
+                download_behavior=DownloadBehaviorAllowed(
+                    type='allowed', destinationFolder=download_dir
+                )
+            )
+        )
+
+        loop = asyncio.get_event_loop()
+        will_begin: asyncio.Future[bool] = loop.create_future()
+        done: asyncio.Future[bool] = loop.create_future()
+        state: dict[str, Optional[str]] = {
+            'url': None,
+            'suggestedFilename': None,
+            'filePath': None,
+        }
+
+        async def on_will_begin(event: dict) -> None:
+            params = event['params']
+            if params.get('context') != self._context_id:
+                return
+            state['url'] = params.get('url')
+            state['suggestedFilename'] = params.get('suggestedFilename')
+            if not will_begin.done():
+                will_begin.set_result(True)
+            logger.info(
+                f'Download will begin: url={state["url"]}, filename={state["suggestedFilename"]}'
+            )
+
+        async def on_end(event: dict) -> None:
+            params = event['params']
+            if params.get('context') != self._context_id or params.get('status') != 'complete':
+                return
+            file_path = params.get('filepath')
+            suggested = state.get('suggestedFilename')
+            if not file_path and suggested:
+                file_path = str(Path(download_dir) / suggested)
+            state['filePath'] = file_path
+            if not done.done():
+                done.set_result(True)
+            logger.info(f'Download completed: {file_path}')
+
+        cb_will_begin = await self.on(Event.DOWNLOAD_STARTED, on_will_begin)
+        cb_end = await self.on(Event.DOWNLOAD_COMPLETED, on_end)
+
+        handle = _BiDiDownloadHandle(
+            state=state,
+            will_begin_future=will_begin,
+            done_future=done,
+            timeout=download_timeout,
+        )
+
+        try:
+            yield handle
+            try:
+                await asyncio.wait_for(done, timeout=download_timeout)
+            except asyncio.TimeoutError as exc:
+                raise DownloadTimeout() from exc
+        finally:
+            await self._cleanup_download_context(
+                [cb_will_begin, cb_end], cleanup_dir, state, download_dir
+            )
+
+    async def _cleanup_download_context(
+        self,
+        callback_ids: list[int],
+        cleanup_dir: bool,
+        state: dict[str, Optional[str]],
+        download_dir: str,
+    ) -> None:
+        for callback_id in callback_ids:
+            await self.remove_callback(callback_id)
+        await self._execute_command(
+            BrowserCommands.set_download_behavior(download_behavior=None)
+        )
+        if cleanup_dir:
+            file_path = state['filePath']
+            if file_path:
+                Path(file_path).unlink(missing_ok=True)
+            shutil.rmtree(download_dir, ignore_errors=True)
+
     @property
     def network_logs(self) -> list:
         """Access captured network logs."""
@@ -637,3 +753,41 @@ def _deserialize_remote_pairs(value: object) -> dict[object, object]:
                 else raw_value
             )
         return result
+
+
+class _BiDiDownloadHandle:
+    """Handle returned by BiDiTab.expect_download to access the downloaded file."""
+
+    def __init__(
+        self,
+        state: dict[str, Optional[str]],
+        will_begin_future: asyncio.Future[bool],
+        done_future: asyncio.Future[bool],
+        timeout: float,
+    ) -> None:
+        self._state = state
+        self._will_begin_future = will_begin_future
+        self._done_future = done_future
+        self._timeout = timeout
+
+    @property
+    def file_path(self) -> Optional[str]:
+        return self._state.get('filePath')
+
+    async def wait_started(self, timeout: Optional[float] = None) -> None:
+        await asyncio.wait_for(self._will_begin_future, timeout=timeout or self._timeout)
+
+    async def wait_finished(self, timeout: Optional[float] = None) -> None:
+        await asyncio.wait_for(self._done_future, timeout=timeout or self._timeout)
+
+    async def read_bytes(self) -> bytes:
+        await self.wait_finished()
+        file_path = self.file_path
+        if not file_path:
+            raise FileNotFoundError('Download file path not available')
+        async with aiofiles.open(file_path, 'rb') as f:
+            return await f.read()
+
+    async def read_base64(self) -> str:
+        data = await self.read_bytes()
+        return _b64.b64encode(data).decode('ascii')
diff --git a/pydoll/browser/protocols.py b/pydoll/browser/protocols.py
@@ -180,6 +180,26 @@ async def intercept_requests(
     async def remove_intercept(self, intercept_id: str) -> None: ...
 
 
+@runtime_checkable
+class DownloadHandleProtocol(Protocol):
+    """Handle to a download captured by ``Tab.expect_download`` (CDP and BiDi).
+
+    Yielded inside the ``expect_download`` block to await completion and read the
+    downloaded bytes, regardless of the underlying protocol.
+    """
+
+    @property
+    def file_path(self) -> Optional[str]: ...
+
+    async def wait_started(self, timeout: Optional[float] = None) -> None: ...
+
+    async def wait_finished(self, timeout: Optional[float] = None) -> None: ...
+
+    async def read_bytes(self) -> bytes: ...
+
+    async def read_base64(self) -> str: ...
+
+
 @runtime_checkable
 class TabProtocol(Protocol):
     """Portable tab-level contract shared by Tab (CDP) and BiDiTab (BiDi).
@@ -283,6 +303,12 @@ def expect_file_chooser(
         self, files: Union[str, Path, list[Union[str, Path]]]
     ) -> AbstractAsyncContextManager[None]: ...
 
+    def expect_download(
+        self,
+        keep_file_at: Optional[Union[str, Path]] = None,
+        timeout: Optional[float] = None,
+    ) -> AbstractAsyncContextManager[DownloadHandleProtocol]: ...
+
     async def on(
         self,
         event_name: Event | str,
diff --git a/tests/bidi/integration/test_bidi_download_integration.py b/tests/bidi/integration/test_bidi_download_integration.py
@@ -0,0 +1,111 @@
+"""Real-Firefox (WebDriver BiDi) integration tests for Tab.expect_download.
+
+A throwaway ``ThreadingHTTPServer`` serves a page with a downloadable attachment.
+Clicking the link inside the ``expect_download`` block must route the file to the
+managed directory and expose it through the handle — the BiDi counterpart of
+tests/cdp/integration/test_tab_io_integration.py::TestExpectDownload. Assertions
+look at the downloaded bytes and the saved file, not the BiDi commands used.
+"""
+
+from __future__ import annotations
+
+import base64
+import socket
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import pytest
+import pytest_asyncio
+
+from pydoll.browser import Firefox
+
+DOWNLOAD_BODY = b'downloaded content from server'
+
+_INDEX_HTML = (
+    '<!DOCTYPE html><html><head><meta charset="utf-8"></head>'
+    '<body><a id="download-link" href="/download" download="hello.txt">download</a>'
+    '</body></html>'
+)
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe:
+        probe.bind(('127.0.0.1', 0))
+        return probe.getsockname()[1]
+
+
+class _DownloadHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        if self.path == '/':
+            body = _INDEX_HTML.encode()
+            self.send_response(200)
+            self.send_header('Content-Type', 'text/html')
+            self.send_header('Content-Length', str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+        elif self.path == '/download':
+            self.send_response(200)
+            self.send_header('Content-Type', 'application/octet-stream')
+            self.send_header('Content-Disposition', 'attachment; filename="hello.txt"')
+            self.send_header('Content-Length', str(len(DOWNLOAD_BODY)))
+            self.end_headers()
+            self.wfile.write(DOWNLOAD_BODY)
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def log_message(self, *args):
+        pass
+
+
+@pytest_asyncio.fixture
+async def served_tab(ci_firefox_options):
+    port = _find_free_port()
+    server = ThreadingHTTPServer(('127.0.0.1', port), _DownloadHandler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    base = f'http://127.0.0.1:{port}'
+    try:
+        async with Firefox(options=ci_firefox_options) as browser:
+            tab = await browser.start()
+            await tab.go_to(f'{base}/')
+            yield tab, base
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join(timeout=5)
+
+
+@pytest.mark.asyncio
+async def test_download_to_temp_dir_is_readable_inside_context(served_tab):
+    tab, _ = served_tab
+    async with tab.expect_download(timeout=30) as handle:
+        link = await tab.find(id='download-link')
+        await link.click()
+        data = await handle.read_bytes()
+        assert data == DOWNLOAD_BODY
+        assert handle.file_path is not None
+
+
+@pytest.mark.asyncio
+async def test_download_base64_matches_bytes(served_tab):
+    tab, _ = served_tab
+    async with tab.expect_download(timeout=30) as handle:
+        link = await tab.find(id='download-link')
+        await link.click()
+        b64 = await handle.read_base64()
+    assert base64.b64decode(b64) == DOWNLOAD_BODY
+
+
+@pytest.mark.asyncio
+async def test_keep_file_at_persists_after_context(served_tab, tmp_path):
+    tab, _ = served_tab
+    target = tmp_path / 'downloads'
+    async with tab.expect_download(keep_file_at=target, timeout=30) as handle:
+        link = await tab.find(id='download-link')
+        await link.click()
+        await handle.wait_finished()
+
+    saved = target / 'hello.txt'
+    assert saved.exists()
+    assert saved.read_bytes() == DOWNLOAD_BODY
