Fix that UIAutomator may not return output dump on a real device if the device screen is off.

android_env/components/adb_call_parser.py

@@ -60,6 +60,7 @@ def __init__(self, adb_controller: adb_control.AdbController):
         'generic': self._handle_generic,
         'package_manager': self._handle_package_manager,
         'dumpsys': self._handle_dumpsys,
+        'uiautomator': self._handle_uiautomator,
     }
 
   def _execute_command(
@@ -905,3 +906,32 @@ def _handle_dumpsys(
         cmd, timeout=timeout)
 
     return response
+
+  def _handle_uiautomator(
+      self, request: adb_pb2.AdbRequest, timeout: float | None = None
+  ) -> adb_pb2.AdbResponse:
+    """Handles UIAUtomatorRequest messages.
+
+    Args:
+      request: The request with the `.uiautomator` field set containing
+        sub-commands to `adb shell uiautomator dump` shell command..
+      timeout: Optional time limit in seconds.
+
+    Returns:
+      An AdbResponse.
+    """
+    request = request.uiautomator
+    cmd = ['shell', 'uiautomator', 'dump']
+
+    if request.file:
+      cmd.append(request.file)
+
+    response, cmd_output = self._execute_command(cmd, timeout=timeout)
+
+    if cmd_output.startswith('UI hierchary dumped to'.encode('utf-8')):
+      response.status = adb_pb2.AdbResponse.Status.OK
+      response.uiautomator.output = cmd_output
+    else:
+      response.status = adb_pb2.AdbResponse.Status.ADB_ERROR
+      response.error_message = cmd_output
+    return response

android_env/loader.py

@@ -25,6 +25,7 @@
 from android_env.components import task_manager as task_manager_lib
 from android_env.components.simulators.emulator import emulator_simulator
 from android_env.components.simulators.fake import fake_simulator
+from android_env.components.simulators.real_device import real_device_simulator
 from android_env.proto import task_pb2
 
 from google.protobuf import text_format
@@ -56,6 +57,10 @@ def load(config: config_classes.AndroidEnvConfig) -> environment.AndroidEnv:
       simulator = emulator_simulator.EmulatorSimulator(config=config.simulator)
     case config_classes.FakeSimulatorConfig():
       simulator = fake_simulator.FakeSimulator(config=config.simulator)
+    case config_classes.RealDeviceConfig():
+      simulator = real_device_simulator.RealDeviceSimulator(
+          config=config.simulator
+      )
     case _:
       raise ValueError('Unsupported simulator config: {config.simulator}')
 

android_env/proto/adb.proto

@@ -286,6 +286,13 @@ message AdbRequest {
     repeated string skip_services = 9;
   }
 
+  // For executing `uiautomator dump`.
+  message UIAutomatorRequest {
+    // The file to dump the accessibility tree to. If empty, the accessibility
+    // tree will be dumped to the device's /sdcard/window_dump.xml.
+    string file = 1;
+  }
+
   oneof command {
     InstallApk install_apk = 1;
     StartActivity start_activity = 2;
@@ -304,6 +311,7 @@ message AdbRequest {
     PackageManagerRequest package_manager = 23;
     DumpsysRequest dumpsys = 26;
     SendBroadcast send_broadcast = 25;
+    UIAutomatorRequest uiautomator = 27;
   }
 
   // Optional (soft) deadline in seconds for completing this command.
@@ -417,6 +425,12 @@ message AdbResponse {
     bytes output = 1;
   }
 
+  // Response for UIAutomatorRequests.
+  message UIAutomatorResponse {
+    // The output, if any, of the `uiautomator dump` command.
+    string output = 1;
+  }
+
   oneof payload {
     GetCurrentActivityResponse get_current_activity = 10;
     StartActivityResponse start_activity = 11;
@@ -429,5 +443,6 @@ message AdbResponse {
     PackageManagerResponse package_manager = 18;
     GetOrientationResponse get_orientation = 19;
     DumpsysResponse dumpsys = 21;
+    UIAutomatorResponse uiautomator = 22;
   }
 }