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

android_env/components/coordinator_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.
@@ -23,11 +23,14 @@ from absl.testing import absltest
 from absl.testing import parameterized
 from android_env.components import action_type
 from android_env.components import adb_call_parser
+from android_env.components import config_classes
 from android_env.components import coordinator as coordinator_lib
+from android_env.components import device_settings as device_settings_lib
 from android_env.components import errors
 from android_env.components import task_manager
 from android_env.components.simulators import base_simulator
 from android_env.proto import adb_pb2
+from android_env.proto import state_pb2
 from android_env.proto import task_pb2
 import dm_env
 import numpy as np
@@ -51,30 +54,50 @@ class CoordinatorTest(parameterized.TestCase):
             'AdbCallParser',
             autospec=True,
             return_value=self._adb_call_parser))
-    self.enter_context(mock.patch.object(time, 'sleep', autospec=True))
     self._coordinator = coordinator_lib.Coordinator(
         simulator=self._simulator,
         task_manager=self._task_manager,
-        num_fingers=1,
-        periodic_restart_time_min=0,
-        check_services_max_tries=0)
+        device_settings=device_settings_lib.DeviceSettings(self._simulator),
+    )
 
-  def test_restart_simulator(self):
-    restart_count = self._coordinator.stats()['restart_count']
-    self._coordinator._restart_simulator()
-    self.assertEqual(self._coordinator.stats()['restart_count'],
-                     restart_count + 1)
+  def tearDown(self):
+    super().tearDown()
+    self._coordinator.close()
 
-  def test_reset(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_relaunch_simulator(self, unused_mock_sleep):
+    relaunch_count = self._coordinator.stats()['relaunch_count']
+    self._coordinator._launch_simulator()
+    self.assertEqual(self._coordinator.stats()['relaunch_count'],
+                     relaunch_count + 1)
+
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_reset(self, unused_mock_sleep):
+    """'relaunch_count_execute_action' should be zero if there's no error."""
+
+    self._coordinator.rl_reset()
+    stats = self._coordinator.stats()
+    self.assertIn('relaunch_count_execute_action', stats)
+    self.assertEqual(stats['relaunch_count_execute_action'], 0)
+
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_reset_error_sending_action(self, unused_mock_sleep):
+    """'relaunch_count_execute_action' should be positive if there's an error."""
+
+    self._simulator.send_touch.side_effect = errors.SendActionError()
     self._coordinator.rl_reset()
+    stats = self._coordinator.stats()
+    self.assertIn('relaunch_count_execute_action', stats)
+    self.assertEqual(stats['relaunch_count_execute_action'], 1)
 
-  def test_lift_all_fingers(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_lift_all_fingers(self, unused_mock_sleep):
     self._coordinator = coordinator_lib.Coordinator(
         simulator=self._simulator,
         task_manager=self._task_manager,
-        num_fingers=3,
-        periodic_restart_time_min=0,
-        check_services_max_tries=0)
+        device_settings=device_settings_lib.DeviceSettings(self._simulator),
+        config=config_classes.CoordinatorConfig(num_fingers=3),
+    )
     self._coordinator.rl_reset()
     expected_actions = [
         # (x, y, is_down, identifier).
@@ -86,7 +109,8 @@ class CoordinatorTest(parameterized.TestCase):
     for actual, expected in zip(actual_actions, expected_actions):
       np.testing.assert_array_equal(actual, expected)
 
-  def test_process_action(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_process_action(self, unused_mock_sleep):
 
     def fake_rl_step(simulator_signals):
       return dm_env.transition(
@@ -114,7 +138,8 @@ class CoordinatorTest(parameterized.TestCase):
     self.assertEqual(obs['extras'], {'extra': [0.0]})
     self.assertFalse(timestep.last())
 
-  def test_process_action_error(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_process_action_error(self, unused_mock_sleep):
 
     def fake_rl_step(simulator_signals):
       self.assertFalse(simulator_signals['simulator_healthy'])
@@ -131,7 +156,8 @@ class CoordinatorTest(parameterized.TestCase):
     self.assertEqual(timestep.reward, 0.0)
     self.assertTrue(timestep.last())
 
-  def test_execute_action_touch(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_execute_action_touch(self, unused_mock_sleep):
 
     def fake_rl_step(simulator_signals):
       return dm_env.transition(
@@ -156,13 +182,14 @@ class CoordinatorTest(parameterized.TestCase):
                             self._random_screenshot)
     self._simulator.send_touch.assert_called_once_with([(300, 400, True, 0)])
 
-  def test_execute_multitouch_action(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_execute_multitouch_action(self, unused_mock_sleep):
     self._coordinator = coordinator_lib.Coordinator(
         simulator=self._simulator,
         task_manager=self._task_manager,
-        num_fingers=3,
-        periodic_restart_time_min=0,
-        check_services_max_tries=0)
+        device_settings=device_settings_lib.DeviceSettings(self._simulator),
+        config=config_classes.CoordinatorConfig(num_fingers=3),
+    )
 
     def fake_rl_step(simulator_signals):
       return dm_env.transition(
@@ -193,8 +220,8 @@ class CoordinatorTest(parameterized.TestCase):
     np.testing.assert_equal(timestep.observation['pixels'],
                             self._random_screenshot)
 
-  def test_execute_action_repeat(self):
-
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_execute_action_repeat(self, unused_mock_sleep):
     def fake_rl_step(simulator_signals):
       return dm_env.transition(
           reward=10.0,
@@ -214,8 +241,8 @@ class CoordinatorTest(parameterized.TestCase):
     np.testing.assert_equal(timestep.observation['pixels'],
                             self._random_screenshot)
 
-  def test_execute_action_error(self):
-
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_execute_action_error(self, unused_mock_sleep):
     def fake_rl_step(simulator_signals):
       self.assertFalse(simulator_signals['simulator_healthy'])
       return dm_env.truncation(reward=0.0, observation=None)
@@ -228,16 +255,18 @@ class CoordinatorTest(parameterized.TestCase):
     })
     self.assertIsNone(timestep.observation)
 
-  def test_max_restarts_setup_steps(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_max_restarts_setup_steps(self, unused_mock_sleep):
     init_fn_call = self._task_manager.setup_task.call_count
     self._task_manager.setup_task.side_effect = errors.StepCommandError
     self.assertRaises(errors.TooManyRestartsError,
-                      self._coordinator._restart_simulator)
-    # The method was called three more times when attempting to restart.
+                      self._coordinator._launch_simulator)
+    # The method was called three more times when attempting to relaunch.
     self.assertEqual(init_fn_call + 3,
                      self._task_manager.setup_task.call_count)
 
-  def test_execute_adb_call(self):
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_execute_adb_call(self, unused_mock_sleep):
     call = adb_pb2.AdbRequest(
         force_stop=adb_pb2.AdbRequest.ForceStop(package_name='blah'))
     expected_response = adb_pb2.AdbResponse(
@@ -249,175 +278,6 @@ class CoordinatorTest(parameterized.TestCase):
     self.assertEqual(response, expected_response)
     self._adb_call_parser.parse.assert_called_with(call)
 
-  @mock.patch.object(tempfile, 'gettempdir', autospec=True)
-  def test_with_tmp_dir_no_tempfile_call(self, mock_gettempdir):
-    """If passing a `tmp_dir`, `tempfile.gettempdir()` should not be called."""
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        periodic_restart_time_min=0,
-        tmp_dir=absltest.get_default_test_tmpdir(),
-        check_services_max_tries=0)
-    mock_gettempdir.assert_not_called()
-
-  @mock.patch.object(tempfile, 'gettempdir', autospec=True)
-  def test_no_tmp_dir_calls_tempfile(self, mock_gettempdir):
-    """If not passing a `tmp_dir`, `tempfile.gettempdir()` should be called."""
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        periodic_restart_time_min=0,
-        check_services_max_tries=0)
-    mock_gettempdir.assert_called_once()
-
-  @parameterized.parameters(
-      (True, '1'),
-      (False, '0'),
-  )
-  def test_touch_indicator(self, show, expected_value):
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        check_services_max_tries=0,
-        show_touches=show)
-    self._adb_call_parser.parse.assert_any_call(
-        adb_pb2.AdbRequest(
-            settings=adb_pb2.AdbRequest.SettingsRequest(
-                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
-                put=adb_pb2.AdbRequest.SettingsRequest.Put(
-                    key='show_touches', value=expected_value))))
-
-  @parameterized.parameters(
-      (True, '1'),
-      (False, '0'),
-  )
-  def test_pointer_location(self, show, expected_value):
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        check_services_max_tries=0,
-        show_pointer_location=show)
-    self._adb_call_parser.parse.assert_any_call(
-        adb_pb2.AdbRequest(
-            settings=adb_pb2.AdbRequest.SettingsRequest(
-                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
-                put=adb_pb2.AdbRequest.SettingsRequest.Put(
-                    key='pointer_location', value=expected_value))))
-
-  @parameterized.parameters(
-      (True, True, 'null*'),
-      (True, False, 'immersive.status=*'),
-      (False, True, 'immersive.navigation=*'),
-      (False, False, 'immersive.full=*'),
-      (None, None, 'immersive.full=*'),  # Defaults to hiding both.
-  )
-  def test_bar_visibility(self, show_navigation_bar, show_status_bar,
-                          expected_value):
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator,
-        task_manager=self._task_manager,
-        check_services_max_tries=0,
-        show_navigation_bar=show_navigation_bar,
-        show_status_bar=show_status_bar)
-    self._adb_call_parser.parse.assert_any_call(
-        adb_pb2.AdbRequest(
-            settings=adb_pb2.AdbRequest.SettingsRequest(
-                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,
-                put=adb_pb2.AdbRequest.SettingsRequest.Put(
-                    key='policy_control', value=expected_value))))
-
-  def test_wait_for_device(self):
-    """Ensures that all required services are up and running."""
-
-    def _parse_fn(adb_call: adb_pb2.AdbRequest):
-      _parse_fn.call_count += 1
-      if adb_call.generic.args:
-        # Initially, return "not found".
-        if _parse_fn.call_count < 3:
-          return adb_pb2.AdbResponse(
-              generic=adb_pb2.AdbResponse.GenericResponse(
-                  output=b'Service window: NOT found'))
-
-        # Return "found" after 2 calls.
-        service = adb_call.generic.args[3]
-        return adb_pb2.AdbResponse(
-            status=adb_pb2.AdbResponse.Status.OK,
-            generic=adb_pb2.AdbResponse.GenericResponse(
-                output=f'Service {service}: found'.encode('utf-8')))
-
-    _parse_fn.call_count = 0
-
-    self._adb_call_parser.parse.side_effect = _parse_fn
-
-    _ = coordinator_lib.Coordinator(
-        simulator=self._simulator, task_manager=self._task_manager)
-    self._adb_call_parser.parse.assert_has_calls([
-        mock.call(
-            adb_pb2.AdbRequest(
-                generic=adb_pb2.AdbRequest.GenericRequest(
-                    args=['shell', 'service', 'check', 'window']))),
-        mock.call(
-            adb_pb2.AdbRequest(
-                generic=adb_pb2.AdbRequest.GenericRequest(
-                    args=['shell', 'service', 'check', 'package']))),
-        mock.call(
-            adb_pb2.AdbRequest(
-                generic=adb_pb2.AdbRequest.GenericRequest(
-                    args=['shell', 'service', 'check', 'input']))),
-        mock.call(
-            adb_pb2.AdbRequest(
-                generic=adb_pb2.AdbRequest.GenericRequest(
-                    args=['shell', 'service', 'check', 'display']))),
-    ],
-                                                 any_order=True)
-
-  def test_wait_for_device_exceeded_max_tries(self):
-    """If the required services are not found, it should raise an exception."""
-
-    def _parse_fn(adb_call: adb_pb2.AdbRequest):
-      if adb_call.generic.args:
-        # Return "not found" repeatedly.
-        return adb_pb2.AdbResponse(
-            generic=adb_pb2.AdbResponse.GenericResponse(
-                output=b'Service window: NOT found'))
-
-    self._adb_call_parser.parse.side_effect = _parse_fn
-
-    self.assertRaises(
-        errors.AdbControllerDeviceTimeoutError,
-        coordinator_lib.Coordinator,
-        simulator=self._simulator,
-        task_manager=self._task_manager)
-
-  def test_update_task_succeeds(self):
-    task = task_pb2.Task(id='my_task')
-    stop_call_count = self._task_manager.stop_task.call_count
-    setup_call_count = self._task_manager.setup_task.call_count
-    success = self._coordinator.update_task(task)
-    self.assertEqual(1,
-                     self._task_manager.stop_task.call_count - stop_call_count)
-    self.assertEqual(
-        1, self._task_manager.setup_task.call_count - setup_call_count)
-    self._task_manager.update_task.assert_called_once_with(task)
-    self.assertTrue(success)
-    self._task_manager.stats.return_value = {}
-    self.assertEqual(0, self._coordinator.stats()['failed_task_updates'])
-
-  def test_update_task_fails(self):
-    task = task_pb2.Task(id='my_task')
-    self._task_manager.setup_task.side_effect = errors.StepCommandError
-    stop_call_count = self._task_manager.stop_task.call_count
-    setup_call_count = self._task_manager.setup_task.call_count
-    success = self._coordinator.update_task(task)
-    self.assertEqual(1,
-                     self._task_manager.stop_task.call_count - stop_call_count)
-    self.assertEqual(
-        1, self._task_manager.setup_task.call_count - setup_call_count)
-    self._task_manager.update_task.assert_called_once_with(task)
-    self.assertFalse(success)
-    self._task_manager.stats.return_value = {}
-    self.assertEqual(1, self._coordinator.stats()['failed_task_updates'])
-
 
 if __name__ == '__main__':
   absltest.main()

android_env/components/device_settings.py

@@ -0,0 +1,174 @@
+# coding=utf-8
+# 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Sets and gets some global settings on an Android device."""
+
+from typing import Final
+from unittest import mock
+
+from absl import logging
+from android_env.components import adb_call_parser
+from android_env.components import config_classes
+from android_env.components.simulators import base_simulator
+from android_env.proto import adb_pb2
+import numpy as np
+
+
+# The internal `AdbCallParser` instance is lazily instantiated within
+# `DeviceSettings`. If we make it optional (i.e. `| None`), pytype will think
+# that it could be `None`, requiring either explicit runtime checks or escape
+# hatches in every actual call, even if it's never actually `None` if reached
+# via the public API.
+# The trick here is to create this dummy instance of the right type that's used
+# as a sentinel to indicate that it hasn't been initialized yet.
+_PLACEHOLDER_ADB_CALL_PARSER: Final[adb_call_parser.AdbCallParser] = (
+    mock.create_autospec(adb_call_parser.AdbCallParser)
+)
+
+
+class DeviceSettings:
+  """An abstraction for general properties and settings of an Android device."""
+
+  def __init__(self, simulator: base_simulator.BaseSimulator):
+    self._simulator = simulator
+    self._adb_call_parser = _PLACEHOLDER_ADB_CALL_PARSER
+
+    # The size of the device screen in pixels.
+    self._screen_width: int = 0
+    self._screen_height: int = 0
+    # The device orientation.
+    self._orientation = np.zeros(4, dtype=np.uint8)
+
+  def update(self, config: config_classes.DeviceSettingsConfig) -> None:
+    """Sets the configuration of the device according to `config`."""
+
+    if self._adb_call_parser is _PLACEHOLDER_ADB_CALL_PARSER:
+      self._adb_call_parser = adb_call_parser.AdbCallParser(
+          adb_controller=self._simulator.create_adb_controller()
+      )
+
+    self._update_screen_size()
+    self._set_show_touches(config.show_touches)
+    self._set_show_pointer_location(config.show_pointer_location)
+    self._set_status_navigation_bars(
+        config.show_navigation_bar, config.show_status_bar
+    )
+
+  def screen_width(self) -> int:
+    """The screen width in pixels. Only valid after `update()` is called."""
+
+    return self._screen_width
+
+  def screen_height(self) -> int:
+    """The screen height in pixels. Only valid after `update()` is called."""
+
+    return self._screen_height
+
+  def get_orientation(self) -> np.ndarray:
+    """Returns the device orientation. Please see specs.py for details."""
+
+    if self._adb_call_parser is _PLACEHOLDER_ADB_CALL_PARSER:
+      self._adb_call_parser = adb_call_parser.AdbCallParser(
+          adb_controller=self._simulator.create_adb_controller()
+      )
+
+    self._update_orientation()
+    return self._orientation
+
+  def _update_screen_size(self) -> None:
+    """Sets the screen size from a screenshot ignoring the color channel."""
+
+    screenshot = self._simulator.get_screenshot()
+    self._screen_height = screenshot.shape[0]
+    self._screen_width = screenshot.shape[1]
+
+  def _set_show_touches(self, show: bool) -> None:
+    """Whether to display circles indicating the touch position."""
+
+    self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            settings=adb_pb2.AdbRequest.SettingsRequest(
+                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
+                put=adb_pb2.AdbRequest.SettingsRequest.Put(
+                    key='show_touches', value='1' if show else '0'
+                ),
+            )
+        )
+    )
+
+  def _set_show_pointer_location(self, show: bool) -> None:
+    """Whether to display blue lines on the screen indicating touch position."""
+
+    self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            settings=adb_pb2.AdbRequest.SettingsRequest(
+                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.SYSTEM,
+                put=adb_pb2.AdbRequest.SettingsRequest.Put(
+                    key='pointer_location', value='1' if show else '0'
+                ),
+            )
+        )
+    )
+
+  def _set_status_navigation_bars(
+      self, show_navigation: bool, show_status: bool
+  ) -> None:
+    """Whether to display the status (top) and navigation (bottom) bars."""
+
+    if show_navigation and show_status:
+      policy_control_value = 'null*'
+    elif show_navigation and not show_status:
+      policy_control_value = 'immersive.status=*'
+    elif not show_navigation and show_status:
+      policy_control_value = 'immersive.navigation=*'
+    else:
+      policy_control_value = 'immersive.full=*'
+
+    self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            settings=adb_pb2.AdbRequest.SettingsRequest(
+                name_space=adb_pb2.AdbRequest.SettingsRequest.Namespace.GLOBAL,
+                put=adb_pb2.AdbRequest.SettingsRequest.Put(
+                    key='policy_control', value=policy_control_value
+                ),
+            )
+        )
+    )
+
+  def _update_orientation(self) -> None:
+    """Updates the current device orientation."""
+
+    # Skip fetching the orientation if we already have it.
+    if not np.all(self._orientation == np.zeros(4)):
+      return
+
+    orientation_response = self._adb_call_parser.parse(
+        adb_pb2.AdbRequest(
+            get_orientation=adb_pb2.AdbRequest.GetOrientationRequest()
+        )
+    )
+    if orientation_response.status != adb_pb2.AdbResponse.Status.OK:
+      logging.error('Got bad orientation: %r', orientation_response)
+      return
+
+    orientation = orientation_response.get_orientation.orientation
+    if orientation not in {0, 1, 2, 3}:
+      logging.error('Got bad orientation: %r', orientation)
+      return
+
+    # Transform into one-hot format.
+    orientation_onehot = np.zeros([4], dtype=np.uint8)
+    orientation_onehot[orientation] = 1
+    self._orientation = orientation_onehot