android-env 1.2.1__py3-none-any.whl → 1.2.3__py3-none-any.whl

android_env/components/adb_controller.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,33 +18,20 @@
 import os
 import subprocess
 import time
-from typing import List, Optional
 
 from absl import logging
+from android_env.components import config_classes
 from android_env.components import errors
 
 
-class AdbController():
+class AdbController:
   """Manages communication with adb."""
 
-  def __init__(self,
-               device_name: str = '',
-               adb_path: str = 'adb',
-               adb_server_port: int = 5037,
-               default_timeout: float = 120.0):
-    """Instantiates an AdbController object.
+  def __init__(self, config: config_classes.AdbControllerConfig):
+    """Instantiates an AdbController object."""
 
-    Args:
-      device_name: Name of the device to communicate with.
-      adb_path: Path to the adb binary.
-      adb_server_port: Port for adb server.
-      default_timeout: Default timeout in seconds.
-    """
-
-    self._device_name = device_name
-    self._adb_path = adb_path
-    self._adb_server_port = str(adb_server_port)
-    self._default_timeout = default_timeout
+    self._config = config
+    logging.info('config: %r', self._config)
 
     # Unset problematic environment variables. ADB commands will fail if these
     # are set. They are normally exported by AndroidStudio.
@@ -53,14 +40,25 @@ class AdbController():
     if 'ANDROID_ADB_SERVER_PORT' in os.environ:
       del os.environ['ANDROID_ADB_SERVER_PORT']
 
-  def command_prefix(self) -> List[str]:
+    # Explicitly expand the $HOME environment variable.
+    self._os_env_vars = dict(os.environ).copy()
+    self._os_env_vars.update(
+        {'HOME': os.path.expandvars(self._os_env_vars.get('HOME', ''))}
+    )
+    logging.info('self._os_env_vars: %r', self._os_env_vars)
+
+  def command_prefix(self, include_device_name: bool = True) -> list[str]:
     """The command for instantiating an adb client to this server."""
-    command_prefix = [self._adb_path, '-P', self._adb_server_port]
-    if self._device_name:
-      command_prefix.extend(['-s', self._device_name])
+    command_prefix = [
+        self._config.adb_path,
+        '-P',
+        str(self._config.adb_server_port),
+    ]
+    if include_device_name:
+      command_prefix.extend(['-s', self._config.device_name])
     return command_prefix
 
-  def init_server(self, timeout: Optional[float] = None):
+  def init_server(self, timeout: float | None = None):
     """Initialize the ADB server deamon on the given port.
 
     This function should be called immediately after initializing the first
@@ -71,16 +69,34 @@ class AdbController():
         timeout set on the constructor will be used.
     """
     # Make an initial device-independent call to ADB to start the deamon.
-    device_name_tmp = self._device_name
-    self._device_name = ''
-    self.execute_command(['devices'], timeout=timeout)
+    self.execute_command(['devices'], timeout, device_specific=False)
+    time.sleep(0.2)
+
+  def _restart_server(self, timeout: float | None = None):
+    """Kills and restarts the adb server.
+
+    Args:
+      timeout: A timeout to use for this operation. If not set the default
+        timeout set on the constructor will be used.
+    """
+    logging.info('Restarting adb server.')
+    self.execute_command(
+        ['kill-server'], timeout=timeout, device_specific=False)
+    time.sleep(0.2)
+    cmd_output = self.execute_command(
+        ['start-server'], timeout=timeout, device_specific=False)
+    logging.info('start-server output: %r', cmd_output.decode('utf-8'))
+    time.sleep(2.0)
+    self.execute_command(
+        ['devices'], timeout=timeout, device_specific=False)
     time.sleep(0.2)
-    # Subsequent calls will use the device name.
-    self._device_name = device_name_tmp
 
-  def execute_command(self,
-                      args: List[str],
-                      timeout: Optional[float] = None) -> bytes:
+  def execute_command(
+      self,
+      args: list[str],
+      timeout: float | None = None,
+      device_specific: bool = True,
+  ) -> bytes:
     """Executes an adb command.
 
     Args:
@@ -88,26 +104,48 @@ class AdbController():
           For example: ['install', '/my/app.apk']
       timeout: A timeout to use for this operation. If not set the default
         timeout set on the constructor will be used.
+      device_specific: Whether the call is device-specific or independent.
 
     Returns:
       The output of running such command as a binary string.
     """
-    timeout = self._default_timeout if timeout is None else timeout
-    command = self.command_prefix() + args
-    command_str = ' '.join(command)
-    logging.info('Executing ADB command: [%s]', command_str)
-
-    try:
-      cmd_output = subprocess.check_output(
-          command, stderr=subprocess.STDOUT, timeout=timeout)
-      logging.info('Done executing ADB command: [%s]', command_str)
-      return cmd_output
-    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error:
-      if error.stdout is not None:
-        logging.error('**stdout**:')
-        for line in error.stdout.splitlines():
-          logging.error(line)
-      raise errors.AdbControllerError(
-          f'Error executing adb command: [{command_str}]\n'
-          f'Caused by {error}\n'
-          f'adb output: [{error.stdout}]') from error
+    timeout = self._config.default_timeout if timeout is None else timeout
+    command = self.command_prefix(include_device_name=device_specific) + args
+    command_str = 'adb ' + ' '.join(command[1:])
+
+    n_retries = 2
+    n_tries = 1
+    latest_error = None
+    while n_tries <= n_retries:
+      try:
+        logging.info('Executing ADB command: [%s]', command_str)
+        cmd_output = subprocess.check_output(
+            command,
+            stderr=subprocess.STDOUT,
+            timeout=timeout,
+            env=self._os_env_vars,
+        )
+        logging.debug('ADB command output: %s', cmd_output)
+        return cmd_output
+      except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as e:
+        logging.exception(
+            'Failed to execute ADB command (try %r of 3): [%s]',
+            n_tries, command_str)
+        if e.stdout is not None:
+          logging.error('**stdout**:')
+          for line in e.stdout.splitlines():
+            logging.error('    %s', line)
+        if e.stderr is not None:
+          logging.error('**stderr**:')
+          for line in e.stderr.splitlines():
+            logging.error('    %s', line)
+        n_tries += 1
+        latest_error = e
+        if device_specific and n_tries <= n_retries:
+          self._restart_server(timeout=timeout)
+
+    raise errors.AdbControllerError(
+        f'Error executing adb command: [{command_str}]\n'
+        f'Caused by: {latest_error}\n'
+        f'adb stdout: [{latest_error.stdout}]\n'
+        f'adb stderr: [{latest_error.stderr}]') from latest_error

android_env/components/adb_controller_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,11 +16,14 @@
 """Tests for android_env.components.adb_controller."""
 
 import os
+import subprocess
 import time
 from unittest import mock
 
 from absl.testing import absltest
-from android_env.components import adb_controller
+from android_env.components import adb_controller as adb_controller_lib
+from android_env.components import config_classes
+from android_env.components import errors
 
 # Timeout to be used by default in tests below. Set to a small value to avoid
 # hanging on a failed test.
@@ -31,30 +34,198 @@ class AdbControllerTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
-    self._mock_execute_command = self.enter_context(
-        mock.patch.object(
-            adb_controller.AdbController, 'execute_command', autospec=True))
-    self._adb_controller = adb_controller.AdbController(
-        adb_path='my_adb', device_name='awesome_device', adb_server_port=9999)
+    # Set two env vars.
+    os.environ['MY_ENV_VAR'] = '/some/path/'
+    os.environ['HOME'] = '$MY_ENV_VAR'
+    self._env_before = os.environ
+    self._adb_controller = adb_controller_lib.AdbController(
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
 
+  @mock.patch.object(subprocess, 'check_output', autospec=True)
   @mock.patch.object(time, 'sleep', autospec=True)
-  def test_init_server(self, mock_sleep):
-    self._adb_controller.init_server(timeout=_TIMEOUT)
-    self._mock_execute_command.assert_called_once_with(self._adb_controller,
-                                                       ['devices'], _TIMEOUT)
+  def test_init_server(self, mock_sleep, mock_check_output):
+    # Arrange.
+    adb_controller = adb_controller_lib.AdbController(
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
+
+    # Act.
+    adb_controller.init_server(timeout=_TIMEOUT)
+
+    # Assert.
+    expected_env = self._env_before
+    expected_env['HOME'] = '/some/path/'
+    mock_check_output.assert_called_once_with(
+        ['my_adb', '-P', '9999', 'devices'],
+        stderr=subprocess.STDOUT,
+        timeout=_TIMEOUT,
+        env=expected_env,
+    )
     mock_sleep.assert_called_once()
 
+  @mock.patch.object(subprocess, 'check_output', autospec=True)
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_restart_server(self, mock_sleep, mock_check_output):
+    # Arrange.
+    mock_check_output.side_effect = [
+        subprocess.CalledProcessError(returncode=1, cmd='blah'),
+    ] + ['fake_output'.encode('utf-8')] * 4
+    adb_controller = adb_controller_lib.AdbController(
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
+
+    # Act.
+    adb_controller.execute_command(['my_command'], timeout=_TIMEOUT)
+
+    # Assert.
+    expected_env = self._env_before
+    expected_env['HOME'] = '/some/path/'
+    mock_check_output.assert_has_calls([
+        mock.call(
+            ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+            stderr=subprocess.STDOUT,
+            timeout=_TIMEOUT,
+            env=expected_env,
+        ),
+        mock.call(
+            ['my_adb', '-P', '9999', 'kill-server'],
+            stderr=subprocess.STDOUT,
+            timeout=_TIMEOUT,
+            env=expected_env,
+        ),
+        mock.call(
+            ['my_adb', '-P', '9999', 'start-server'],
+            stderr=subprocess.STDOUT,
+            timeout=_TIMEOUT,
+            env=expected_env,
+        ),
+        mock.call(
+            ['my_adb', '-P', '9999', 'devices'],
+            stderr=subprocess.STDOUT,
+            timeout=_TIMEOUT,
+            env=expected_env,
+        ),
+        mock.call(
+            ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+            stderr=subprocess.STDOUT,
+            timeout=_TIMEOUT,
+            env=expected_env,
+        ),
+    ])
+    mock_sleep.assert_has_calls(
+        [mock.call(0.2), mock.call(2.0), mock.call(0.2)])
+
+  @mock.patch.object(subprocess, 'check_output', autospec=True)
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_invalid_command(self, mock_sleep, mock_check_output):
+    # Arrange.
+    restart_sequence = ['fake_output'.encode('utf-8')] * 3
+    mock_check_output.side_effect = (
+        [
+            subprocess.CalledProcessError(returncode=1, cmd='blah'),
+        ]
+        + restart_sequence
+        + [subprocess.CalledProcessError(returncode=1, cmd='blah')]
+        # Don't restart if last call fails.
+    )
+    adb_controller = adb_controller_lib.AdbController(
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
+
+    # Act.
+    with self.assertRaises(errors.AdbControllerError):
+      adb_controller.execute_command(['my_command'], timeout=_TIMEOUT)
+
+    # Assert.
+    expected_env = self._env_before
+    expected_env['HOME'] = '/some/path/'
+    mock_check_output.assert_has_calls(
+        [
+            mock.call(
+                ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'kill-server'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'start-server'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'devices'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+        ],
+        any_order=False,
+    )
+    mock_sleep.assert_has_calls(
+        [mock.call(0.2), mock.call(2.0), mock.call(0.2)]
+    )
+
+  @mock.patch.object(subprocess, 'check_output', autospec=True)
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_avoid_infinite_recursion(self, mock_sleep, mock_check_output):
+    del mock_sleep
+    mock_check_output.side_effect = subprocess.CalledProcessError(
+        returncode=1, cmd='blah')
+    adb_controller = adb_controller_lib.AdbController(
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+        )
+    )
+    self.assertRaises(
+        errors.AdbControllerError,
+        adb_controller.execute_command, ['my_command'], timeout=_TIMEOUT)
+
 
 class AdbControllerInitTest(absltest.TestCase):
 
   def test_deletes_problem_env_vars(self):
     os.environ['ANDROID_HOME'] = '/usr/local/Android/Sdk'
     os.environ['ANDROID_ADB_SERVER_PORT'] = '1337'
-    adb_controller.AdbController(
-        adb_path='my_adb',
-        device_name='awesome_device',
-        adb_server_port=9999,
-        default_timeout=_TIMEOUT)
+    adb_controller_lib.AdbController(
+        config_classes.AdbControllerConfig(
+            adb_path='my_adb',
+            device_name='awesome_device',
+            adb_server_port=9999,
+            default_timeout=_TIMEOUT,
+        )
+    )
     self.assertNotIn('ANDROID_HOME', os.environ)
     self.assertNotIn('ANDROID_ADB_SERVER_PORT', os.environ)
 

android_env/components/adb_log_stream.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,8 +16,8 @@
 """Class for a stream of logs output by a locally running emulator."""
 
 import subprocess
-from typing import List
 
+from absl import logging
 from android_env.components import log_stream
 
 
@@ -27,11 +27,19 @@ _LOGCAT_COMMAND = ['logcat', '-v', 'epoch']
 class AdbLogStream(log_stream.LogStream):
   """Manages adb logcat process for a locally running emulator."""
 
-  def __init__(self, adb_command_prefix: List[str], *args, **kwargs):
-    super().__init__(*args, **kwargs)
+  def __init__(self, adb_command_prefix: list[str], verbose: bool = False):
+    super().__init__(verbose=verbose)
     self._adb_command_prefix = adb_command_prefix
 
   def _get_stream_output(self):
+
+    # Before spawning a long-lived process, we issue `logcat -b all -c` to clear
+    # all buffers to avoid interference from previous runs.
+    clear_buffer_output = subprocess.check_output(
+        self._adb_command_prefix + ['logcat', '-b', 'all', '-c'],
+        stderr=subprocess.STDOUT,
+        timeout=100)
+    logging.info('clear_buffer_output: %r', clear_buffer_output)
     cmd = self._adb_command_prefix + _LOGCAT_COMMAND + self._filters
     self._adb_subprocess = subprocess.Popen(
         cmd,
@@ -42,4 +50,8 @@ class AdbLogStream(log_stream.LogStream):
     return self._adb_subprocess.stdout
 
   def stop_stream(self):
-    self._adb_subprocess.kill()
+    if not hasattr(self, '_adb_subprocess') or self._adb_subprocess is None:
+      logging.error('`stop_stream()` called before `get_stream_output()`. '
+                    'This violates the `LogStream` API.')
+    else:
+      self._adb_subprocess.kill()

android_env/components/adb_log_stream_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -22,7 +22,7 @@ from absl.testing import absltest
 from android_env.components import adb_log_stream
 
 
-class FakeAdbSubprocess():
+class FakeAdbSubprocess:
 
   @property
   def stdout(self):
@@ -34,8 +34,9 @@ class FakeAdbSubprocess():
 
 class AdbLogStreamTest(absltest.TestCase):
 
+  @mock.patch.object(subprocess, 'check_output', return_value=b'')
   @mock.patch.object(subprocess, 'Popen', return_value=FakeAdbSubprocess())
-  def test_get_stream_output(self, mock_popen):
+  def test_get_stream_output(self, mock_popen, unused_mock_check_output):
     stream = adb_log_stream.AdbLogStream(adb_command_prefix=['foo'])
     stream.set_log_filters(['bar'])
     stream_output = stream.get_stream_output()
@@ -50,6 +51,19 @@ class AdbLogStreamTest(absltest.TestCase):
         bufsize=1,
         universal_newlines=True)
 
+  def test_stop_stream_before_get_stream_output(self):
+    """Calling `stop_stream()` before `get_stream_output()` should not crash."""
+
+    # Arrange.
+    stream = adb_log_stream.AdbLogStream(adb_command_prefix=['foo'])
+
+    # Act.
+    stream.stop_stream()
+
+    # Assert.
+    # Nothing to assert. The test should just finish without raising an
+    # exception.
+
 
 if __name__ == '__main__':
   absltest.main()

android_env/components/app_screen_checker.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -15,23 +15,23 @@
 
 """Determines if the current app screen matches an expected app screen."""
 
+from collections.abc import Callable, Sequence
 import enum
 import re
 import time
-from typing import Callable, List, Optional, Sequence, Pattern
+from typing import Self
 
 from absl import logging
-
 from android_env.components import adb_call_parser as adb_call_parser_lib
 from android_env.components import errors
 from android_env.proto import adb_pb2
 from android_env.proto import task_pb2
 
 
-class _DumpsysNode():
+class _DumpsysNode:
   """A node in a dumpsys tree."""
 
-  def __init__(self, data: Optional[str] = None):
+  def __init__(self, data: str):
     self._children = []
     self._data = data
 
@@ -40,12 +40,12 @@ class _DumpsysNode():
     return self._data
 
   @property
-  def children(self) -> List['_DumpsysNode']:
+  def children(self) -> list[Self]:
     return self._children
 
-  def find_child(self,
-                 predicate: Callable[['_DumpsysNode'], bool],
-                 max_levels: int = 0) -> Optional['_DumpsysNode']:
+  def find_child(
+      self, predicate: Callable[[Self], bool], max_levels: int = 0
+  ) -> Self | None:
     """Returns the first direct child that matches `predicate`, None otherwise.
 
     Args:
@@ -126,9 +126,11 @@ def build_tree_from_dumpsys_output(dumpsys_output: str) -> _DumpsysNode:
   return root
 
 
-def matches_path(dumpsys_activity_output: str,
-                 expected_view_hierarchy_path: Sequence[Pattern[str]],
-                 max_levels: int = 0) -> bool:
+def matches_path(
+    dumpsys_activity_output: str,
+    expected_view_hierarchy_path: Sequence[re.Pattern[str]],
+    max_levels: int = 0,
+) -> bool:
   """Returns True if the current dumpsys output matches the expected path.
 
   Args:
@@ -150,7 +152,7 @@ def matches_path(dumpsys_activity_output: str,
         'view_hierarchy is None. Dumpsys activity output: %s. tree: %r',
         str(dumpsys_activity_output), root.print_tree())
     logging.error('Tree root: %s', str(root))
-    return None
+    return False
 
   current_node = view_hierarchy
   for i, regex in enumerate(expected_view_hierarchy_path):
@@ -165,13 +167,13 @@ def matches_path(dumpsys_activity_output: str,
                     regex.pattern, current_node)
       logging.error('Dumpsys activity output: %s', str(dumpsys_activity_output))
       logging.error('Tree root: %s', str(root))
-      return None
+      return False
     else:
       current_node = child
   return True
 
 
-class AppScreenChecker():
+class AppScreenChecker:
   """Checks that the current app screen matches an expected screen."""
 
   class Outcome(enum.IntEnum):

android_env/components/app_screen_checker_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2022 DeepMind Technologies Limited.
+# Copyright 2024 DeepMind Technologies Limited.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,7 +16,6 @@
 """Tests for android_env.components.app_screen_checker."""
 
 import re
-from typing import Sequence
 from unittest import mock
 
 from absl.testing import absltest
@@ -27,13 +26,13 @@ from android_env.proto import adb_pb2
 from android_env.proto import task_pb2
 
 
-def flatten_tree(tree: app_screen_checker._DumpsysNode,
-                 flat_tree: Sequence[str],
-                 indent: int = 2):
+def _flatten_tree(
+    tree: app_screen_checker._DumpsysNode, flat_tree: list[str], indent: int = 2
+):
   """Appends a list of strings to `flat_tree` from `tree`."""
   flat_tree.append(' ' * indent + tree.data)
   for c in tree.children:
-    flatten_tree(c, flat_tree, indent + 2)
+    _flatten_tree(c, flat_tree, indent + 2)
 
 
 class AppScreenCheckerTest(absltest.TestCase):
@@ -66,7 +65,7 @@ Queen Elizabeth II
 """
     tree = app_screen_checker.build_tree_from_dumpsys_output(dumpsys_output)
     flat_tree = []
-    flatten_tree(tree, flat_tree, indent=2)
+    _flatten_tree(tree, flat_tree, indent=2)
     self.assertEqual(flat_tree, [
         '  ___root___',
         '    Queen Elizabeth II',
@@ -116,7 +115,7 @@ Tree2
 """
     tree = app_screen_checker.build_tree_from_dumpsys_output(dumpsys_output)
     flat_tree = []
-    flatten_tree(tree, flat_tree, indent=2)
+    _flatten_tree(tree, flat_tree, indent=2)
     self.assertEqual(flat_tree, [
         '  ___root___',
         '    Tree1',