From 4f2bc784f5b858d65b306c5b0a0116d71ec4a612 Mon Sep 17 00:00:00 2001 From: Tony Li Date: Wed, 8 Jul 2026 10:24:49 -0700 Subject: [PATCH] Refactor android_env wrappers to use datetime.timedelta for timing. PiperOrigin-RevId: 944570797 --- android_env/components/adb_call_parser.py | 79 ++++++++++++++----- .../components/adb_call_parser_test.py | 19 +++++ android_env/components/adb_controller.py | 14 ++-- android_env/components/adb_controller_test.py | 27 ++++--- android_env/components/app_screen_checker.py | 35 +++++--- .../components/app_screen_checker_test.py | 8 +- android_env/components/config_classes.py | 5 +- .../components/setup_step_interpreter.py | 8 +- .../components/simulators/base_simulator.py | 17 ++-- .../simulators/base_simulator_test.py | 13 ++- .../simulators/emulator/emulator_simulator.py | 24 ++++-- .../emulator/emulator_simulator_test.py | 3 +- .../simulators/fake/fake_simulator.py | 3 +- android_env/wrappers/a11y_grpc_wrapper.py | 32 +++++--- .../wrappers/a11y_grpc_wrapper_test.py | 17 +++- android_env/wrappers/rate_limit_wrapper.py | 36 +++++---- .../wrappers/rate_limit_wrapper_test.py | 35 ++++---- 17 files changed, 250 insertions(+), 125 deletions(-) diff --git a/android_env/components/adb_call_parser.py b/android_env/components/adb_call_parser.py index c38bf3fc..ee401e2c 100644 --- a/android_env/components/adb_call_parser.py +++ b/android_env/components/adb_call_parser.py @@ -15,6 +15,7 @@ """Processes adb_pb2.AdbRequest commands.""" +import datetime import os import re import subprocess @@ -63,7 +64,7 @@ def __init__(self, adb_controller: adb_control.AdbController): } def _execute_command( - self, command_args: list[str], timeout: float | None + self, command_args: list[str], timeout: datetime.timedelta | None ) -> tuple[adb_pb2.AdbResponse, bytes]: """Executes the command, catches errors and populates the response status. @@ -107,11 +108,17 @@ def parse(self, request: adb_pb2.AdbRequest) -> adb_pb2.AdbResponse: f'Got: {request.timeout_sec}') return response - timeout: float | None = request.timeout_sec or None + timeout = ( + datetime.timedelta(seconds=request.timeout_sec) + if request.timeout_sec > 0 + else None + ) return self._handlers[command_type](request, timeout) def _force_stop( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Stops an application. @@ -136,7 +143,7 @@ def _force_stop( return response def _fetch_current_task_id( - self, full_activity_name: str, timeout: float | None = None + self, full_activity_name: str, timeout: datetime.timedelta | None = None ) -> int: """Returns the task ID of the given `full_activity_name`. @@ -183,7 +190,9 @@ def _fetch_current_task_id( return -1 def _start_screen_pinning( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Pins an application. @@ -217,7 +226,9 @@ def _start_screen_pinning( return response def _send_broadcast( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Sends a broadcast. @@ -250,7 +261,9 @@ def _send_broadcast( return response def _install_apk( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Installs an app given its local path in the filesystem. @@ -305,7 +318,9 @@ def _install_apk( return response def _start_activity( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Starts a given activity. @@ -353,7 +368,9 @@ def _start_activity( return response def _press_button( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Presses a keyboard key. @@ -380,7 +397,9 @@ def _press_button( return response def _handle_uninstall_package( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Handles UninstallPackage messages. @@ -419,7 +438,9 @@ def _handle_uninstall_package( return response def _get_current_activity( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Fetches current activity. @@ -473,7 +494,9 @@ def _get_current_activity( return response def _get_orientation( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Fetches current device orientation. @@ -534,7 +557,9 @@ def _get_orientation( return response def _push( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Uploads contents to the device. @@ -565,7 +590,9 @@ def _push( return response def _pull( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Downloads file content from the device. @@ -598,7 +625,9 @@ def _pull( return response def _input_text( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Inserts text as keyboard events. @@ -621,7 +650,9 @@ def _input_text( return response def _tap( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Taps the device screen. @@ -651,7 +682,9 @@ def _tap( return response def _handle_settings( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Handles SettingsRequest messages. @@ -767,7 +800,9 @@ def _handle_settings( ) def _handle_generic( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Handles GenericRequest messages. @@ -786,7 +821,9 @@ def _handle_generic( return response def _handle_package_manager( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Handles PackageManagerRequest messages. @@ -880,7 +917,9 @@ def _handle_package_manager( ) def _handle_dumpsys( - self, request: adb_pb2.AdbRequest, timeout: float | None = None + self, + request: adb_pb2.AdbRequest, + timeout: datetime.timedelta | None = None, ) -> adb_pb2.AdbResponse: """Handles DumpsysRequest messages. diff --git a/android_env/components/adb_call_parser_test.py b/android_env/components/adb_call_parser_test.py index c3f4cfb8..5c62aec4 100644 --- a/android_env/components/adb_call_parser_test.py +++ b/android_env/components/adb_call_parser_test.py @@ -14,6 +14,7 @@ # limitations under the License. import builtins +import datetime import os import subprocess import sys @@ -775,6 +776,24 @@ def test_generic_command(self): self.assertEqual(response.generic.output, expected_output) adb.execute_command.assert_called_once_with(args, None) + def test_generic_command_with_timeout_successful(self): + adb = mock.create_autospec(adb_controller.AdbController) + expected_output = b'broadcast_output' + args = ['shell', 'am', 'broadcast', '-n', 'receiver', '-a', 'action'] + adb.execute_command.return_value = expected_output + parser = adb_call_parser.AdbCallParser(adb) + + generic_request = adb_pb2.AdbRequest.GenericRequest(args=args) + request = adb_pb2.AdbRequest(generic=generic_request, timeout_sec=5.0) + response = parser.parse(request) + + self.assertEqual(adb_pb2.AdbResponse.Status.OK, response.status) + self.assertEmpty(response.error_message) + self.assertEqual(response.generic.output, expected_output) + adb.execute_command.assert_called_once_with( + args, datetime.timedelta(seconds=5.0) + ) + def test_generic_command_adb_error(self): adb = mock.create_autospec(adb_controller.AdbController) args = ['shell', 'am', 'broadcast', '-n', 'receiver', '-a', 'action'] diff --git a/android_env/components/adb_controller.py b/android_env/components/adb_controller.py index 82737571..c68b3329 100644 --- a/android_env/components/adb_controller.py +++ b/android_env/components/adb_controller.py @@ -15,6 +15,7 @@ """A class to manage and control an external ADB process.""" +import datetime import os import subprocess import time @@ -68,7 +69,7 @@ def command_prefix(self, include_device_name: bool = True) -> list[str]: command_prefix.extend(['-s', self._config.device_name]) return command_prefix - def init_server(self, timeout: float | None = None): + def init_server(self, timeout: datetime.timedelta | None = None): """Initialize the ADB server deamon on the given port. This function should be called immediately after initializing the first @@ -82,7 +83,7 @@ def init_server(self, timeout: float | None = None): self.execute_command(['devices'], timeout, device_specific=False) time.sleep(0.2) - def _restart_server(self, timeout: float | None = None): + def _restart_server(self, timeout: datetime.timedelta | None = None): """Kills and restarts the adb server. Args: @@ -105,7 +106,7 @@ def _restart_server(self, timeout: float | None = None): def execute_command( self, args: list[str], - timeout: float | None = None, + timeout: datetime.timedelta | None = None, device_specific: bool = True, ) -> bytes: """Executes an adb command. @@ -141,14 +142,17 @@ def execute_command( ) from latest_error def _execute_attempt( - self, command: list[str], command_str: str, timeout: float + self, + command: list[str], + command_str: str, + timeout: datetime.timedelta | None, ) -> bytes: """Executes a single adb command attempt.""" logging.info('Executing ADB command: [%s]', command_str) cmd_output = subprocess.check_output( command, stderr=subprocess.STDOUT, - timeout=timeout, + timeout=timeout.total_seconds() if timeout is not None else None, env=self._os_env_vars, ) logging.debug('ADB command output: %s', cmd_output) diff --git a/android_env/components/adb_controller_test.py b/android_env/components/adb_controller_test.py index 71a3a4e6..aeb0af2d 100644 --- a/android_env/components/adb_controller_test.py +++ b/android_env/components/adb_controller_test.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime import os import subprocess import time @@ -25,7 +26,7 @@ # Timeout to be used by default in tests below. Set to a small value to avoid # hanging on a failed test. -_TIMEOUT = 2 +_TIMEOUT = datetime.timedelta(seconds=2) class AdbControllerTest(absltest.TestCase): @@ -67,7 +68,7 @@ def test_init_server(self, mock_sleep, mock_check_output): mock_check_output.assert_called_once_with( ['my_adb', 'devices'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ) mock_sleep.assert_called_once() @@ -104,7 +105,7 @@ def test_init_server_with_adb_server_port_from_os_env( mock_check_output.assert_called_once_with( ['my_adb', 'devices'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ) mock_sleep.assert_called_once() @@ -138,31 +139,31 @@ def test_restart_server(self, mock_sleep, mock_check_output): mock.call( ['my_adb', '-s', 'awesome_device', 'my_command'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', 'kill-server'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', 'start-server'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', 'devices'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', '-s', 'awesome_device', 'my_command'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), ]) @@ -211,31 +212,31 @@ def test_invalid_command(self, mock_sleep, mock_check_output): mock.call( ['my_adb', '-s', 'awesome_device', 'my_command'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', 'kill-server'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', 'start-server'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', 'devices'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), mock.call( ['my_adb', '-s', 'awesome_device', 'my_command'], stderr=subprocess.STDOUT, - timeout=_TIMEOUT, + timeout=_TIMEOUT.total_seconds(), env=expected_env, ), ], diff --git a/android_env/components/app_screen_checker.py b/android_env/components/app_screen_checker.py index ada7c33e..f4ebc637 100644 --- a/android_env/components/app_screen_checker.py +++ b/android_env/components/app_screen_checker.py @@ -16,6 +16,7 @@ """Determines if the current app screen matches an expected app screen.""" from collections.abc import Callable, Sequence +import datetime import enum import re import time @@ -233,32 +234,40 @@ def matches_current_app_screen(self) -> enum.IntEnum: return AppScreenChecker.Outcome.SUCCESS - def wait_for_app_screen(self, timeout_sec: float) -> float: + def wait_for_app_screen( + self, timeout: datetime.timedelta + ) -> datetime.timedelta: """Waits for `self._expected_app_screen` to be the current screen. Args: - timeout_sec: Maximum total time to wait for the screen to pop up. + timeout: Maximum total time to wait for the screen to pop up. Returns: - The total amount of time in seconds spent waiting for the screen to pop - up. + The total amount of time spent waiting for the screen to pop up. Raises: errors.WaitForAppScreenError if the screen does not pop up within - `timeout_sec`. + `timeout`. """ logging.info('Waiting for app screen...') - start_time = time.time() - while time.time() - start_time < timeout_sec: + start_time = time.monotonic() + timeout_sec = timeout.total_seconds() + while time.monotonic() - start_time < timeout_sec: if self.matches_current_app_screen() == AppScreenChecker.Outcome.SUCCESS: - wait_time = time.time() - start_time - logging.info('Successfully waited for app screen in %r seconds: [%r]', - wait_time, self._expected_app_screen) + wait_time = datetime.timedelta(seconds=time.monotonic() - start_time) + logging.info( + 'Successfully waited for app screen in %r: [%r]', + wait_time, + self._expected_app_screen, + ) return wait_time time.sleep(0.1) - wait_time = time.time() - start_time - logging.error('Failed to wait for app screen in %r seconds: [%r].', - wait_time, self._expected_app_screen) + wait_time = datetime.timedelta(seconds=time.monotonic() - start_time) + logging.error( + 'Failed to wait for app screen in %r: [%r].', + wait_time, + self._expected_app_screen, + ) raise errors.WaitForAppScreenError() diff --git a/android_env/components/app_screen_checker_test.py b/android_env/components/app_screen_checker_test.py index a11ee0ef..1a29b1f6 100644 --- a/android_env/components/app_screen_checker_test.py +++ b/android_env/components/app_screen_checker_test.py @@ -15,6 +15,7 @@ """Tests for android_env.components.app_screen_checker.""" +import datetime import re import timeit from unittest import mock @@ -249,7 +250,8 @@ def test_wait_for_app_screen_zero_timeout(self): self.assertRaises( errors.WaitForAppScreenError, screen_checker.wait_for_app_screen, - timeout_sec=0.0) + timeout=datetime.timedelta(seconds=0.0), + ) def test_wait_for_app_screen_successful(self): """Ensures that with the right conditions, the app screen should pop up.""" @@ -262,8 +264,8 @@ def test_wait_for_app_screen_successful(self): screen_checker = app_screen_checker.AppScreenChecker( call_parser, app_screen) - timeout = 1.0 - wait_time = screen_checker.wait_for_app_screen(timeout_sec=timeout) + timeout = datetime.timedelta(seconds=1.0) + wait_time = screen_checker.wait_for_app_screen(timeout=timeout) # The call should not generate an exception and the return value should be # less than the timeout given. diff --git a/android_env/components/config_classes.py b/android_env/components/config_classes.py index 528974f5..72461fdc 100644 --- a/android_env/components/config_classes.py +++ b/android_env/components/config_classes.py @@ -16,6 +16,7 @@ """Dataclass definitions used for instantiating AndroidEnv components.""" import dataclasses +import datetime import enum @@ -31,7 +32,7 @@ class AdbControllerConfig: # Port for adb server. adb_server_port: int = 5037 # Default timeout in seconds for internal commands. - default_timeout: float = 120.0 + default_timeout: datetime.timedelta = datetime.timedelta(seconds=120.0) # Name of the device to communicate with. device_name: str = '' # Whether to use adb server port set in OS Environment variables. @@ -84,7 +85,7 @@ class SimulatorConfig: # How often to (asynchronously) grab the screenshot from the simulator. # If <= 0, stepping the environment blocks on fetching the screenshot (the # environment is synchronous). - interaction_rate_sec: float = 0.0 + interaction_rate_sec: datetime.timedelta = datetime.timedelta(seconds=0.0) @enum.unique diff --git a/android_env/components/setup_step_interpreter.py b/android_env/components/setup_step_interpreter.py index eb80f7bb..10d27a64 100644 --- a/android_env/components/setup_step_interpreter.py +++ b/android_env/components/setup_step_interpreter.py @@ -17,6 +17,7 @@ from collections.abc import Sequence import copy +import datetime import time from typing import Any @@ -43,7 +44,7 @@ def __init__(self, adb_call_parser: adb_call_parser_lib.AdbCallParser): 'error_count_wait_for_app_screen': 0, 'error_count_check_install': 0, 'error_count_wait_for_message': 0, - 'total_time_waiting_for_app_screen': 0.0, + 'total_time_waiting_for_app_screen': datetime.timedelta(seconds=0.0), } def stats(self) -> dict[str, Any]: @@ -149,9 +150,8 @@ def _check_success( adb_call_parser=self._adb_call_parser, expected_app_screen=wait_for_app_screen.app_screen, ) - wait_time = screen_checker.wait_for_app_screen( - timeout_sec=wait_for_app_screen.timeout_sec - ) + timeout = datetime.timedelta(seconds=wait_for_app_screen.timeout_sec) + wait_time = screen_checker.wait_for_app_screen(timeout=timeout) self._stats['total_time_waiting_for_app_screen'] += wait_time case 'check_install': self._check_install(success_condition.check_install) diff --git a/android_env/components/simulators/base_simulator.py b/android_env/components/simulators/base_simulator.py index 66b8e79c..f7aee84e 100644 --- a/android_env/components/simulators/base_simulator.py +++ b/android_env/components/simulators/base_simulator.py @@ -17,6 +17,7 @@ import abc from collections.abc import Callable +import datetime import threading import time @@ -85,7 +86,7 @@ def launch(self) -> None: ) from error # Start interaction thread. - if self._config.interaction_rate_sec > 0: + if self._config.interaction_rate_sec > datetime.timedelta(seconds=0): self._interaction_thread = InteractionThread( self._get_screenshot_impl, self._config.interaction_rate_sec ) @@ -151,7 +152,7 @@ def save_state( def get_screenshot(self) -> np.ndarray: """Returns pixels representing the current screenshot of the simulator.""" - if self._config.interaction_rate_sec > 0: + if self._config.interaction_rate_sec > datetime.timedelta(seconds=0): assert self._interaction_thread is not None return self._interaction_thread.screenshot() # Async mode. else: @@ -180,7 +181,7 @@ class InteractionThread(threading.Thread): def __init__( self, get_screenshot_fn: Callable[[], np.ndarray], - interaction_rate_sec: float, + interaction_rate_sec: datetime.timedelta, ): super().__init__() self._get_screenshot_fn = get_screenshot_fn @@ -189,15 +190,15 @@ def __init__( self._screenshot = self._get_screenshot_fn() def run(self): - last_read = time.time() + last_read = time.monotonic() while not self._should_stop.is_set(): self._screenshot = self._get_screenshot_fn() - now = time.time() - elapsed = now - last_read + now = time.monotonic() + elapsed = datetime.timedelta(seconds=now - last_read) last_read = now sleep_time = self._interaction_rate_sec - elapsed - if sleep_time > 0.0: - time.sleep(sleep_time) + if sleep_time > datetime.timedelta(seconds=0): + time.sleep(sleep_time.total_seconds()) logging.info('InteractionThread.run() finished.') def stop(self): diff --git a/android_env/components/simulators/base_simulator_test.py b/android_env/components/simulators/base_simulator_test.py index d0ee5455..8c282956 100644 --- a/android_env/components/simulators/base_simulator_test.py +++ b/android_env/components/simulators/base_simulator_test.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime import itertools import time from unittest import mock @@ -79,7 +80,9 @@ def test_get_screenshot_error_async(self): errors.ReadObservationError() ) simulator = fake_simulator.FakeSimulator( - config_classes.FakeSimulatorConfig(interaction_rate_sec=0.5) + config_classes.FakeSimulatorConfig( + interaction_rate_sec=datetime.timedelta(seconds=0.5) + ) ) with mock.patch.object( base_simulator, @@ -99,7 +102,7 @@ def test_get_screenshot_faster_than_screenshot_impl(self): """Return same screenshot when step is faster than the interaction rate.""" # Arrange. - slow_rate = 0.5 + slow_rate = datetime.timedelta(seconds=0.5) simulator = fake_simulator.FakeSimulator( config_classes.FakeSimulatorConfig(interaction_rate_sec=slow_rate) ) @@ -126,7 +129,7 @@ def test_get_screenshot_slower_than_screenshot_impl(self): """Return different screenshots when step slower than the interaction rate.""" # Arrange. - fast_rate = 0.01 + fast_rate = datetime.timedelta(seconds=0.01) simulator = fake_simulator.FakeSimulator( config_classes.FakeSimulatorConfig(interaction_rate_sec=fast_rate) ) @@ -155,7 +158,9 @@ def test_interaction_thread_closes_upon_relaunch(self): # Arrange. simulator = fake_simulator.FakeSimulator( - config_classes.FakeSimulatorConfig(interaction_rate_sec=0.01) + config_classes.FakeSimulatorConfig( + interaction_rate_sec=datetime.timedelta(seconds=0.01) + ) ) mock_interaction_thread = mock.create_autospec( base_simulator.InteractionThread diff --git a/android_env/components/simulators/emulator/emulator_simulator.py b/android_env/components/simulators/emulator/emulator_simulator.py index 4e7af5c4..e49c60d9 100644 --- a/android_env/components/simulators/emulator/emulator_simulator.py +++ b/android_env/components/simulators/emulator/emulator_simulator.py @@ -16,6 +16,7 @@ """A class that manages an Android Emulator.""" from collections.abc import Callable, Mapping +import datetime import functools import os import platform @@ -370,7 +371,7 @@ def save_state( def _connect_to_emulator( self, grpc_port: int, - timeout_sec: int = 100, + timeout: datetime.timedelta = datetime.timedelta(seconds=100), ) -> tuple[ emulator_controller_pb2_grpc.EmulatorControllerStub, snapshot_service_pb2_grpc.SnapshotServiceStub, @@ -390,7 +391,9 @@ def _connect_to_emulator( try: self._channel = grpc.secure_channel(port, creds, options=options) - grpc.channel_ready_future(self._channel).result(timeout=timeout_sec) + grpc.channel_ready_future(self._channel).result( + timeout=timeout.total_seconds() + ) except (grpc.RpcError, grpc.FutureTimeoutError) as grpc_error: logging.exception('Failed to connect to the emulator.') raise EmulatorBootError( @@ -404,16 +407,19 @@ def _connect_to_emulator( return emulator_controller_stub, snapshot_stub @_reconnect_on_grpc_error - def _confirm_booted(self, startup_wait_time_sec: int = 300): + def _confirm_booted( + self, + startup_wait_time: datetime.timedelta = datetime.timedelta(seconds=300), + ): """Waits until the emulator is fully booted.""" assert ( self._emulator_stub is not None ), 'Emulator stub has not been initialized yet.' - start_time = time.time() - deadline = start_time + startup_wait_time_sec + start_time = time.monotonic() + deadline = start_time + startup_wait_time.total_seconds() success = False - while time.time() < deadline: + while time.monotonic() < deadline: emu_status = self._emulator_stub.getStatus(empty_pb2.Empty()) logging.info('Waiting for emulator (%r) to start... (%rms)', self.adb_device_name(), emu_status.uptime) @@ -422,10 +428,12 @@ def _confirm_booted(self, startup_wait_time_sec: int = 300): break time.sleep(5.0) - elapsed_time = time.time() - start_time + elapsed_time = time.monotonic() - start_time if not success: raise EmulatorCrashError( - f'The emulator failed to boot after {startup_wait_time_sec} seconds') + 'The emulator failed to boot after' + f' {startup_wait_time.total_seconds()} seconds' + ) logging.info('Done booting the emulator (in %f seconds).', elapsed_time) logging.info('********** Emulator logs **********') diff --git a/android_env/components/simulators/emulator/emulator_simulator_test.py b/android_env/components/simulators/emulator/emulator_simulator_test.py index 708a274f..8824d5db 100644 --- a/android_env/components/simulators/emulator/emulator_simulator_test.py +++ b/android_env/components/simulators/emulator/emulator_simulator_test.py @@ -16,6 +16,7 @@ """Tests for android_env.components.emulator_simulator.""" import builtins +import datetime import os import time from unittest import mock @@ -342,7 +343,7 @@ def test_get_screenshot(self): def test_get_screenshot_reconnects_on_grpc_error(self): config = config_classes.EmulatorConfig( - interaction_rate_sec=0.0, + interaction_rate_sec=datetime.timedelta(seconds=0.0), emulator_launcher=config_classes.EmulatorLauncherConfig( grpc_port=1234, tmp_dir=self.create_tempdir().full_path ), diff --git a/android_env/components/simulators/fake/fake_simulator.py b/android_env/components/simulators/fake/fake_simulator.py index 4209b17c..8893e315 100644 --- a/android_env/components/simulators/fake/fake_simulator.py +++ b/android_env/components/simulators/fake/fake_simulator.py @@ -15,6 +15,7 @@ """Fake Simulator for testing AndroidEnv infrastructure.""" +import datetime import random import threading import time @@ -67,7 +68,7 @@ class FakeAdbController(adb_controller.AdbController): def execute_command( self, args: list[str], - timeout: float | None = None, + timeout: datetime.timedelta | None = None, device_specific: bool = True, ) -> bytes: """Returns fake output for adb commands.""" diff --git a/android_env/wrappers/a11y_grpc_wrapper.py b/android_env/wrappers/a11y_grpc_wrapper.py index b6e3d965..a859fede 100644 --- a/android_env/wrappers/a11y_grpc_wrapper.py +++ b/android_env/wrappers/a11y_grpc_wrapper.py @@ -16,6 +16,7 @@ """Wraps AndroidEnv to retrieve accessibility messages from gRPC.""" from concurrent import futures +import datetime import time from typing import Any @@ -84,7 +85,7 @@ def __init__( start_a11y_service: bool = True, enable_a11y_tree_info: bool = False, add_latest_a11y_info_to_obs: bool = False, - a11y_info_timeout: float | None = None, + a11y_info_timeout: datetime.timedelta | None = None, max_enable_networking_attempts: int = 10, latest_a11y_info_only: bool = False, grpc_server_ip: str = '10.0.2.2', @@ -158,7 +159,10 @@ def __init__( self._add_latest_a11y_info_to_obs = add_latest_a11y_info_to_obs self._a11y_info_timeout = a11y_info_timeout self._parent_action_spec = self._env.action_spec() - if self._a11y_info_timeout is not None and self._a11y_info_timeout > 0.0: + if ( + self._a11y_info_timeout is not None + and self._a11y_info_timeout > datetime.timedelta(seconds=0.0) + ): if 'action_type' not in self._parent_action_spec.keys(): raise ValueError( 'action_type not in the parent action spec: ' @@ -344,7 +348,9 @@ def _configure_grpc(self) -> None: ) def _accumulate_and_return_a11y_info( - self, timer: float | None = None, get_env_observation: bool = True + self, + timer: datetime.timedelta | None = None, + get_env_observation: bool = True, ) -> dict[str, Any]: """Accumulates and returns the latest a11y tree info and observation. @@ -358,9 +364,9 @@ def _accumulate_and_return_a11y_info( a dict with a11y forest under key 'a11y_forest'. All other fields will provide the observation, if requested. """ - timer = timer or 0.0 - if timer > 0.0: - time.sleep(timer) + timer = timer or datetime.timedelta() + if timer > datetime.timedelta(seconds=0.0): + time.sleep(timer.total_seconds()) if get_env_observation: # Fetch observation. @@ -383,9 +389,11 @@ def _accumulate_and_return_a11y_info( return observation def _fetch_task_extras_and_update_observation( - self, observation: dict[str, Any], timeout: float = 0.0 + self, + observation: dict[str, Any], + timeout: datetime.timedelta = datetime.timedelta(seconds=0.0), ) -> dict[str, Any]: - if timeout > 0.0: + if timeout > datetime.timedelta(seconds=0.0): observation = self._accumulate_and_return_a11y_info( timeout, get_env_observation=True ) @@ -406,7 +414,7 @@ def reset(self) -> dm_env.TimeStep: self._configure_grpc() self._relaunch_count = self._env.stats()['relaunch_count'] self._accumulated_extras = {} - timeout = self._a11y_info_timeout or 0.0 + timeout = self._a11y_info_timeout or datetime.timedelta() new_observation = self._fetch_task_extras_and_update_observation( timestep.observation, timeout ) @@ -414,7 +422,11 @@ def reset(self) -> dm_env.TimeStep: return timestep def step(self, action: Any) -> dm_env.TimeStep: - timeout = float(action.pop('wait_time', self._a11y_info_timeout or 0.0)) + wait_time = action.pop('wait_time', None) + if wait_time is not None: + timeout = datetime.timedelta(seconds=float(wait_time)) + else: + timeout = self._a11y_info_timeout or datetime.timedelta() timestep = self._env.step(action) new_observation = self._fetch_task_extras_and_update_observation( timestep.observation, timeout=timeout diff --git a/android_env/wrappers/a11y_grpc_wrapper_test.py b/android_env/wrappers/a11y_grpc_wrapper_test.py index abe0901a..2ffdad3b 100644 --- a/android_env/wrappers/a11y_grpc_wrapper_test.py +++ b/android_env/wrappers/a11y_grpc_wrapper_test.py @@ -15,6 +15,7 @@ """Tests for a11y_grpc_wrapper.""" +import datetime import time from unittest import mock @@ -536,7 +537,9 @@ def test_a11y_info_disabled( observation={'dummy': 42}, reward=0.0 ) wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper( - base_env, add_latest_a11y_info_to_obs=False, a11y_info_timeout=1.0 + base_env, + add_latest_a11y_info_to_obs=False, + a11y_info_timeout=datetime.timedelta(seconds=1.0), ) with mock.patch.object(wrapped_env, '_fetch_task_extras'): wrapped_env._fetch_task_extras.return_value = { @@ -569,7 +572,9 @@ def test_a11y_info_with_timer_info_present( observation={'dummy': 42}, reward=0.0 ) wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper( - base_env, add_latest_a11y_info_to_obs=True, a11y_info_timeout=1.0 + base_env, + add_latest_a11y_info_to_obs=True, + a11y_info_timeout=datetime.timedelta(seconds=1.0), ) with mock.patch.object(wrapped_env, '_fetch_task_extras'): wrapped_env._fetch_task_extras.side_effect = [{ @@ -600,7 +605,9 @@ def test_a11y_info_with_timer_task_extra_returned( observation={'dummy': 42}, reward=0.0 ) wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper( - base_env, add_latest_a11y_info_to_obs=True, a11y_info_timeout=1.0 + base_env, + add_latest_a11y_info_to_obs=True, + a11y_info_timeout=datetime.timedelta(seconds=1.0), ) with mock.patch.object(wrapped_env, '_fetch_task_extras'): wrapped_env._fetch_task_extras.side_effect = [ @@ -635,7 +642,9 @@ def test_a11y_info_with_timer_from_action( observation={'dummy': 42}, reward=0.0 ) wrapped_env = a11y_grpc_wrapper.A11yGrpcWrapper( - base_env, add_latest_a11y_info_to_obs=True, a11y_info_timeout=0.0 + base_env, + add_latest_a11y_info_to_obs=True, + a11y_info_timeout=datetime.timedelta(seconds=0.0), ) with mock.patch.object(wrapped_env, '_fetch_task_extras'): wrapped_env._fetch_task_extras.side_effect = [ diff --git a/android_env/wrappers/rate_limit_wrapper.py b/android_env/wrappers/rate_limit_wrapper.py index ab3895a3..c0395655 100644 --- a/android_env/wrappers/rate_limit_wrapper.py +++ b/android_env/wrappers/rate_limit_wrapper.py @@ -15,6 +15,7 @@ """Limits interactions with the environment to a given rate.""" +import datetime import enum import time @@ -46,15 +47,17 @@ class SleepType(enum.IntEnum): # because the sleep does not account for the time taken by step(). AFTER_WITH_REPEAT = 2 - def __init__(self, - env: env_interface.AndroidEnvInterface, - rate: float, - sleep_type: SleepType = SleepType.AFTER_WITH_REPEAT): + def __init__( + self, + env: env_interface.AndroidEnvInterface, + max_wait: datetime.timedelta, + sleep_type: SleepType = SleepType.AFTER_WITH_REPEAT, + ): """Initializes this wrapper. Args: env: The underlying environment to which this wrapper is applied. - rate: The desired rate in Hz to interact with the environment. If <=0.0, + max_wait: The desired maximum wait time between interactions. If <=0.0, this wrapper will be disabled. sleep_type: This determines how the wrapper will interact with the underlying AndroidEnv environment. @@ -62,7 +65,7 @@ def __init__(self, super().__init__(env) self._assert_base_env() self._last_step_time = None - self._max_wait = 1.0 / rate if rate > 0.0 else 0.0 + self._max_wait = max_wait self._sleep_type = sleep_type def _assert_base_env(self): @@ -74,14 +77,14 @@ def _assert_base_env(self): def reset(self): timestep = self._env.reset() - self._last_step_time = time.time() + self._last_step_time = time.monotonic() return timestep def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep: """Takes a step while maintaining a steady interaction rate.""" # If max_wait is non-positive, the wrapper has no effect. - if self._max_wait <= 0.0: + if self._max_wait <= datetime.timedelta(seconds=0.0): return self._env.step(action) if self._sleep_type == RateLimitWrapper.SleepType.BEFORE: @@ -105,13 +108,18 @@ def step(self, action: dict[str, np.ndarray]) -> dm_env.TimeStep: elif self._sleep_type == RateLimitWrapper.SleepType.AFTER: self._wait() - self._last_step_time = time.time() + self._last_step_time = time.monotonic() return timestep def _wait(self) -> None: - if self._max_wait > 0.0 and self._last_step_time is not None: - time_since_step = time.time() - self._last_step_time - sec_to_wait = self._max_wait - time_since_step - if sec_to_wait > 0.0: - time.sleep(sec_to_wait) + if ( + self._max_wait > datetime.timedelta(seconds=0.0) + and self._last_step_time is not None + ): + time_since_step = datetime.timedelta( + seconds=time.monotonic() - self._last_step_time + ) + to_wait = self._max_wait - time_since_step + if to_wait > datetime.timedelta(seconds=0.0): + time.sleep(to_wait.total_seconds()) diff --git a/android_env/wrappers/rate_limit_wrapper_test.py b/android_env/wrappers/rate_limit_wrapper_test.py index 5e1b8ad1..e9d69f7d 100644 --- a/android_env/wrappers/rate_limit_wrapper_test.py +++ b/android_env/wrappers/rate_limit_wrapper_test.py @@ -15,6 +15,7 @@ """Tests for rate_limit_wrapper.""" +import datetime import time from typing import Any, Protocol from unittest import mock @@ -61,14 +62,14 @@ def _with_timestamp(fn: Any) -> _FnWithTimestamps: class RateLimitWrapperTest(parameterized.TestCase): @parameterized.named_parameters( - ('zero_rate', 0), - ('negative_rate', -50), + ('zero_wait', datetime.timedelta(seconds=0.0)), + ('negative_wait', datetime.timedelta(seconds=-1.0)), ) @mock.patch.object(time, 'sleep', autospec=True) - def test_disabled(self, rate, mock_sleep): + def test_disabled(self, max_wait, mock_sleep): """With a non-positive rate, this wrapper should do nothing.""" env = _get_base_env() - wrapper = rate_limit_wrapper.RateLimitWrapper(env, rate=rate) + wrapper = rate_limit_wrapper.RateLimitWrapper(env, max_wait=max_wait) _ = wrapper.reset() mock_sleep.assert_not_called() _ = wrapper.step({ @@ -85,7 +86,9 @@ def test_enabled(self, mock_sleep): env = _get_base_env() env.step.return_value = dm_env.transition(reward=None, observation=None) - wrapper = rate_limit_wrapper.RateLimitWrapper(env, rate=1/33.33) + wrapper = rate_limit_wrapper.RateLimitWrapper( + env, max_wait=datetime.timedelta(seconds=33.33) + ) _ = wrapper.reset() mock_sleep.assert_not_called() # It should never sleep during reset(). @@ -113,8 +116,9 @@ def test_enabled_sleep_type_before(self, mock_sleep): env = _get_base_env() wrapper = rate_limit_wrapper.RateLimitWrapper( env, - rate=1/33.33, - sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.BEFORE) + max_wait=datetime.timedelta(seconds=33.33), + sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.BEFORE, + ) _ = wrapper.reset() mock_sleep.assert_not_called() # It should never sleep during reset(). @@ -154,8 +158,9 @@ def test_enabled_sleep_type_after(self, mock_sleep): env = _get_base_env() wrapper = rate_limit_wrapper.RateLimitWrapper( env, - rate=1/33.33, - sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.AFTER) + max_wait=datetime.timedelta(seconds=33.33), + sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.AFTER, + ) _ = wrapper.reset() mock_sleep.assert_not_called() # It should never sleep during reset(). @@ -194,9 +199,9 @@ def test_enabled_sleep_type_after_with_repeat(self, mock_sleep): env = _get_base_env() wrapper = rate_limit_wrapper.RateLimitWrapper( env, - rate=1/33.33, - sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType - .AFTER_WITH_REPEAT) + max_wait=datetime.timedelta(seconds=33.33), + sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.AFTER_WITH_REPEAT, + ) _ = wrapper.reset() mock_sleep.assert_not_called() # It should never sleep during reset(). @@ -249,9 +254,9 @@ def test_enabled_sleep_type_after_with_repeat_last(self, mock_sleep): env = _get_base_env() wrapper = rate_limit_wrapper.RateLimitWrapper( env, - rate=1/33.33, - sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType - .AFTER_WITH_REPEAT) + max_wait=datetime.timedelta(seconds=33.33), + sleep_type=rate_limit_wrapper.RateLimitWrapper.SleepType.AFTER_WITH_REPEAT, + ) _ = wrapper.reset() mock_sleep.assert_not_called() # It should never sleep during reset().