Refactor android_env simulator and emulator timeouts to use datetime.timedelta.

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.
 

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']

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)