diff --git a/android_env/components/adb_call_parser.py b/android_env/components/adb_call_parser.py index c38bf3f..ee401e2 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 c3f4cfb..5c62aec 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 8273757..c68b332 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 71a3a4e..aeb0af2 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/config_classes.py b/android_env/components/config_classes.py index 528974f..1a41b8f 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. diff --git a/android_env/components/simulators/fake/fake_simulator.py b/android_env/components/simulators/fake/fake_simulator.py index 4209b17..8893e31 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."""