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

android_env/components/device_settings_test.py

@@ -0,0 +1,228 @@
+# 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.
+
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from android_env.components import config_classes
+from android_env.components import device_settings as device_settings_lib
+from android_env.components.simulators import base_simulator
+import numpy as np
+
+
+class DeviceSettingsTest(parameterized.TestCase):
+
+  def test_screen_size_before_update(self):
+    """The screen size should be 0x0 without calling `update()`."""
+
+    # Arrange.
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
+    device_settings = device_settings_lib.DeviceSettings(simulator)
+
+    # Act.
+    height = device_settings.screen_height()
+    width = device_settings.screen_width()
+
+    # Assert.
+    self.assertEqual(height, 0)
+    self.assertEqual(width, 0)
+
+  def test_screen_size_after_update(self):
+    """The screen size should be set after calling `update()`."""
+
+    # Arrange.
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
+    simulator.get_screenshot.return_value = np.random.randint(
+        low=0, high=255, size=(123, 456, 3), dtype=np.uint8
+    )
+    adb_controller = simulator.create_adb_controller.return_value
+    adb_controller.execute_command.return_value = b''
+    device_settings = device_settings_lib.DeviceSettings(simulator)
+
+    # Act.
+    device_settings.update(config_classes.DeviceSettingsConfig())
+    height = device_settings.screen_height()
+    width = device_settings.screen_width()
+
+    # Assert.
+    self.assertEqual(height, 123)
+    self.assertEqual(width, 456)
+
+  @parameterized.named_parameters(
+      (
+          'show_touches',
+          config_classes.DeviceSettingsConfig(show_touches=True),
+          mock.call(
+              ['shell', 'settings', 'put', 'system', 'show_touches', '1'],
+              timeout=None,
+          ),
+      ),
+      (
+          'show_touches_false',
+          config_classes.DeviceSettingsConfig(show_touches=False),
+          mock.call(
+              ['shell', 'settings', 'put', 'system', 'show_touches', '0'],
+              timeout=None,
+          ),
+      ),
+      (
+          'show_pointer_location',
+          config_classes.DeviceSettingsConfig(show_pointer_location=True),
+          mock.call(
+              ['shell', 'settings', 'put', 'system', 'pointer_location', '1'],
+              timeout=None,
+          ),
+      ),
+      (
+          'show_pointer_location_false',
+          config_classes.DeviceSettingsConfig(show_pointer_location=False),
+          mock.call(
+              ['shell', 'settings', 'put', 'system', 'pointer_location', '0'],
+              timeout=None,
+          ),
+      ),
+      (
+          'show_navigation_and_status',
+          config_classes.DeviceSettingsConfig(
+              show_navigation_bar=True, show_status_bar=True
+          ),
+          mock.call(
+              ['shell', 'settings', 'put', 'global', 'policy_control', 'null*'],
+              timeout=None,
+          ),
+      ),
+      (
+          'show_navigation_and_no_status',
+          config_classes.DeviceSettingsConfig(
+              show_navigation_bar=True, show_status_bar=False
+          ),
+          mock.call(
+              [
+                  'shell',
+                  'settings',
+                  'put',
+                  'global',
+                  'policy_control',
+                  'immersive.status=*',
+              ],
+              timeout=None,
+          ),
+      ),
+      (
+          'show_no_navigation_and_status',
+          config_classes.DeviceSettingsConfig(
+              show_navigation_bar=False, show_status_bar=True
+          ),
+          mock.call(
+              [
+                  'shell',
+                  'settings',
+                  'put',
+                  'global',
+                  'policy_control',
+                  'immersive.navigation=*',
+              ],
+              timeout=None,
+          ),
+      ),
+      (
+          'show_no_navigation_and_no_status',
+          config_classes.DeviceSettingsConfig(
+              show_navigation_bar=False, show_status_bar=False
+          ),
+          mock.call(
+              [
+                  'shell',
+                  'settings',
+                  'put',
+                  'global',
+                  'policy_control',
+                  'immersive.full=*',
+              ],
+              timeout=None,
+          ),
+      ),
+  )
+  def test_update(
+      self, settings: config_classes.DeviceSettingsConfig, expected_call
+  ):
+    """We expect the right call for each setting."""
+
+    # Arrange.
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
+    adb_controller = simulator.create_adb_controller.return_value
+    adb_controller.execute_command.return_value = b''
+    device_settings = device_settings_lib.DeviceSettings(simulator)
+
+    # Act.
+    device_settings.update(settings)
+
+    # Assert.
+    adb_controller.execute_command.assert_has_calls(
+        [expected_call], any_order=True
+    )
+
+  def test_get_orientation_bad_response(self):
+    """The orientation should be unset if the underlying response is bad."""
+
+    # Arrange.
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
+    adb_controller = simulator.create_adb_controller.return_value
+    adb_controller.execute_command.return_value = b''
+    device_settings = device_settings_lib.DeviceSettings(simulator)
+
+    # Act.
+    orientation = device_settings.get_orientation()
+
+    # Assert.
+    np.testing.assert_array_equal(orientation, np.zeros(4))
+
+  def test_get_orientation_bad_orientation(self):
+    """The orientation should be unset if the underlying orientation is bad."""
+
+    # Arrange.
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
+    adb_controller = simulator.create_adb_controller.return_value
+    adb_controller.execute_command.return_value = b' InputDeviceOrientation: 9'
+    device_settings = device_settings_lib.DeviceSettings(simulator)
+
+    # Act.
+    orientation = device_settings.get_orientation()
+
+    # Assert.
+    np.testing.assert_array_equal(orientation, np.zeros(4))
+
+  def test_get_orientation_success(self):
+    """Checks that the orientation comes back as expected."""
+
+    # Arrange.
+    simulator = mock.create_autospec(base_simulator.BaseSimulator)
+    adb_controller = simulator.create_adb_controller.return_value
+    adb_controller.execute_command.return_value = b' InputDeviceOrientation: 3'
+    device_settings = device_settings_lib.DeviceSettings(simulator)
+
+    # Act.
+    orientation = device_settings.get_orientation()
+    # The output should be idempotent if the underlying system did not change.
+    orientation_again = device_settings.get_orientation()
+
+    # Assert.
+    np.testing.assert_array_equal(orientation, np.array([0, 0, 0, 1]))
+    np.testing.assert_array_equal(orientation, orientation_again)
+
+
+if __name__ == '__main__':
+  absltest.main()

android_env/components/dumpsys_thread.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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 @@
 """A ThreadFunction that runs and parses adb dumpsys."""
 
 import concurrent.futures
-from typing import Optional
 
 from absl import logging
 from android_env.components import app_screen_checker as app_screen_checker_lib
@@ -56,9 +55,9 @@ class DumpsysThread:
     self._check_frequency = check_frequency
     self._max_failed_activity_extraction = max_failed_current_activity
     self._num_failed_activity_extraction = 0
-    self._latest_check: Optional[concurrent.futures.Future] = None
+    self._latest_check: concurrent.futures.Future | None = None
 
-  def check_user_exited(self, timeout: Optional[float] = None) -> bool:
+  def check_user_exited(self, timeout: float | None = None) -> bool:
     """Returns True if the user is not in the expected screen.
 
     Args:

android_env/components/dumpsys_thread_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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.

android_env/components/errors.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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,8 +15,6 @@
 
 """Definitions of exceptions used by AndroidEnv."""
 
-from typing import Optional
-
 
 class AndroidEnvError(Exception):
   """Base class for all known errors generated by AndroidEnv."""
@@ -80,7 +78,7 @@ class CheckInstallError(StepCommandError):
   ERROR_CODE = 9
 
 
-def from_code(code: int, msg: str = '') -> Optional[AndroidEnvError]:
+def from_code(code: int, msg: str = '') -> AndroidEnvError | None:
   """Returns an AndroidEnvError instance from the given arguments."""
 
   code_to_error = {
@@ -98,4 +96,3 @@ def from_code(code: int, msg: str = '') -> Optional[AndroidEnvError]:
 
   if code in code_to_error:
     return code_to_error[code](msg)
-

android_env/components/errors_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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.

android_env/components/log_stream.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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 @@
 """Abstract class for handling a stream of logs from a simulator."""
 
 import abc
+from collections.abc import Generator, Sequence
 import threading
-from typing import Generator, Sequence
 from absl import logging
 
 

android_env/components/log_stream_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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.

android_env/components/logcat_thread.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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,20 +15,21 @@
 
 """A class that launches a thread to read Android log outputs."""
 
+from collections.abc import Callable
+import dataclasses
 import re
 import threading
-# `typing.Pattern` has been deprecated in Python 3.9 in favor of `re.Pattern`,
-# but it is not available even in slightly older Python versions.
-# Please see https://www.python.org/dev/peps/pep-0585/
-from typing import Callable, Match, NamedTuple, Pattern
 
 from absl import logging
 from android_env.components import log_stream as log_stream_lib
 
 
-class EventListener(NamedTuple):
-  regexp: Pattern[str]
-  handler_fn: Callable[[Pattern[str], Match[str]], None]
+@dataclasses.dataclass
+class EventListener:
+  """A function that's called when an event is triggered."""
+
+  regexp: re.Pattern[str]
+  handler_fn: Callable[[re.Pattern[str], re.Match[str]], None]
 
 
 class LogcatThread:

android_env/components/logcat_thread_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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.
@@ -17,7 +17,6 @@
 
 import re
 import threading
-from typing import Match, Pattern
 
 from absl.testing import absltest
 from android_env.components import log_stream
@@ -108,7 +107,7 @@ class LogcatThreadTest(absltest.TestCase):
     # Set up a listener that modifies an arbitrary state.
     some_state = threading.Event()
 
-    def my_handler(event: Pattern[str], match: Match[str]):
+    def my_handler(event: re.Pattern[str], match: re.Match[str]):
       del event, match
       nonlocal some_state
       some_state.set()

android_env/components/{utils.py → pixel_fns.py}

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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,11 +15,10 @@
 
 """Utils for AndroidEnv."""
 
-from typing import Sequence
+from collections.abc import Sequence
 
 from dm_env import specs
 import numpy as np
-import numpy.typing as np_typing
 
 
 def touch_position_to_pixel_position(
@@ -39,27 +38,27 @@ def transpose_pixels(frame: np.ndarray) -> np.ndarray:
 
 def orient_pixels(frame: np.ndarray, orientation: int) -> np.ndarray:
   """Rotates screen pixels according to the given orientation."""
-  if orientation == 0:  # PORTRAIT_90
-    return frame
-  elif orientation == 1:  # LANDSCAPE_90
-    return np.rot90(frame, k=3, axes=(0, 1))
-  elif orientation == 2:  # PORTRAIT_180
-    return np.rot90(frame, k=2, axes=(0, 1))
-  elif orientation == 3:  # LANDSCAPE_270
-    return np.rot90(frame, k=1, axes=(0, 1))
-  else:
-    raise ValueError(
-        'Orientation must be an integer in [0, 3] but is %r' % orientation)
+
+  match orientation:
+    case 0:  # PORTRAIT_90
+      return frame
+    case 1:  # LANDSCAPE_90
+      return np.rot90(frame, k=3, axes=(0, 1))
+    case 2:  # PORTRAIT_180
+      return np.rot90(frame, k=2, axes=(0, 1))
+    case 3:  # LANDSCAPE_270
+      return np.rot90(frame, k=1, axes=(0, 1))
+    case _:
+      raise ValueError(
+          'Orientation must be an integer in [0, 3] but is %r' % orientation
+      )
 
 
-def convert_int_to_float(data: np.ndarray,
-                         data_spec: specs.Array,
-                         float_type: np_typing.DTypeLike = np.float32):
+def convert_int_to_float(data: np.ndarray, data_spec: specs.Array):
   """Converts an array of int values to floats between 0 and 1."""
+
   if not np.issubdtype(data.dtype, np.integer):
     raise TypeError(f'{data.dtype} is not an integer type')
-  if not np.issubdtype(float_type, np.floating):
-    raise TypeError(f'{float_type} is not a floating-point type')
   if isinstance(data_spec, specs.BoundedArray):
     value_min = data_spec.minimum
     value_max = data_spec.maximum
@@ -68,4 +67,4 @@ def convert_int_to_float(data: np.ndarray,
     iinfo = np.iinfo(data_spec.dtype)
     value_min = iinfo.min
     value_max = iinfo.max
-  return float_type(1.0 * (data - value_min) / (value_max - value_min))
+  return np.float32(1.0 * (data - value_min) / (value_max - value_min))

android_env/components/{utils_test.py → pixel_fns_test.py}

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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.
@@ -13,11 +13,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for android_env.components.utils."""
+"""Tests for pixel_fns."""
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from android_env.components import utils
+from android_env.components import pixel_fns
 from dm_env import specs
 import numpy as np
 
@@ -32,38 +32,42 @@ class UtilsTest(parameterized.TestCase):
       )
   def test_touch_position_to_pixel_position(
       self, touch_pos, width_height, pixel_pos):
-    self.assertEqual(utils.touch_position_to_pixel_position(
-        np.array(touch_pos), width_height), pixel_pos)
+    self.assertEqual(
+        pixel_fns.touch_position_to_pixel_position(
+            np.array(touch_pos), width_height
+        ),
+        pixel_pos,
+    )
 
   def test_transpose_pixels(self):
     image = np.reshape(np.array(range(12)), (3, 2, 2))
     expected = [[[0, 1], [4, 5], [8, 9]], [[2, 3], [6, 7], [10, 11]]]
-    self.assertEqual(utils.transpose_pixels(image).shape, (2, 3, 2))
-    self.assertTrue((utils.transpose_pixels(image) == expected).all())
+    self.assertEqual(pixel_fns.transpose_pixels(image).shape, (2, 3, 2))
+    self.assertTrue((pixel_fns.transpose_pixels(image) == expected).all())
 
   def test_orient_pixels(self):
     image = np.reshape(np.array(range(12)), (3, 2, 2))
 
     expected_90 = [[[8, 9], [4, 5], [0, 1]], [[10, 11], [6, 7], [2, 3]]]
     rot_90 = 1  # LANDSCAPE_90
-    rotated = utils.orient_pixels(image, rot_90)
+    rotated = pixel_fns.orient_pixels(image, rot_90)
     self.assertEqual(rotated.shape, (2, 3, 2))
     self.assertTrue((rotated == expected_90).all())
 
     expected_180 = [[[10, 11], [8, 9]], [[6, 7], [4, 5]], [[2, 3], [0, 1]]]
     rot_180 = 2  # PORTRAIT_180
-    rotated = utils.orient_pixels(image, rot_180)
+    rotated = pixel_fns.orient_pixels(image, rot_180)
     self.assertEqual(rotated.shape, (3, 2, 2))
     self.assertTrue((rotated == expected_180).all())
 
     expected_270 = [[[2, 3], [6, 7], [10, 11]], [[0, 1], [4, 5], [8, 9]]]
     rot_270 = 3  # LANDSCAPE_270
-    rotated = utils.orient_pixels(image, rot_270)
+    rotated = pixel_fns.orient_pixels(image, rot_270)
     self.assertEqual(rotated.shape, (2, 3, 2))
     self.assertTrue((rotated == expected_270).all())
 
     rot_0 = 0  # PORTRAIT_0
-    rotated = utils.orient_pixels(image, rot_0)
+    rotated = pixel_fns.orient_pixels(image, rot_0)
     self.assertEqual(rotated.shape, (3, 2, 2))
     self.assertTrue((rotated == image).all())
 
@@ -75,15 +79,16 @@ class UtilsTest(parameterized.TestCase):
         maximum=[5, 5, 20, 2],
         name='bounded_array')
     data = np.array([2, 2, 10, 0], dtype=np.int32)
-    float_data = utils.convert_int_to_float(data, spec, np.float64)
+    float_data = pixel_fns.convert_int_to_float(data, spec)
     np.testing.assert_equal(
-        np.array([2. / 5., 1. / 4., 0., 0.5], dtype=np.float64), float_data)
+        np.array([2.0 / 5.0, 1.0 / 4.0, 0.0, 0.5], dtype=np.float32), float_data
+    )
 
   def test_convert_int_to_float_bounded_array_broadcast(self):
     spec = specs.BoundedArray(
         shape=(3,), dtype=np.int16, minimum=2, maximum=4, name='bounded_array')
     data = np.array([2, 3, 4], dtype=np.int16)
-    float_data = utils.convert_int_to_float(data, spec, np.float32)
+    float_data = pixel_fns.convert_int_to_float(data, spec)
     np.testing.assert_equal(
         np.array([0.0, 0.5, 1.0], dtype=np.float32), float_data)
 
@@ -93,7 +98,7 @@ class UtilsTest(parameterized.TestCase):
         dtype=np.int8,  # int8 implies min=-128, max=127
         name='bounded_array')
     data = np.array([-128, 0, 127], dtype=np.int16)
-    float_data = utils.convert_int_to_float(data, spec, np.float32)
+    float_data = pixel_fns.convert_int_to_float(data, spec)
     np.testing.assert_equal(
         np.array([0.0, 128. / 255., 1.0], dtype=np.float32), float_data)
 

android_env/components/setup_step_interpreter.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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,9 +15,10 @@
 
 """A component that parses and processes SetupSteps."""
 
+from collections.abc import Sequence
 import copy
 import time
-from typing import Any, Optional, Sequence
+from typing import Any
 
 from absl import logging
 from android_env.components import adb_call_parser as adb_call_parser_lib
@@ -110,45 +111,52 @@ class SetupStepInterpreter:
         f'Step failed: [{step_cmd}]') from latest_error
 
   def _execute_step_cmd(
-      self, step_cmd: task_pb2.SetupStep,
-      step_type: Optional[str]) -> Optional[adb_pb2.AdbResponse]:
+      self, step_cmd: task_pb2.SetupStep, step_type: str | None
+  ) -> adb_pb2.AdbResponse | None:
     """Executes a step command of given type."""
-    if not step_type:
-      return
 
-    if step_type == 'sleep':
-      time.sleep(step_cmd.sleep.time_sec)
-    elif step_type == 'adb_request':
-      response = self._adb_call_parser.parse(step_cmd.adb_request)
-      if response.status != adb_pb2.AdbResponse.Status.OK:
-        raise errors.AdbControllerError(
-            f'Failed to execute AdbRequest [{step_cmd.adb_request}].\n'
-            f'Status: {response.status}\n'
-            f'Error: {response.error_message}')
-      return response
-    else:
-      raise NotImplementedError('No step command of type [%s].' % step_type)
-
-  def _check_success(self,
-                     success_check: Optional[str],
-                     success_condition: task_pb2.SuccessCondition) -> None:
+    match step_type:
+      case None:
+        return None
+      case 'sleep':
+        time.sleep(step_cmd.sleep.time_sec)
+        return None
+      case 'adb_request':
+        response = self._adb_call_parser.parse(step_cmd.adb_request)
+        if response.status != adb_pb2.AdbResponse.Status.OK:
+          raise errors.AdbControllerError(
+              f'Failed to execute AdbRequest [{step_cmd.adb_request}].\n'
+              f'Status: {response.status}\n'
+              f'Error: {response.error_message}'
+          )
+        return response
+      case _:
+        raise NotImplementedError(f'No step command of type [{step_type}].')
+
+  def _check_success(
+      self,
+      success_check: str | None,
+      success_condition: task_pb2.SuccessCondition,
+  ) -> None:
     """Checks whether the given success condition was met."""
 
-    if not success_check:
-      return
-
-    if success_check == 'wait_for_app_screen':
-      wait_for_app_screen = success_condition.wait_for_app_screen
-      screen_checker = app_screen_checker.AppScreenChecker(
-          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)
-      self._stats['total_time_waiting_for_app_screen'] += wait_time
-    elif success_check == 'check_install':
-      self._check_install(success_condition.check_install)
-    else:
-      raise NotImplementedError('No success check called [%s].' % success_check)
+    match success_check:
+      case None:
+        return None
+      case 'wait_for_app_screen':
+        wait_for_app_screen = success_condition.wait_for_app_screen
+        screen_checker = app_screen_checker.AppScreenChecker(
+            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
+        )
+        self._stats['total_time_waiting_for_app_screen'] += wait_time
+      case 'check_install':
+        self._check_install(success_condition.check_install)
+      case _:
+        raise NotImplementedError(f'No success check called [{success_check}].')
 
   def _check_install(self, check_install: task_pb2.CheckInstall) -> None:
     """Checks that the given package is installed."""

android_env/components/setup_step_interpreter_test.py

@@ -1,5 +1,5 @@
 # coding=utf-8
-# Copyright 2023 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.