diff --git a/sdk/harambe/core.py b/sdk/harambe/core.py index af2b952..847c946 100644 --- a/sdk/harambe/core.py +++ b/sdk/harambe/core.py @@ -215,7 +215,10 @@ async def paginate( return async def capture_url( - self, clickable: ElementHandle, resource_type: ResourceType = "document" + self, + clickable: ElementHandle, + resource_type: ResourceType = "document", + timeout: Optional[int] = 10, ) -> URL | None: """ Capture the url of a click event. This will click the element and return the url @@ -224,11 +227,12 @@ async def capture_url( :param clickable: the element to click :param resource_type: the type of resource to capture + :param timeout: the time to wait for the new page to open (in seconds) :return url: the url of the captured resource or None if no match was found :raises ValueError: if more than one request matches """ async with ResourceRequestHandler( - self.page, resource_type=resource_type + self.page, resource_type=resource_type, timeout=timeout ) as handler: await clickable.click() diff --git a/sdk/harambe/handlers.py b/sdk/harambe/handlers.py index 81b7897..ea67896 100644 --- a/sdk/harambe/handlers.py +++ b/sdk/harambe/handlers.py @@ -38,12 +38,13 @@ def __init__( self, page: Page, resource_type: ResourceType, + timeout: int, url_pattern: str = "**/*", ): self.page = page self.url_pattern = url_pattern self.resource_type = resource_type - + self.timeout = timeout self._initial_pages = [p.url for p in page.context.pages] self._new_pages: list[str] = [] @@ -61,11 +62,13 @@ async def __aexit__(self, *_: Any, **__: Any) -> None: self._new_pages.append(page.url) await page.close() except TimeoutError: - raise TimeoutError("No new page opened within the timeout.") + raise TimeoutError( + f"No new page opened within the {self.timeout} seconds timeout." + ) - async def _wait_for_new_page(self, timeout: int = 10) -> Page | None: + async def _wait_for_new_page(self) -> Page | None: start_time = time.monotonic() - while time.monotonic() - start_time < timeout: + while time.monotonic() - start_time < self.timeout: current_pages = self.page.context.pages for page in current_pages: if page.url not in self._initial_pages: