Merge pull request #63 from cleitonleonel/main

Handle case where Browser.WindowID is not found and improve error handling.

Author: thalissonvs

pydoll/browser/base.py

@@ -77,7 +77,9 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
             exc_val: The exception value, if raised.
             exc_tb: The traceback, if an exception was raised.
         """
-        await self.stop()
+        if await self._is_browser_running():
+            await self.stop()
+
         await self._connection_handler.close()
 
     async def start(self) -> None:
@@ -137,6 +139,7 @@ async def get_page(self) -> Page:
         page_id = (
             await self.new_page() if not self._pages else self._pages.pop()
         )
+
         return Page(self._connection_port, page_id)
 
     async def delete_all_cookies(self):
@@ -243,10 +246,25 @@ async def get_window_id(self):
         Retrieves the ID of the current browser window.
 
         Returns:
-            str: The ID of the current browser window.
+            int: The ID of the current browser window.
+
+        Raises:
+            RuntimeError: If unable to retrieve the window ID.
         """
-        response = await self._execute_command(BrowserCommands.get_window_id())
-        return response['result']['windowId']
+        command = BrowserCommands.get_window_id()
+        response = await self._execute_command(command)
+
+        if response.get('error'):
+            pages = await self.get_targets()
+            target_id = await self._get_valid_target_id(pages)
+            response = await self._execute_command(
+                BrowserCommands.get_window_id_by_target(target_id)
+            )
+
+        if window_id := response.get('result', {}).get('windowId'):
+            return window_id
+
+        raise RuntimeError(response.get('error', {}))
 
     async def set_window_bounds(self, bounds: dict):
         """
@@ -519,7 +537,7 @@ def _is_valid_page(page: dict) -> bool:
             'url', ''
         )
 
-    async def _get_valid_page(self, pages) -> str:
+    async def _get_valid_page(self, pages: list) -> str:
         """
         Gets the ID of a valid page or creates a new one.
 
@@ -541,6 +559,31 @@ async def _get_valid_page(self, pages) -> str:
 
         return await self.new_page()
 
+    @staticmethod
+    async def _get_valid_target_id(pages: list) -> str:
+        """
+        Retrieves the target ID of a valid attached browser page.
+
+        Returns:
+            str: The target ID of a valid page.
+
+        """
+
+        valid_page = next(
+            (page for page in pages
+             if page.get('type') == 'page' and page.get('attached')),
+            None
+        )
+
+        if not valid_page:
+            raise RuntimeError("No valid attached browser page found.")
+
+        target_id = valid_page.get('targetId')
+        if not target_id:
+            raise RuntimeError("Valid page found but missing 'targetId'.")
+
+        return target_id
+
     async def _is_browser_running(self, timeout: int = 10) -> bool:
         """
         Checks if the browser process is currently running.
@@ -553,9 +596,10 @@ async def _is_browser_running(self, timeout: int = 10) -> bool:
             if await self._connection_handler.ping():
                 return True
             await asyncio.sleep(1)
+
         return False
 
-    async def _execute_command(self, command: str):
+    async def _execute_command(self, command: dict):
         """
         Executes a command through the connection handler.
 

pydoll/browser/page.py

@@ -157,6 +157,7 @@ async def has_dialog(self) -> bool:
         """
         if self._connection_handler.dialog:
             return True
+
         return False
 
     async def get_dialog_message(self) -> str:

pydoll/commands/browser.py

@@ -17,6 +17,10 @@ class BrowserCommands:
 
     CLOSE = {'method': 'Browser.close'}
     GET_WINDOW_ID = {'method': 'Browser.WindowID'}
+    GET_WINDOW_ID_BY_TARGET = {
+        'method': 'Browser.getWindowForTarget',
+        'params': {},
+    }
     SET_WINDOW_BOUNDS_TEMPLATE = {
         'method': 'Browser.setWindowBounds',
         'params': {},
@@ -62,6 +66,21 @@ def get_window_id(cls) -> dict:
         """
         return cls.GET_WINDOW_ID
 
+    @classmethod
+    def get_window_id_by_target(cls, target_id: str) -> dict:
+        """
+        Generates the command to get the ID of the current window.
+
+        Args:
+            target_id (str): The target_id to set for the window.
+
+        Returns:
+            dict: The command to be sent to the browser.
+        """
+        command = cls.GET_WINDOW_ID_BY_TARGET.copy()
+        command['params']['targetId'] = target_id
+        return command
+
     @classmethod
     def set_window_bounds(cls, window_id: int, bounds: dict) -> dict:
         """

tests/test_browser_window.py

@@ -0,0 +1,108 @@
+import pytest
+import pytest_asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+from pydoll.browser.base import Browser
+from pydoll.commands.browser import BrowserCommands
+
+
+class ConcreteBrowser(Browser):
+    def _get_default_binary_location(self) -> str:
+        return '/fake/path/to/browser'
+
+
+@pytest_asyncio.fixture
+async def mock_browser():
+    with patch.multiple(
+            Browser,
+            _get_default_binary_location=MagicMock(
+                return_value='/fake/path/to/browser'
+            ),
+    ), patch(
+        'pydoll.connection.connection.ConnectionHandler',
+        autospec=True,
+    ) as mock_conn_handler:
+        browser = ConcreteBrowser()
+        browser._connection_handler = mock_conn_handler.return_value
+        browser._connection_handler.execute_command = AsyncMock()
+        browser._connection_handler.register_callback = AsyncMock()
+
+        browser._pages = ['page1']
+
+        yield browser
+
+
+@pytest.mark.asyncio
+async def test_get_window_id_success(mock_browser):
+    mock_browser._connection_handler.execute_command.return_value = {
+        'result': {'windowId': 123}
+    }
+    result = await mock_browser.get_window_id()
+    assert result == 123
+
+
+@pytest.mark.asyncio
+async def test_get_window_id_with_error_and_retry(mock_browser):
+    mock_browser._execute_command = AsyncMock(side_effect=[
+        {'error': 'some error'},
+        {'result': {
+            'targetInfos': [{
+                'type': 'page',
+                'attached': True,
+                'targetId': 'target1'
+            }]
+        }},
+        {'result': {'windowId': 123}}
+    ])
+
+    result = await mock_browser.get_window_id()
+    assert result == 123
+
+
+@pytest.mark.asyncio
+async def test_get_window_id_failure(mock_browser):
+    mock_browser._connection_handler.execute_command.return_value = {
+        'error': 'some error'
+    }
+    mock_browser.get_targets = AsyncMock(return_value=[])
+    with pytest.raises(RuntimeError):
+        await mock_browser.get_window_id()
+
+
+@pytest.mark.asyncio
+async def test_get_valid_target_id_success(mock_browser):
+    pages = [{
+        'type': 'page',
+        'attached': True,
+        'targetId': 'target1'
+    }]
+    result = await mock_browser._get_valid_target_id(pages)
+    assert result == 'target1'
+
+
+@pytest.mark.asyncio
+async def test_get_valid_target_id_no_valid_page(mock_browser):
+    pages = []
+    with pytest.raises(RuntimeError, match="No valid attached browser page found."):
+        await mock_browser._get_valid_target_id(pages)
+
+
+@pytest.mark.asyncio
+async def test_get_valid_target_id_missing_target_id(mock_browser):
+    pages = [{
+        'type': 'page',
+        'attached': True,
+        'targetId': None
+    }]
+    with pytest.raises(RuntimeError, match="Valid page found but missing 'targetId'."):
+        await mock_browser._get_valid_target_id(pages)
+
+
+@pytest.mark.asyncio
+async def test_get_window_id_by_target(mock_browser):
+    expected_command = {
+        'method': 'Browser.getWindowForTarget',
+        'params': {
+            'targetId': 'target1',
+        },
+    }
+    assert BrowserCommands.get_window_id_by_target('target1') == expected_command