Merge pull request #281 from autoscrape-labs/fix/expect-download

Fix expect_download context manager

Author: thalissonvs

pydoll/browser/tab.py

@@ -51,7 +51,6 @@
 )
 from pydoll.protocol.base import EmptyResponse, Response
 from pydoll.protocol.browser.events import (
-    BrowserEvent,
     DownloadProgressEvent,
     DownloadWillBeginEvent,
 )
@@ -851,9 +850,13 @@ async def expect_download(
             behavior=DownloadBehavior.ALLOW,
             download_path=download_dir,
             browser_context_id=self._browser_context_id,
-            events_enabled=True,
         )
 
+        _page_events_was_enabled = True
+        if not self._page_events_enabled:
+            _page_events_was_enabled = False
+            await self.enable_page_events()
+
         loop = asyncio.get_event_loop()
         will_begin: asyncio.Future[bool] = loop.create_future()
         done: asyncio.Future[bool] = loop.create_future()
@@ -889,13 +892,13 @@ async def on_progress(event: DownloadProgressEvent):
             if not done.done():
                 done.set_result(True)
 
-        cb_id_will_begin = await self.on(
-            BrowserEvent.DOWNLOAD_WILL_BEGIN,
+        await self.on(
+            PageEvent.DOWNLOAD_WILL_BEGIN,
             cast(Callable[[dict], Awaitable[Any]], on_will_begin),
             True,
         )
         cb_id_progress = await self.on(
-            BrowserEvent.DOWNLOAD_PROGRESS,
+            PageEvent.DOWNLOAD_PROGRESS,
             cast(Callable[[dict], Awaitable[Any]], on_progress),
             False,
         )
@@ -914,19 +917,37 @@ async def on_progress(event: DownloadProgressEvent):
             except asyncio.TimeoutError as exc:
                 raise DownloadTimeout() from exc
         finally:
-            await self.remove_callback(cb_id_progress)
-            await self.remove_callback(cb_id_will_begin)
-            await self._browser.set_download_behavior(
-                behavior=DownloadBehavior.DEFAULT,
-                browser_context_id=self._browser_context_id,
+            await self._cleanup_download_context(
+                cb_id_progress,
+                _page_events_was_enabled,
+                cleanup_dir,
+                state,
+                download_dir,
             )
 
-            if cleanup_dir:
-                file_path = state['filePath']
-                if not file_path:
-                    return
-                Path(file_path).unlink(missing_ok=True)
-                shutil.rmtree(download_dir, ignore_errors=True)
+    async def _cleanup_download_context(
+        self,
+        cb_id_progress: int,
+        page_events_was_enabled: bool,
+        cleanup_dir: bool,
+        state: dict[str, Any],
+        download_dir: str,
+    ) -> None:
+        await self.remove_callback(cb_id_progress)
+        await self._browser.set_download_behavior(
+            behavior=DownloadBehavior.DEFAULT,
+            browser_context_id=self._browser_context_id,
+        )
+
+        if cleanup_dir:
+            file_path = state['filePath']
+            if not file_path:
+                return
+            Path(file_path).unlink(missing_ok=True)
+            shutil.rmtree(download_dir, ignore_errors=True)
+
+        if not page_events_was_enabled:
+            await self.disable_page_events()
 
     @overload
     async def on(

pydoll/protocol/page/events.py

@@ -273,6 +273,8 @@ class PageEvent(str, Enum):
     Args:
         visible (bool): True if the page is visible.
     """
+    DOWNLOAD_WILL_BEGIN = 'Page.downloadWillBegin'
+    DOWNLOAD_PROGRESS = 'Page.downloadProgress'
 
 
 class DomContentEventFiredEventParams(TypedDict):

tests/test_browser/test_browser_tab.py

@@ -11,7 +11,7 @@
 from pydoll.protocol.fetch.types import RequestStage
 from pydoll.constants import By
 from pydoll.browser.tab import Tab
-from pydoll.protocol.browser.events import BrowserEvent
+from pydoll.protocol.page.events import PageEvent
 from pydoll.protocol.browser.types import DownloadBehavior
 from pydoll.exceptions import DownloadTimeout, InvalidTabInitialization
 from pydoll.exceptions import (
@@ -1077,13 +1077,13 @@ async def test_expect_download_keeps_file_when_path_provided(self, tab, tmp_path
 
         async def fake_on(event_name, handler, temporary=False):
             handlers[event_name] = handler
-            return 100 if event_name == BrowserEvent.DOWNLOAD_WILL_BEGIN else 101
+            return 100 if event_name == PageEvent.DOWNLOAD_WILL_BEGIN else 101
 
         with patch.object(tab, 'on', fake_on):
             async with tab.expect_download(keep_file_at=str(target_dir)) as download:
                 # Simulate willBegin
-                await handlers[BrowserEvent.DOWNLOAD_WILL_BEGIN]({
-                    'method': BrowserEvent.DOWNLOAD_WILL_BEGIN,
+                await handlers[PageEvent.DOWNLOAD_WILL_BEGIN]({
+                    'method': PageEvent.DOWNLOAD_WILL_BEGIN,
                     'params': {
                         'frameId': 'frame-1',
                         'guid': 'guid-1',
@@ -1092,8 +1092,8 @@ async def fake_on(event_name, handler, temporary=False):
                     }
                 })
                 # Simulate progress Completed without filePath (fallback to suggested)
-                await handlers[BrowserEvent.DOWNLOAD_PROGRESS]({
-                    'method': BrowserEvent.DOWNLOAD_PROGRESS,
+                await handlers[PageEvent.DOWNLOAD_PROGRESS]({
+                    'method': PageEvent.DOWNLOAD_PROGRESS,
                     'params': {
                         'guid': 'guid-1',
                         'totalBytes': 10,
@@ -1122,14 +1122,14 @@ async def test_expect_download_timeout_raises(self, tab, tmp_path):
 
         async def fake_on(event_name, handler, temporary=False):
             handlers[event_name] = handler
-            return 200 if event_name == BrowserEvent.DOWNLOAD_WILL_BEGIN else 201
+            return 200 if event_name == PageEvent.DOWNLOAD_WILL_BEGIN else 201
 
         with patch.object(tab, 'on', fake_on):
             with pytest.raises(DownloadTimeout):
                 async with tab.expect_download(keep_file_at=str(tmp_path), timeout=0.01):
                     # Trigger will begin but never complete
-                    await handlers[BrowserEvent.DOWNLOAD_WILL_BEGIN]({
-                        'method': BrowserEvent.DOWNLOAD_WILL_BEGIN,
+                    await handlers[PageEvent.DOWNLOAD_WILL_BEGIN]({
+                        'method': PageEvent.DOWNLOAD_WILL_BEGIN,
                         'params': {
                             'frameId': 'frame-1',
                             'guid': 'guid-2',
@@ -1147,22 +1147,22 @@ async def test_expect_download_cleans_temp_directory(self, tab, tmp_path):
 
         async def fake_on(event_name, handler, temporary=False):
             handlers[event_name] = handler
-            return 300 if event_name == BrowserEvent.DOWNLOAD_WILL_BEGIN else 301
+            return 300 if event_name == PageEvent.DOWNLOAD_WILL_BEGIN else 301
 
         with patch.object(tab, 'on', fake_on):
             # Use None to create temp dir and ensure cleanup occurs
             async with tab.expect_download(keep_file_at=None) as download:
-                await handlers[BrowserEvent.DOWNLOAD_WILL_BEGIN]({
-                    'method': BrowserEvent.DOWNLOAD_WILL_BEGIN,
+                await handlers[PageEvent.DOWNLOAD_WILL_BEGIN]({
+                    'method': PageEvent.DOWNLOAD_WILL_BEGIN,
                     'params': {
                         'frameId': 'frame-1',
                         'guid': 'guid-3',
                         'url': 'https://example.com/tmp.txt',
                         'suggestedFilename': 'tmp.txt',
                     }
                 })
-                await handlers[BrowserEvent.DOWNLOAD_PROGRESS]({
-                    'method': BrowserEvent.DOWNLOAD_PROGRESS,
+                await handlers[PageEvent.DOWNLOAD_PROGRESS]({
+                    'method': PageEvent.DOWNLOAD_PROGRESS,
                     'params': {
                         'guid': 'guid-3',
                         'totalBytes': 3,
@@ -1182,6 +1182,140 @@ async def fake_on(event_name, handler, temporary=False):
             # We cannot know the exact temp dir path (random), but ensure file is gone
             assert not file_path.exists()
 
+    @pytest.mark.asyncio
+    async def test_expect_download_ignores_progress_with_different_guid(self, tab, tmp_path):
+        tab._browser.set_download_behavior = AsyncMock()
+
+        handlers = {}
+
+        async def fake_on(event_name, handler, temporary=False):
+            handlers[event_name] = handler
+            return 400 if event_name == PageEvent.DOWNLOAD_WILL_BEGIN else 401
+
+        with patch.object(tab, 'on', fake_on):
+            async with tab.expect_download(keep_file_at=str(tmp_path)) as download:
+                await handlers[PageEvent.DOWNLOAD_WILL_BEGIN]({
+                    'method': PageEvent.DOWNLOAD_WILL_BEGIN,
+                    'params': {
+                        'frameId': 'frame-1',
+                        'guid': 'guid-x',
+                        'url': 'https://example.com/file.bin',
+                        'suggestedFilename': 'file.bin',
+                    }
+                })
+
+                # Wrong guid should be ignored and not mark as done
+                await handlers[PageEvent.DOWNLOAD_PROGRESS]({
+                    'method': PageEvent.DOWNLOAD_PROGRESS,
+                    'params': {
+                        'guid': 'wrong-guid',
+                        'totalBytes': 1,
+                        'receivedBytes': 1,
+                        'state': 'completed',
+                    }
+                })
+
+                # Still not finished
+                assert download.file_path is None
+
+                # Correct guid completes
+                await handlers[PageEvent.DOWNLOAD_PROGRESS]({
+                    'method': PageEvent.DOWNLOAD_PROGRESS,
+                    'params': {
+                        'guid': 'guid-x',
+                        'totalBytes': 10,
+                        'receivedBytes': 10,
+                        'state': 'completed',
+                        'filePath': str(tmp_path / 'file.bin'),
+                    }
+                })
+
+                await download.wait_finished()
+
+    @pytest.mark.asyncio
+    async def test_expect_download_page_events_auto_enable_disable(self, tab, tmp_path):
+        """When page events are disabled, expect_download should enable and then disable them."""
+        tab._browser.set_download_behavior = AsyncMock()
+        tab._page_events_enabled = False
+
+        enable_page_events = AsyncMock()
+        disable_page_events = AsyncMock()
+
+        handlers = {}
+
+        async def fake_on(event_name, handler, temporary=False):
+            handlers[event_name] = handler
+            return 500 if event_name == PageEvent.DOWNLOAD_WILL_BEGIN else 501
+
+        with patch.object(tab, 'enable_page_events', enable_page_events), \
+             patch.object(tab, 'disable_page_events', disable_page_events), \
+             patch.object(tab, 'on', fake_on):
+            async with tab.expect_download(keep_file_at=str(tmp_path)):
+                await handlers[PageEvent.DOWNLOAD_WILL_BEGIN]({
+                    'method': PageEvent.DOWNLOAD_WILL_BEGIN,
+                    'params': {
+                        'frameId': 'frame-1',
+                        'guid': 'guid-y',
+                        'url': 'https://example.com/auto.bin',
+                        'suggestedFilename': 'auto.bin',
+                    }
+                })
+                await handlers[PageEvent.DOWNLOAD_PROGRESS]({
+                    'method': PageEvent.DOWNLOAD_PROGRESS,
+                    'params': {
+                        'guid': 'guid-y',
+                        'totalBytes': 2,
+                        'receivedBytes': 2,
+                        'state': 'completed',
+                        'filePath': str(tmp_path / 'auto.bin'),
+                    }
+                })
+
+        enable_page_events.assert_awaited_once()
+        disable_page_events.assert_awaited_once()
+
+    @pytest.mark.asyncio
+    async def test_expect_download_keeps_page_events_enabled_when_already_enabled(self, tab, tmp_path):
+        """When page events already enabled, expect_download should not disable them on exit."""
+        tab._browser.set_download_behavior = AsyncMock()
+        tab._page_events_enabled = True
+
+        enable_page_events = AsyncMock()
+        disable_page_events = AsyncMock()
+
+        handlers = {}
+
+        async def fake_on(event_name, handler, temporary=False):
+            handlers[event_name] = handler
+            return 600 if event_name == PageEvent.DOWNLOAD_WILL_BEGIN else 601
+
+        with patch.object(tab, 'enable_page_events', enable_page_events), \
+             patch.object(tab, 'disable_page_events', disable_page_events), \
+             patch.object(tab, 'on', fake_on):
+            async with tab.expect_download(keep_file_at=str(tmp_path)):
+                await handlers[PageEvent.DOWNLOAD_WILL_BEGIN]({
+                    'method': PageEvent.DOWNLOAD_WILL_BEGIN,
+                    'params': {
+                        'frameId': 'frame-1',
+                        'guid': 'guid-z',
+                        'url': 'https://example.com/enabled.bin',
+                        'suggestedFilename': 'enabled.bin',
+                    }
+                })
+                await handlers[PageEvent.DOWNLOAD_PROGRESS]({
+                    'method': PageEvent.DOWNLOAD_PROGRESS,
+                    'params': {
+                        'guid': 'guid-z',
+                        'totalBytes': 2,
+                        'receivedBytes': 2,
+                        'state': 'completed',
+                        'filePath': str(tmp_path / 'enabled.bin'),
+                    }
+                })
+
+        enable_page_events.assert_not_awaited()
+        disable_page_events.assert_not_awaited()
+
     @pytest.mark.asyncio
     async def test_disable_auto_solve_cloudflare_captcha(self, tab):
         """Test disabling auto-solve Cloudflare captcha."""