From 3dbd855567189466ecdb05f29acb106b1607d533 Mon Sep 17 00:00:00 2001 From: Dominik Klotz Date: Sat, 1 Nov 2025 15:04:20 +0100 Subject: [PATCH 1/3] feat: implement wait until --- src/askui/agent_base.py | 90 ++++++++-- src/askui/models/exceptions.py | 22 ++- tests/e2e/agent/test_wait.py | 291 +++++++++++++++++++++++++++++++++ 3 files changed, 391 insertions(+), 12 deletions(-) create mode 100644 tests/e2e/agent/test_wait.py diff --git a/src/askui/agent_base.py b/src/askui/agent_base.py index 2483e6a8..9e8186c0 100644 --- a/src/askui/agent_base.py +++ b/src/askui/agent_base.py @@ -1,3 +1,4 @@ +import dis import logging import time import types @@ -5,6 +6,7 @@ from typing import Annotated, Literal, Optional, Type, overload from dotenv import load_dotenv +from exceptiongroup import catch from pydantic import ConfigDict, Field, field_validator, validate_call from pydantic_settings import BaseSettings, SettingsConfigDict from typing_extensions import Self @@ -22,7 +24,7 @@ from askui.utils.source_utils import InputSource, load_image_source from .models import ModelComposition -from .models.exceptions import ElementNotFoundError +from .models.exceptions import ElementNotFoundError, WaitUntilError from .models.model_router import ModelRouter, initialize_default_model_registry from .models.models import ( ModelChoice, @@ -399,6 +401,7 @@ def _locate( self, locator: str | Locator, screenshot: Optional[InputSource] = None, + retry: Optional[Retry] = None, model: ModelComposition | str | None = None, ) -> PointList: def locate_with_screenshot() -> PointList: @@ -410,8 +413,8 @@ def locate_with_screenshot() -> PointList: locator=locator, model=self._get_model(model, "locate"), ) - - points = self._retry.attempt(locate_with_screenshot) + retry = retry or self._retry + points = retry.attempt(locate_with_screenshot) self._reporter.add_message("ModelRouter", f"locate {len(points)} elements") logger.debug("ModelRouter locate: %d elements", len(points)) return points @@ -454,7 +457,7 @@ def locate( "VisionAgent received instruction to locate first matching element %s", locator, ) - return self._locate(locator, screenshot, model)[0] + return self._locate(locator=locator, screenshot=screenshot, model=model)[0] @telemetry.record_call(exclude={"locator", "screenshot"}) @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) @@ -497,30 +500,95 @@ def locate_all( "VisionAgent received instruction to locate all matching UI elements %s", locator, ) - return self._locate(locator, screenshot, model) + return self._locate(locator=locator, screenshot=screenshot, model=model) - @telemetry.record_call() - @validate_call + @telemetry.record_call(exclude={"locator"}) + @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def wait( self, - sec: Annotated[float, Field(gt=0.0)], + until: Annotated[float, Field(gt=0.0)] | str | Locator, + retry_count: Optional[int] = None, + delay: Optional[int] = None, + until_condition: Literal["appear", "disappear"] = "appear", + model: ModelComposition | str | None = None, ) -> None: """ - Pauses the execution of the program for the specified number of seconds. + Pauses execution or waits until a UI element appears or disappears on the screen. Args: - sec (float): The number of seconds to wait. Must be greater than `0.0`. + until (float | str | Locator): If a float, pauses execution for the + specified number of seconds (must be greater than 0.0). If a string + or Locator, waits until the specified UI element appears or disappears on screen. + retry_count (int | None): Number of retries when waiting for a UI element. + Defaults to 3 if None. + delay (int | None): Sleep duration in seconds between retries when waiting + for a UI element. Defaults to 1 second if None. + until_condition (Literal["appear", "disappear"]): The condition to wait until + the element satisfies. Defaults to "appear". + model (ModelComposition | str | None, optional): The composition or name + of the model(s) to be used for locating the element using the `until` locator. + + Raises: + WaitUntilError: If the UI element is not found after all retries. Example: ```python from askui import VisionAgent + from askui.locators import loc with VisionAgent() as agent: + # Wait for a specific duration agent.wait(5) # Pauses execution for 5 seconds agent.wait(0.5) # Pauses execution for 500 milliseconds + + # Wait for a UI element to appear + agent.wait("Submit button", retry_count=5, delay=2) + agent.wait("Login form") # Uses default retries and sleep time + agent.wait(loc.Text("Password")) # Uses default retries and sleep time + + # Wait for a UI element to disappear + agent.wait("Loading spinner", until_condition="disappear") + + # Wait using a specific model + agent.wait("Submit button", model="custom_model") ``` """ - time.sleep(sec) + if isinstance(until, float) or isinstance(until, int): + time.sleep(until) + return + + retry_count = retry_count if retry_count is not None else 3 + delay = delay if delay is not None else 1 + if until_condition == "appear": + try: + self._locate(until, model=model, + retry=ConfigurableRetry( + strategy="Fixed", + base_delay=delay*1000, + retry_count=retry_count, + on_exception_types=(ElementNotFoundError,) + )) + return + except ElementNotFoundError as e: + raise WaitUntilError(e.locator, e.locator_serialized, retry_count, delay, until_condition) from e + else: + for i in range(retry_count): + try: + self._locate(until, model=model, + retry=ConfigurableRetry( + strategy="Fixed", + base_delay=delay*1000, + retry_count=1, + on_exception_types=(ElementNotFoundError,) + )) + logger.debug( + "Element still present, retrying... %d/%d", i + 1, retry_count + ) + time.sleep(delay) + except ElementNotFoundError: + return + raise WaitUntilError(until, str(until), retry_count, delay, until_condition) + @telemetry.record_call() def close(self) -> None: diff --git a/src/askui/models/exceptions.py b/src/askui/models/exceptions.py index 6afbec8a..3692d6e1 100644 --- a/src/askui/models/exceptions.py +++ b/src/askui/models/exceptions.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Literal from askui.locators.locators import Locator @@ -43,6 +43,26 @@ def __init__(self, locator: str | Locator, locator_serialized: Any) -> None: super().__init__(f"Element not found: {self.locator}") +class WaitUntilError(AutomationError): + """Exception raised when an element cannot be located within the given time. + + Args: + locator (str | Locator): The locator that was used. + locator_serialized (Any): The locator serialized for the specific model + """ + + def __init__(self, locator: str | Locator, locator_serialized: Any, retry_count: int, + delay: float, + until_condition: Literal["appear", "disappear"]) -> None: + self.locator = locator + self.locator_serialized = locator_serialized + self.retry_count = retry_count + self.delay = delay + self.until_condition = until_condition + + super().__init__(f"Wait until condition '{self.until_condition}' not met for locator: {self.locator} " + f"after {self.retry_count} retries with {self.delay} seconds delay") + class QueryUnexpectedResponseError(AutomationError): """Exception raised when a query returns an unexpected response. diff --git a/tests/e2e/agent/test_wait.py b/tests/e2e/agent/test_wait.py new file mode 100644 index 00000000..4954bdf7 --- /dev/null +++ b/tests/e2e/agent/test_wait.py @@ -0,0 +1,291 @@ +"""Tests for VisionAgent wait functionality.""" + +import time +from typing import TYPE_CHECKING +from unittest.mock import Mock, patch + +import pytest +from PIL import Image as PILImage + +from askui.agent import VisionAgent +from askui.locators import Element +from askui.models import ModelName +from askui.models.exceptions import WaitUntilError + + + +@pytest.mark.parametrize( + "model", + [ + ModelName.ASKUI, + ], +) +class TestVisionAgentWait: + """Test class for VisionAgent wait functionality.""" + + def test_wait_duration_float(self, vision_agent: VisionAgent, model: str) -> None: + """Test waiting for a specific duration (float).""" + start_time = time.time() + wait_duration = 0.5 + + vision_agent.wait(wait_duration) + + elapsed_time = time.time() - start_time + # Allow small tolerance for timing + assert elapsed_time >= wait_duration - 0.1 + assert elapsed_time <= wait_duration + 0.2 + + def test_wait_duration_int(self, vision_agent: VisionAgent, model: str) -> None: + """Test waiting for a specific duration (int).""" + start_time = time.time() + wait_duration = 1 + + vision_agent.wait(wait_duration) + + elapsed_time = time.time() - start_time + # Allow small tolerance for timing + assert elapsed_time >= wait_duration - 0.1 + assert elapsed_time <= wait_duration + 0.2 + + def test_wait_for_element_appear_success( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting for an element to appear (successful case).""" + locator = "Forgot password?" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should not raise an exception since element exists + vision_agent.wait(locator, retries=3, delay=1, until_condition="appear", model=model) + + # Verify screenshot was called + mock_screenshot.assert_called() + + def test_wait_for_element_appear_failure( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting for an element to appear (failure case).""" + locator = "Non-existent element" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should raise WaitUntilError since element doesn't exist + with pytest.raises(WaitUntilError) as exc_info: + vision_agent.wait(locator, retries=2, delay=1, until_condition="appear", model=model) + + assert "appear" in str(exc_info.value) + mock_screenshot.assert_called() + + def test_wait_for_element_disappear_success( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting for an element to disappear (successful case).""" + locator = "Non-existent element" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should not raise an exception since element doesn't exist (already "disappeared") + vision_agent.wait(locator, retries=2, delay=1, until_condition="disappear", model=model) + + mock_screenshot.assert_called() + + def test_wait_for_element_disappear_failure( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting for an element to disappear (failure case).""" + locator = "Forgot password?" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should raise WaitUntilError since element exists and won't disappear + with pytest.raises(WaitUntilError) as exc_info: + vision_agent.wait(locator, retries=2, delay=1, until_condition="disappear", model=model) + + assert "disappear" in str(exc_info.value) + mock_screenshot.assert_called() + + def test_wait_with_locator_object( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting with a Locator object.""" + locator = Element("textfield") + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should not raise an exception since textfield exists + vision_agent.wait(locator, retries=2, delay=1, until_condition="appear", model=model) + + mock_screenshot.assert_called() + + def test_wait_with_default_parameters( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting with default parameters.""" + locator = "Forgot password?" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should use default retries=3, delay=1, until_condition="appear" + vision_agent.wait(locator, model=model) + + mock_screenshot.assert_called() + + def test_wait_with_custom_retries_and_delay( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting with custom retries and delay values.""" + locator = "Forgot password?" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + start_time = time.time() + vision_agent.wait(locator, retries=1, delay=2, until_condition="appear", model=model) + elapsed_time = time.time() - start_time + + # Should complete quickly since element exists on first try + assert elapsed_time < 1.0 + mock_screenshot.assert_called() + + def test_wait_with_custom_model( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting with a custom model parameter.""" + locator = "Forgot password?" + custom_model = "custom_model_name" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should not raise an exception + vision_agent.wait( + locator, + retries=1, + delay=1, + until_condition="appear", + model=custom_model + ) + + mock_screenshot.assert_called() + + @patch("time.sleep") + def test_wait_disappear_timing( + self, + mock_sleep: Mock, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test that wait for disappear calls sleep with correct delay.""" + locator = "Forgot password?" + delay = 2 + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should raise WaitUntilError and call sleep with correct delay + with pytest.raises(WaitUntilError): + vision_agent.wait( + locator, + retries=2, + delay=delay, + until_condition="disappear", + model=model + ) + + # Verify sleep was called with the correct delay + expected_calls = delay * (2 - 1) # (retries - 1) * delay + assert mock_sleep.call_count == 1 + mock_sleep.assert_called_with(delay) + + def test_wait_until_error_contains_correct_info( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test that WaitUntilError contains correct information.""" + locator = "Non-existent element" + retries = 3 + delay = 1 + until_condition = "appear" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + with pytest.raises(WaitUntilError) as exc_info: + vision_agent.wait( + locator, + retries=retries, + delay=delay, + until_condition=until_condition, + model=model + ) + + error = exc_info.value + # Verify error contains the expected information + assert locator in str(error) + assert until_condition in str(error) + + def test_wait_zero_retries( + self, + vision_agent: VisionAgent, + github_login_screenshot: PILImage.Image, + model: str, + ) -> None: + """Test waiting with zero retries.""" + locator = "Non-existent element" + + # Mock screenshot to return the image + mock_screenshot = Mock(return_value=github_login_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + + # Should fail immediately with 0 retries + with pytest.raises(WaitUntilError): + vision_agent.wait( + locator, + retries=0, + delay=1, + until_condition="appear", + model=model + ) From 6e5e88774158c19fa794cf215ad561ea89546e9e Mon Sep 17 00:00:00 2001 From: Dominik Klotz Date: Sun, 2 Nov 2025 16:51:07 +0100 Subject: [PATCH 2/3] feat: implement wait until --- src/askui/agent_base.py | 123 +++++++--- src/askui/models/exceptions.py | 19 +- tests/conftest.py | 7 + tests/e2e/agent/test_wait.py | 283 ++++++++-------------- tests/fixtures/screenshots/white_page.png | Bin 0 -> 8031 bytes 5 files changed, 207 insertions(+), 225 deletions(-) create mode 100644 tests/fixtures/screenshots/white_page.png diff --git a/src/askui/agent_base.py b/src/askui/agent_base.py index 9e8186c0..03b7d85e 100644 --- a/src/askui/agent_base.py +++ b/src/askui/agent_base.py @@ -1,4 +1,3 @@ -import dis import logging import time import types @@ -6,7 +5,6 @@ from typing import Annotated, Literal, Optional, Type, overload from dotenv import load_dotenv -from exceptiongroup import catch from pydantic import ConfigDict, Field, field_validator, validate_call from pydantic_settings import BaseSettings, SettingsConfigDict from typing_extensions import Self @@ -413,6 +411,7 @@ def locate_with_screenshot() -> PointList: locator=locator, model=self._get_model(model, "locate"), ) + retry = retry or self._retry points = retry.attempt(locate_with_screenshot) self._reporter.add_message("ModelRouter", f"locate {len(points)} elements") @@ -507,26 +506,28 @@ def locate_all( def wait( self, until: Annotated[float, Field(gt=0.0)] | str | Locator, - retry_count: Optional[int] = None, - delay: Optional[int] = None, + retry_count: Optional[Annotated[int, Field(gt=0)]] = None, + delay: Optional[Annotated[float, Field(gt=0.0)]] = None, until_condition: Literal["appear", "disappear"] = "appear", model: ModelComposition | str | None = None, ) -> None: """ - Pauses execution or waits until a UI element appears or disappears on the screen. + Pauses execution or waits until a UI element appears or disappears. Args: until (float | str | Locator): If a float, pauses execution for the specified number of seconds (must be greater than 0.0). If a string - or Locator, waits until the specified UI element appears or disappears on screen. - retry_count (int | None): Number of retries when waiting for a UI element. - Defaults to 3 if None. - delay (int | None): Sleep duration in seconds between retries when waiting - for a UI element. Defaults to 1 second if None. - until_condition (Literal["appear", "disappear"]): The condition to wait until - the element satisfies. Defaults to "appear". + or Locator, waits until the specified UI element appears or + disappears on screen. + retry_count (int | None): Number of retries when waiting for a UI + element. Defaults to 3 if None. + delay (int | None): Sleep duration in seconds between retries when + waiting for a UI element. Defaults to 1 second if None. + until_condition (Literal["appear", "disappear"]): The condition to wait + until the element satisfies. Defaults to "appear". model (ModelComposition | str | None, optional): The composition or name - of the model(s) to be used for locating the element using the `until` locator. + of the model(s) to be used for locating the element using the + `until` locator. Raises: WaitUntilError: If the UI element is not found after all retries. @@ -554,41 +555,87 @@ def wait( ``` """ if isinstance(until, float) or isinstance(until, int): + self._reporter.add_message("User", f"wait {until} seconds") time.sleep(until) return + self._reporter.add_message( + "User", f"wait for element '{until}' to {until_condition}" + ) retry_count = retry_count if retry_count is not None else 3 delay = delay if delay is not None else 1 + if until_condition == "appear": - try: - self._locate(until, model=model, - retry=ConfigurableRetry( + self._wait_for_appear(until, model, retry_count, delay) + else: + self._wait_for_disappear(until, model, retry_count, delay) + + def _wait_for_appear( + self, + locator: str | Locator, + model: ModelComposition | str | None, + retry_count: int, + delay: float, + ) -> None: + """Wait for an element to appear on screen.""" + try: + self._locate( + locator, + model=model, + retry=ConfigurableRetry( strategy="Fixed", - base_delay=delay*1000, + base_delay=int(delay * 1000), retry_count=retry_count, - on_exception_types=(ElementNotFoundError,) - )) + on_exception_types=(ElementNotFoundError,), + ), + ) + self._reporter.add_message( + "VisionAgent", f"element '{locator}' appeared successfully" + ) + except ElementNotFoundError as e: + self._reporter.add_message( + "VisionAgent", + f"element '{locator}' failed to appear after {retry_count} retries", + ) + raise WaitUntilError( + e.locator, e.locator_serialized, retry_count, delay, "appear" + ) from e + + def _wait_for_disappear( + self, + locator: str | Locator, + model: ModelComposition | str | None, + retry_count: int, + delay: float, + ) -> None: + """Wait for an element to disappear from screen.""" + for i in range(retry_count): + try: + self._locate( + locator, + model=model, + retry=ConfigurableRetry( + strategy="Fixed", + base_delay=int(delay * 1000), + retry_count=1, + on_exception_types=(), + ), + ) + logger.debug( + "Element still present, retrying... %d/%d", i + 1, retry_count + ) + time.sleep(delay) + except ElementNotFoundError: # noqa: PERF203 + self._reporter.add_message( + "VisionAgent", f"element '{locator}' disappeared successfully" + ) return - except ElementNotFoundError as e: - raise WaitUntilError(e.locator, e.locator_serialized, retry_count, delay, until_condition) from e - else: - for i in range(retry_count): - try: - self._locate(until, model=model, - retry=ConfigurableRetry( - strategy="Fixed", - base_delay=delay*1000, - retry_count=1, - on_exception_types=(ElementNotFoundError,) - )) - logger.debug( - "Element still present, retrying... %d/%d", i + 1, retry_count - ) - time.sleep(delay) - except ElementNotFoundError: - return - raise WaitUntilError(until, str(until), retry_count, delay, until_condition) + self._reporter.add_message( + "VisionAgent", + f"element '{locator}' failed to disappear after {retry_count} retries", + ) + raise WaitUntilError(locator, str(locator), retry_count, delay, "disappear") @telemetry.record_call() def close(self) -> None: diff --git a/src/askui/models/exceptions.py b/src/askui/models/exceptions.py index 3692d6e1..ae7febed 100644 --- a/src/askui/models/exceptions.py +++ b/src/askui/models/exceptions.py @@ -51,17 +51,26 @@ class WaitUntilError(AutomationError): locator_serialized (Any): The locator serialized for the specific model """ - def __init__(self, locator: str | Locator, locator_serialized: Any, retry_count: int, - delay: float, - until_condition: Literal["appear", "disappear"]) -> None: + def __init__( + self, + locator: str | Locator, + locator_serialized: Any, + retry_count: int, + delay: float, + until_condition: Literal["appear", "disappear"], + ) -> None: self.locator = locator self.locator_serialized = locator_serialized self.retry_count = retry_count self.delay = delay self.until_condition = until_condition - super().__init__(f"Wait until condition '{self.until_condition}' not met for locator: {self.locator} " - f"after {self.retry_count} retries with {self.delay} seconds delay") + super().__init__( + f"Wait until condition '{self.until_condition}' not met" + f" for locator: '{self.locator}' after {self.retry_count} retries" + f" with {self.delay} seconds delay" + ) + class QueryUnexpectedResponseError(AutomationError): """Exception raised when a query returns an unexpected response. diff --git a/tests/conftest.py b/tests/conftest.py index 6e386721..0fde4951 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -71,6 +71,13 @@ def github_login_screenshot(path_fixtures_screenshots: pathlib.Path) -> Image.Im return Image.open(screenshot_path) +@pytest.fixture +def white_page_screenshot(path_fixtures_screenshots: pathlib.Path) -> Image.Image: + """Fixture providing the white page screenshot.""" + screenshot_path = path_fixtures_screenshots / "white_page.png" + return Image.open(screenshot_path) + + @pytest.fixture def path_fixtures_github_com__icon(path_fixtures_images: pathlib.Path) -> pathlib.Path: """Fixture providing the path to the github com icon image.""" diff --git a/tests/e2e/agent/test_wait.py b/tests/e2e/agent/test_wait.py index 4954bdf7..45394308 100644 --- a/tests/e2e/agent/test_wait.py +++ b/tests/e2e/agent/test_wait.py @@ -1,29 +1,22 @@ +# mypy: disable-error-code="method-assign" """Tests for VisionAgent wait functionality.""" import time -from typing import TYPE_CHECKING -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest from PIL import Image as PILImage +from pydantic import ValidationError from askui.agent import VisionAgent from askui.locators import Element -from askui.models import ModelName from askui.models.exceptions import WaitUntilError - -@pytest.mark.parametrize( - "model", - [ - ModelName.ASKUI, - ], -) class TestVisionAgentWait: """Test class for VisionAgent wait functionality.""" - def test_wait_duration_float(self, vision_agent: VisionAgent, model: str) -> None: + def test_wait_duration_float(self, vision_agent: VisionAgent) -> None: """Test waiting for a specific duration (float).""" start_time = time.time() wait_duration = 0.5 @@ -35,7 +28,7 @@ def test_wait_duration_float(self, vision_agent: VisionAgent, model: str) -> Non assert elapsed_time >= wait_duration - 0.1 assert elapsed_time <= wait_duration + 0.2 - def test_wait_duration_int(self, vision_agent: VisionAgent, model: str) -> None: + def test_wait_duration_int(self, vision_agent: VisionAgent) -> None: """Test waiting for a specific duration (int).""" start_time = time.time() wait_duration = 1 @@ -51,241 +44,167 @@ def test_wait_for_element_appear_success( self, vision_agent: VisionAgent, github_login_screenshot: PILImage.Image, - model: str, + white_page_screenshot: PILImage.Image, ) -> None: """Test waiting for an element to appear (successful case).""" locator = "Forgot password?" - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + # Mock screenshot to return the image where element exists + mock_screenshot = Mock( + side_effect=[white_page_screenshot, github_login_screenshot] + ) + vision_agent._agent_os.screenshot = mock_screenshot + vision_agent._model_router.locate = Mock( + wraps=vision_agent._model_router.locate + ) # Should not raise an exception since element exists - vision_agent.wait(locator, retries=3, delay=1, until_condition="appear", model=model) + vision_agent.wait(locator, retry_count=3, delay=0.1, until_condition="appear") - # Verify screenshot was called - mock_screenshot.assert_called() + # Verify locate was called + assert vision_agent._model_router.locate.call_count == 2 def test_wait_for_element_appear_failure( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, + self, vision_agent: VisionAgent, white_page_screenshot: PILImage.Image ) -> None: """Test waiting for an element to appear (failure case).""" - locator = "Non-existent element" + locator = "Forgot password?" - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + # Mock screenshot to return white page where element doesn't exist + mock_screenshot = Mock(return_value=white_page_screenshot) + vision_agent._agent_os.screenshot = mock_screenshot + vision_agent._model_router.locate = Mock( + wraps=vision_agent._model_router.locate + ) # Should raise WaitUntilError since element doesn't exist with pytest.raises(WaitUntilError) as exc_info: - vision_agent.wait(locator, retries=2, delay=1, until_condition="appear", model=model) + vision_agent.wait( + locator, retry_count=2, delay=0.1, until_condition="appear" + ) - assert "appear" in str(exc_info.value) - mock_screenshot.assert_called() + assert ( + "Wait until condition 'appear' not met for locator: 'text similar to " + '"Forgot password?" (similarity >= 70%)\' after 2 retries with 0.1 ' + "seconds delay" == str(exc_info.value) + ) + assert vision_agent._model_router.locate.call_count == 2 def test_wait_for_element_disappear_success( self, vision_agent: VisionAgent, github_login_screenshot: PILImage.Image, - model: str, + white_page_screenshot: PILImage.Image, ) -> None: """Test waiting for an element to disappear (successful case).""" - locator = "Non-existent element" + locator = "Forgot password?" - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + # Mock screenshot to first show element exists, then show it's gone + mock_screenshot = Mock( + side_effect=[github_login_screenshot, white_page_screenshot] + ) + vision_agent._agent_os.screenshot = mock_screenshot + vision_agent._model_router.locate = Mock( + wraps=vision_agent._model_router.locate + ) - # Should not raise an exception since element doesn't exist (already "disappeared") - vision_agent.wait(locator, retries=2, delay=1, until_condition="disappear", model=model) + # Should not raise an exception since element disappears + vision_agent.wait( + locator, retry_count=2, delay=0.1, until_condition="disappear" + ) - mock_screenshot.assert_called() + assert vision_agent._model_router.locate.call_count == 2 def test_wait_for_element_disappear_failure( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, + self, vision_agent: VisionAgent, github_login_screenshot: PILImage.Image ) -> None: """Test waiting for an element to disappear (failure case).""" locator = "Forgot password?" - # Mock screenshot to return the image + # Mock screenshot to always show the element exists mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + vision_agent._agent_os.screenshot = mock_screenshot + vision_agent._model_router.locate = Mock( + wraps=vision_agent._model_router.locate + ) # Should raise WaitUntilError since element exists and won't disappear with pytest.raises(WaitUntilError) as exc_info: - vision_agent.wait(locator, retries=2, delay=1, until_condition="disappear", model=model) + vision_agent.wait( + locator, retry_count=2, delay=0.1, until_condition="disappear" + ) - assert "disappear" in str(exc_info.value) - mock_screenshot.assert_called() + assert ( + "Wait until condition 'disappear' not met for locator: 'Forgot password?' " + "after 2 retries with 0.1 seconds delay" == str(exc_info.value) + ) + assert vision_agent._model_router.locate.call_count == 2 def test_wait_with_locator_object( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, + self, vision_agent: VisionAgent, github_login_screenshot: PILImage.Image ) -> None: """Test waiting with a Locator object.""" locator = Element("textfield") - # Mock screenshot to return the image + # Mock screenshot to return the image where textfield exists mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + vision_agent._agent_os.screenshot = mock_screenshot + vision_agent._model_router.locate = Mock( + wraps=vision_agent._model_router.locate + ) # Should not raise an exception since textfield exists - vision_agent.wait(locator, retries=2, delay=1, until_condition="appear", model=model) + vision_agent.wait(locator, retry_count=2, delay=0.1, until_condition="appear") - mock_screenshot.assert_called() + assert vision_agent._model_router.locate.call_count >= 1 def test_wait_with_default_parameters( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, + self, vision_agent: VisionAgent, github_login_screenshot: PILImage.Image ) -> None: """Test waiting with default parameters.""" locator = "Forgot password?" - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] - - # Should use default retries=3, delay=1, until_condition="appear" - vision_agent.wait(locator, model=model) - - mock_screenshot.assert_called() - - def test_wait_with_custom_retries_and_delay( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, - ) -> None: - """Test waiting with custom retries and delay values.""" - locator = "Forgot password?" - - # Mock screenshot to return the image + # Mock screenshot to return the image where element exists mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] - - start_time = time.time() - vision_agent.wait(locator, retries=1, delay=2, until_condition="appear", model=model) - elapsed_time = time.time() - start_time - - # Should complete quickly since element exists on first try - assert elapsed_time < 1.0 - mock_screenshot.assert_called() - - def test_wait_with_custom_model( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, - ) -> None: - """Test waiting with a custom model parameter.""" - locator = "Forgot password?" - custom_model = "custom_model_name" - - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] - - # Should not raise an exception - vision_agent.wait( - locator, - retries=1, - delay=1, - until_condition="appear", - model=custom_model + vision_agent._agent_os.screenshot = mock_screenshot + vision_agent._model_router.locate = Mock( + wraps=vision_agent._model_router.locate ) - mock_screenshot.assert_called() + # Should use default retry_count=3, delay=1, until_condition="appear" + vision_agent.wait(locator) + + assert vision_agent._model_router.locate.call_count >= 1 - @patch("time.sleep") def test_wait_disappear_timing( self, - mock_sleep: Mock, vision_agent: VisionAgent, github_login_screenshot: PILImage.Image, - model: str, + white_page_screenshot: PILImage.Image, ) -> None: """Test that wait for disappear calls sleep with correct delay.""" locator = "Forgot password?" - delay = 2 - - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] - # Should raise WaitUntilError and call sleep with correct delay - with pytest.raises(WaitUntilError): - vision_agent.wait( - locator, - retries=2, - delay=delay, - until_condition="disappear", - model=model - ) - - # Verify sleep was called with the correct delay - expected_calls = delay * (2 - 1) # (retries - 1) * delay - assert mock_sleep.call_count == 1 - mock_sleep.assert_called_with(delay) - - def test_wait_until_error_contains_correct_info( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, - ) -> None: - """Test that WaitUntilError contains correct information.""" - locator = "Non-existent element" - retries = 3 - delay = 1 - until_condition = "appear" - - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] - - with pytest.raises(WaitUntilError) as exc_info: - vision_agent.wait( - locator, - retries=retries, - delay=delay, - until_condition=until_condition, - model=model - ) - - error = exc_info.value - # Verify error contains the expected information - assert locator in str(error) - assert until_condition in str(error) + mock_screenshot = Mock( + side_effect=[ + github_login_screenshot, + github_login_screenshot, + white_page_screenshot, + ] + ) + vision_agent._agent_os.screenshot = mock_screenshot + vision_agent._model_router.locate = Mock( + wraps=vision_agent._model_router.locate + ) - def test_wait_zero_retries( - self, - vision_agent: VisionAgent, - github_login_screenshot: PILImage.Image, - model: str, - ) -> None: - """Test waiting with zero retries.""" - locator = "Non-existent element" + vision_agent.wait( + locator, until_condition="disappear", retry_count=3, delay=0.2 + ) - # Mock screenshot to return the image - mock_screenshot = Mock(return_value=github_login_screenshot) - vision_agent._agent_os.screenshot = mock_screenshot # type: ignore[method-assign] + assert vision_agent._model_router.locate.call_count == 3 - # Should fail immediately with 0 retries - with pytest.raises(WaitUntilError): - vision_agent.wait( - locator, - retries=0, - delay=1, - until_condition="appear", - model=model - ) + def test_wait_zero_retries(self, vision_agent: VisionAgent) -> None: + """Test waiting with zero retry_count.""" + # Should fail immediately with 0 retry_count + with pytest.raises(ValidationError): + vision_agent.wait("test", retry_count=0) diff --git a/tests/fixtures/screenshots/white_page.png b/tests/fixtures/screenshots/white_page.png new file mode 100644 index 0000000000000000000000000000000000000000..884d49fd7b22a30e976c5e99dccc234e57b85cee GIT binary patch literal 8031 zcmeAS@N?(olHy`uVBq!ia0y~yU@c}~VE)a)1{9H$m+b*kjKx9jP7LeL$-D$|SkfJR z9T^xl_H+M9WCij$3p^r$I>b~#n9(yxG+w>J!V84P$1Z1|+U_pj>( zk8hG3GueT97!Fi!W(Ls=-)HOu(~N0)Alji}CO3!{NO;BsqFF|zLBTK@EThQ*lnzF7 z$!O^SOa`N+;%H?7Oa`Ns`5fql5u6{1-oD!M Date: Mon, 3 Nov 2025 16:46:08 +0100 Subject: [PATCH 3/3] Change telemetry exclusion from 'locator' to 'until' --- src/askui/agent_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/askui/agent_base.py b/src/askui/agent_base.py index 03b7d85e..a6a1dfb7 100644 --- a/src/askui/agent_base.py +++ b/src/askui/agent_base.py @@ -501,7 +501,7 @@ def locate_all( ) return self._locate(locator=locator, screenshot=screenshot, model=model) - @telemetry.record_call(exclude={"locator"}) + @telemetry.record_call(exclude={"until"}) @validate_call(config=ConfigDict(arbitrary_types_allowed=True)) def wait( self,