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

android_env/components/simulators/emulator/emulator_simulator_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,13 +23,17 @@ from unittest import mock
 from absl.testing import absltest
 from android_env.components import adb_call_parser
 from android_env.components import adb_controller
+from android_env.components import config_classes
 from android_env.components.simulators.emulator import emulator_launcher
 from android_env.components.simulators.emulator import emulator_simulator
+from android_env.proto import state_pb2
 import grpc
 from PIL import Image
+import portpicker
 
 from android_env.proto import emulator_controller_pb2
 from android_env.proto import emulator_controller_pb2_grpc
+from android_env.proto import snapshot_service_pb2
 
 
 class EmulatorSimulatorTest(absltest.TestCase):
@@ -47,21 +51,16 @@ class EmulatorSimulatorTest(absltest.TestCase):
 
     self._grpc_channel = mock.create_autospec(grpc.Channel)
     mock.patch.object(
-        grpc.aio, 'secure_channel',
-        return_value=self._grpc_channel).start()
+        grpc.aio, 'secure_channel', return_value=self._grpc_channel).start()
     mock.patch.object(
-        grpc, 'secure_channel',
-        return_value=self._grpc_channel).start()
+        grpc, 'secure_channel', return_value=self._grpc_channel).start()
     mock.patch.object(
         grpc, 'local_channel_credentials',
         return_value=self._grpc_channel).start()
     self._mock_future = mock.create_autospec(grpc.Future)
     mock.patch.object(
-        grpc, 'channel_ready_future',
-        return_value=self._mock_future).start()
-    mock.patch.object(
-        time, 'time',
-        return_value=12345).start()
+        grpc, 'channel_ready_future', return_value=self._mock_future).start()
+    mock.patch.object(time, 'time', return_value=12345).start()
 
     mock.patch.object(
         adb_controller, 'AdbController',
@@ -76,67 +75,110 @@ class EmulatorSimulatorTest(absltest.TestCase):
         return_value=self._launcher).start()
 
   def test_adb_device_name_not_empty(self):
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-        })
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
     self.assertNotEmpty(simulator.adb_device_name())
 
   @mock.patch.object(os.path, 'exists', autospec=True, return_value=True)
   @mock.patch.object(builtins, 'open', autospec=True)
   def test_logfile_path(self, mock_open, unused_mock_exists):
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
+    config = config_classes.EmulatorConfig(
         logfile_path='fake/logfile/path',
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-        })
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
     mock_open.return_value.__enter__.return_value.read.return_value = (
         'fake_logs'.encode('utf-8'))
     logs = simulator.get_logs()
     mock_open.assert_called_once_with('fake/logfile/path', 'rb')
     self.assertEqual(logs, 'fake_logs')
 
-  def test_launch(self):
-
-    # Make sure that adb_controller is started before Emulator is launched.
+  @mock.patch.object(portpicker, 'is_port_free', return_value=True)
+  def test_grpc_port(self, unused_mock_portpicker):
+
+    launcher_config = config_classes.EmulatorLauncherConfig(
+        tmp_dir=self.create_tempdir().full_path
+    )
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=launcher_config,
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
+    self.assertEqual(launcher_config.grpc_port, 8554)
+
+  @mock.patch.object(portpicker, 'is_port_free', return_value=False)
+  def test_grpc_port_unavailable(self, unused_mock_portpicker):
+
+    launcher_config = config_classes.EmulatorLauncherConfig(
+        tmp_dir=self.create_tempdir().full_path
+    )
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=launcher_config,
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
+    self.assertNotEqual(launcher_config.grpc_port, 8554)
+
+  def test_launch_operation_order(self):
+    """Makes sure that adb_controller is started before Emulator is launched."""
+
+    # Arrange.
     call_order = []
-    self._adb_controller.init_server.side_effect = (
-        lambda *a, **kw: call_order.append('init_server'))
+    self._adb_controller.init_server.side_effect = lambda: call_order.append(
+        'init_server'
+    )
     self._launcher.launch_emulator_process.side_effect = (
-        lambda *a, **kw: call_order.append('launch_emulator_process'))
-
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-            'prompt_regex': 'awesome>',
-        })
-
-    # The simulator should launch and not crash.
-    simulator.launch()
-
+        lambda: call_order.append('launch_emulator_process')
+    )
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
+
+    # Act.
+    simulator.launch()  # The simulator should launch and not crash.
+
+    # Assert.
+    # The adb server should be initialized before launching the emulator.
     self.assertEqual(call_order, ['init_server', 'launch_emulator_process'])
 
   def test_close(self):
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-        })
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
 
     # The simulator should launch and not crash.
     simulator.launch()
@@ -146,15 +188,36 @@ class EmulatorSimulatorTest(absltest.TestCase):
     # resources.
     simulator.close()
 
-  def test_restart(self):
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-        })
+  def test_value_error_if_launch_attempt_params_incorrect(self):
+    self.assertRaises(
+        ValueError,
+        emulator_simulator.EmulatorSimulator,
+        config=config_classes.EmulatorConfig(
+            emulator_launcher=config_classes.EmulatorLauncherConfig(
+                grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+            ),
+            adb_controller=config_classes.AdbControllerConfig(
+                adb_path='/my/adb',
+                adb_server_port=5037,
+            ),
+            launch_n_times_without_reboot=2,
+            launch_n_times_without_reinstall=1,
+        ),
+    )
+
+  def test_launch_attempt_reboot(self):
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+        launch_n_times_without_reboot=1,
+        launch_n_times_without_reinstall=2,
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
 
     # The simulator should launch and not crash.
     simulator.launch()
@@ -162,20 +225,82 @@ class EmulatorSimulatorTest(absltest.TestCase):
     self._launcher.launch_emulator_process.assert_called_once()
     self._launcher.reset_mock()
 
-    # For whatever reason clients may want to restart the EmulatorSimulator.
-    simulator.restart()
+    # Launch attempt 2.
+    simulator.launch()
     self._launcher.confirm_shutdown.assert_called_once()
+    self._launcher.close.assert_not_called()
+    self._launcher.launch_emulator_process.assert_called_once()
+
+  def test_launch_attempt_reinstall_after_zero_attempts(self):
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+        launch_n_times_without_reboot=0,
+        launch_n_times_without_reinstall=0,
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
+
+    # The simulator should not reboot or reinstall on its very first launch.
+    simulator.launch()
+    self._launcher.launch_emulator_process.assert_called_once()
+    self._launcher.confirm_shutdown.assert_not_called()
+    self._launcher.close.assert_not_called()
+
+    # Every subsequent attempt should reboot and reinstall.
+    self._launcher.reset_mock()
+    simulator.launch()
+    self._launcher.confirm_shutdown.assert_called_once()
+    self._launcher.close.assert_called_once()  # Now this should `close()`.
+    self._launcher.launch_emulator_process.assert_called_once()
+
+  def test_launch_attempt_reinstall(self):
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+        launch_n_times_without_reboot=1,
+        launch_n_times_without_reinstall=2,
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
+
+    # The simulator should launch and not crash.
+    simulator.launch()
+    self._launcher.launch_emulator_process.assert_called_once()
+
+    # Launch attempt 2.
+    self._launcher.reset_mock()
+    simulator.launch()
+    self._launcher.confirm_shutdown.assert_called_once()
+    self._launcher.close.assert_not_called()  # Reboots don't `close()`.
+    self._launcher.launch_emulator_process.assert_called_once()
+
+    # Launch attempt 3.
+    self._launcher.reset_mock()
+    simulator.launch()
+    self._launcher.confirm_shutdown.assert_called_once()
+    self._launcher.close.assert_called_once()  # Now this should `close()`.
     self._launcher.launch_emulator_process.assert_called_once()
 
   def test_get_screenshot(self):
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-        })
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
 
     # The simulator should launch and not crash.
     simulator.launch()
@@ -191,15 +316,113 @@ class EmulatorSimulatorTest(absltest.TestCase):
     # and it should have 3 channels (RGB).
     self.assertEqual(screenshot.shape, (1234, 5678, 3))
 
+  def test_load_state(self):
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
+
+    # The simulator should launch and not crash.
+    simulator.launch()
+
+    with mock.patch.object(
+        simulator, '_snapshot_stub', create_autospec=True
+    ) as mock_snapshot_stub:
+      snapshot_list = snapshot_service_pb2.SnapshotList()
+      snapshot_list.snapshots.add(snapshot_id='snapshot_name_foo')
+      snapshot_list.snapshots.add(snapshot_id='snapshot_name_bar')
+      mock_snapshot_stub.ListSnapshots.return_value = snapshot_list
+      mock_snapshot_stub.LoadSnapshot.return_value = (
+          snapshot_service_pb2.SnapshotPackage(success=True)
+      )
+      load_response = simulator.load_state(
+          request=state_pb2.LoadStateRequest(
+              args={'snapshot_name': 'snapshot_name_foo'}
+          )
+      )
+      self.assertEqual(
+          load_response.status, state_pb2.LoadStateResponse.Status.OK
+      )
+      load_response = simulator.load_state(
+          request=state_pb2.LoadStateRequest(
+              args={'snapshot_name': 'snapshot_name_baz'}
+          )
+      )
+      self.assertEqual(
+          load_response.status, state_pb2.LoadStateResponse.Status.NOT_FOUND
+      )
+      mock_snapshot_stub.LoadSnapshot.return_value = (
+          snapshot_service_pb2.SnapshotPackage(success=False, err=b'error')
+      )
+      load_response = simulator.load_state(
+          request=state_pb2.LoadStateRequest(
+              args={'snapshot_name': 'snapshot_name_bar'}
+          )
+      )
+      self.assertEqual(
+          load_response.status, state_pb2.LoadStateResponse.Status.ERROR
+      )
+      self.assertEqual(load_response.error_message, 'error')
+
+  def test_save_state(self):
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
+
+    # The simulator should launch and not crash.
+    simulator.launch()
+
+    with mock.patch.object(
+        simulator, '_snapshot_stub', create_autospec=True
+    ) as mock_snapshot_stub:
+      mock_snapshot_stub.SaveSnapshot.return_value = (
+          snapshot_service_pb2.SnapshotPackage(success=True)
+      )
+      save_response = simulator.save_state(
+          request=state_pb2.SaveStateRequest(
+              args={'snapshot_name': 'snapshot_name_foo'}
+          )
+      )
+      self.assertEqual(
+          save_response.status, state_pb2.SaveStateResponse.Status.OK
+      )
+      mock_snapshot_stub.SaveSnapshot.return_value = (
+          snapshot_service_pb2.SnapshotPackage(success=False, err=b'error')
+      )
+      save_response = simulator.save_state(
+          request=state_pb2.SaveStateRequest(
+              args={'snapshot_name': 'snapshot_name_bar'}
+          )
+      )
+      self.assertEqual(
+          save_response.status, state_pb2.SaveStateResponse.Status.ERROR
+      )
+      self.assertEqual(save_response.error_message, 'error')
+
   def test_send_touch(self):
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-        })
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
 
     # The simulator should launch and not crash.
     simulator.launch()
@@ -247,14 +470,16 @@ class EmulatorSimulatorTest(absltest.TestCase):
     ])
 
   def test_send_key(self):
-    tmp_dir = absltest.get_default_test_tmpdir()
-    simulator = emulator_simulator.EmulatorSimulator(
-        tmp_dir=tmp_dir,
-        emulator_launcher_args={'grpc_port': 1234},
-        adb_controller_args={
-            'adb_path': '/my/adb',
-            'adb_server_port': 5037,
-        })
+    config = config_classes.EmulatorConfig(
+        emulator_launcher=config_classes.EmulatorLauncherConfig(
+            grpc_port=1234, tmp_dir=self.create_tempdir().full_path
+        ),
+        adb_controller=config_classes.AdbControllerConfig(
+            adb_path='/my/adb',
+            adb_server_port=5037,
+        ),
+    )
+    simulator = emulator_simulator.EmulatorSimulator(config)
 
     # The simulator should launch and not crash.
     simulator.launch()

android_env/components/simulators/fake/__init__.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.

android_env/components/simulators/fake/fake_simulator.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,16 +18,16 @@
 import random
 import threading
 import time
-from typing import List, Optional, Tuple
 
 from absl import logging
 from android_env.components import adb_controller
+from android_env.components import config_classes
 from android_env.components import log_stream
 from android_env.components.simulators import base_simulator
 import numpy as np
 
 
-class FakeStream():
+class FakeStream:
   """This class simulates the logs coming from ADB."""
 
   def __init__(self):
@@ -77,12 +77,15 @@ class FakeLogStream(log_stream.LogStream):
 class FakeAdbController(adb_controller.AdbController):
   """Fake adb controller for FakeSimulator."""
 
-  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:
     """Returns fake output for adb commands."""
 
-    del timeout
+    del timeout, device_specific
 
     # Fake "service is ready" output.
     if args[:3] == ['shell', 'service', 'check']:
@@ -103,18 +106,10 @@ class FakeAdbController(adb_controller.AdbController):
 class FakeSimulator(base_simulator.BaseSimulator):
   """FakeSimulator class."""
 
-  def __init__(self,
-               screen_dimensions: Tuple[int, int] = (480, 320),
-               **kwargs):
-    """FakeSimulator class that can replace EmulatorSimulator in AndroidEnv.
-
-    Args:
-      screen_dimensions: desired screen dimensions in pixels. This determines
-        the shape of the screenshots returned by get_screenshot().
-      **kwargs: other keyword arguments for the base class.
-    """
-    super().__init__(**kwargs)
-    self._screen_dimensions = np.array(screen_dimensions)
+  def __init__(self, config: config_classes.FakeSimulatorConfig):
+    """FakeSimulator class that can replace EmulatorSimulator in AndroidEnv."""
+    super().__init__(config)
+    self._screen_dimensions = np.array(config.screen_dimensions)
     logging.info('Created FakeSimulator.')
 
   def get_logs(self) -> str:
@@ -124,23 +119,20 @@ class FakeSimulator(base_simulator.BaseSimulator):
     return 'fake_simulator'
 
   def create_adb_controller(self):
-    return FakeAdbController()
+    return FakeAdbController(config_classes.AdbControllerConfig())
 
   def create_log_stream(self) -> log_stream.LogStream:
     return FakeLogStream()
 
-  def _restart_impl(self) -> None:
-    pass
-
   def _launch_impl(self) -> None:
     pass
 
-  def send_touch(self, touches: List[Tuple[int, int, bool, int]]) -> None:
+  def send_touch(self, touches: list[tuple[int, int, bool, int]]) -> None:
     del touches
 
   def send_key(self, keycode: np.int32, event_type: str) -> None:
     del keycode, event_type
 
-  def get_screenshot(self) -> np.ndarray:
+  def _get_screenshot_impl(self) -> np.ndarray:
     return np.random.randint(
         low=0, high=255, size=(*self._screen_dimensions, 3), dtype=np.uint8)