Internal change

PiperOrigin-RevId: 686531041

android_world/env/android_world_controller.py

@@ -146,11 +146,12 @@ def __init__(
       self,
       env: env_interface.AndroidEnvInterface,
       a11y_method: A11yMethod = A11yMethod.A11Y_FORWARDER_APP,
+      install_a11y_forwarding_app: bool = True,
   ):
     if a11y_method == A11yMethod.A11Y_FORWARDER_APP:
       self._env = a11y_grpc_wrapper.A11yGrpcWrapper(
           env,
-          install_a11y_forwarding=True,
+          install_a11y_forwarding=install_a11y_forwarding_app,
           start_a11y_service=True,
           enable_a11y_tree_info=True,
           latest_a11y_info_only=True,

android_world/env/setup_device/setup.py

@@ -24,6 +24,7 @@
 import os
 
 from absl import logging
+from android_env import env_interface
 from android_env.components import errors
 from android_world.env import adb_utils
 from android_world.env import interface
@@ -61,10 +62,12 @@
 )
 
 
-def _download_and_install_apk(apk: str, env: interface.AsyncEnv) -> None:
-  """Downloads all APKs from remote location and installs them."""
+def download_and_install_apk(
+    apk: str, raw_env: env_interface.AndroidEnvInterface
+) -> None:
+  """Downloads APK from remote location and installs it."""
   path = apps.download_app_data(apk)
-  adb_utils.install_apk(path, env.controller)
+  adb_utils.install_apk(path, controller)
 
 
 def _install_all_apks(env: interface.AsyncEnv) -> None:
@@ -76,7 +79,7 @@ def _install_all_apks(env: interface.AsyncEnv) -> None:
     apk_installed = False
     for apk_name in app.apk_names:
       try:
-        _download_and_install_apk(apk_name, env)
+        download_and_install_apk(apk_name, env.controller.env)
         apk_installed = True
         break
       except errors.AdbControllerError:

android_world/env/setup_device/setup_test.py

@@ -33,7 +33,7 @@ def setUp(self):
     )
 
   @mock.patch.object(tools, "AndroidToolController")
-  @mock.patch.object(setup, "_download_and_install_apk")
+  @mock.patch.object(setup, "download_and_install_apk")
   @mock.patch.object(app_snapshot, "save_snapshot")
   def test_setup_apps(self, mock_save_snapshot, mock_install_apk, unused_tools):
     env = mock.create_autospec(interface.AsyncEnv)
@@ -46,7 +46,9 @@ def test_setup_apps(self, mock_save_snapshot, mock_install_apk, unused_tools):
 
     for app_class in setup._APPS:
       if app_class.apk_names:  # 1P apps do not have APKs.
-        mock_install_apk.assert_any_call(app_class.apk_names[0], env)
+        mock_install_apk.assert_any_call(
+            app_class.apk_names[0], env.controller.env
+        )
       mock_app_setups[app_class].assert_any_call(env)
       mock_save_snapshot.assert_any_call(app_class.app_name, env.controller)
 
@@ -63,8 +65,8 @@ class InstallApksTest(absltest.TestCase):
   def setUp(self):
     super().setUp()
     self.env = mock.create_autospec(interface.AsyncEnv)
-    self.mock_download_and_install_apk = self.enter_context(
-        mock.patch.object(setup, "_download_and_install_apk")
+    self.mockdownload_and_install_apk = self.enter_context(
+        mock.patch.object(setup, "download_and_install_apk")
     )
     self.apps = [
         _App(apk_names=["apk1", "apk2"], app_name="App1"),
@@ -74,12 +76,15 @@ def setUp(self):
     setup._APPS = self.apps
 
   def test_install_all_apks_success(self):
-    self.mock_download_and_install_apk.return_value = None
+    self.mockdownload_and_install_apk.return_value = None
 
     setup._install_all_apks(self.env)
 
-    expected_calls = [mock.call("apk1", self.env), mock.call("apk3", self.env)]
-    self.mock_download_and_install_apk.assert_has_calls(
+    expected_calls = [
+        mock.call("apk1", self.env.controller.env),
+        mock.call("apk3", self.env.controller.env),
+    ]
+    self.mockdownload_and_install_apk.assert_has_calls(
         expected_calls, any_order=True
     )
 
@@ -90,16 +95,16 @@ def side_effect(apk_name, env):
         raise errors.AdbControllerError
       return None
 
-    self.mock_download_and_install_apk.side_effect = side_effect
+    self.mockdownload_and_install_apk.side_effect = side_effect
 
     setup._install_all_apks(self.env)
 
     expected_calls = [
-        mock.call("apk1", self.env),
-        mock.call("apk2", self.env),
-        mock.call("apk3", self.env),
+        mock.call("apk1", self.env.controller.env),
+        mock.call("apk2", self.env.controller.env),
+        mock.call("apk3", self.env.controller.env),
     ]
-    self.mock_download_and_install_apk.assert_has_calls(expected_calls)
+    self.mockdownload_and_install_apk.assert_has_calls(expected_calls)
 
 
 if __name__ == "__main__":