Support platforms other than Linux (Windows and macOS) when sending key strokes in the Android Emulator simulator.

android_env/components/simulators/emulator/emulator_simulator.py

@@ -15,9 +15,10 @@
 
 """A class that manages an Android Emulator."""
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
 import functools
 import os
+import platform
 import time
 from typing import Any
 
@@ -31,6 +32,7 @@
 from android_env.components.simulators.emulator import emulator_launcher
 from android_env.proto import state_pb2
 import grpc
+import immutabledict
 import numpy as np
 import portpicker
 
@@ -43,6 +45,14 @@
 
 _DEFAULT_SNAPSHOT_NAME: str = 'default_snapshot'
 
+_KEYCODE_TYPE_BY_SYSTEM: Mapping[
+    str, emulator_controller_pb2.KeyboardEvent.KeyCodeType
+] = immutabledict.immutabledict({
+    'Linux': emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
+    'Windows': emulator_controller_pb2.KeyboardEvent.KeyCodeType.Win,
+    'Darwin': emulator_controller_pb2.KeyboardEvent.KeyCodeType.Mac,
+})
+
 
 def _is_existing_emulator_provided(
     launcher_config: config_classes.EmulatorLauncherConfig,
@@ -127,11 +137,25 @@ def wrapper(self: 'EmulatorSimulator', *args: Any, **kwargs: Any) -> Any:
 class EmulatorSimulator(base_simulator.BaseSimulator):
   """Controls an Android Emulator."""
 
-  def __init__(self, config: config_classes.EmulatorConfig):
-    """Instantiates an EmulatorSimulator."""
+  def __init__(
+      self,
+      config: config_classes.EmulatorConfig,
+      *,
+      platform_system: str | None = None,
+  ):
+    """Initializes the instance.
+
+    Args:
+      config: Configuration for the EmulatorSimulator.
+      platform_system: The name of the host operating system, as returned by
+        `platform.system()` (e.g., 'Linux', 'Windows', 'Darwin'). This is used
+        to determine the correct keycode type for `send_key`. If None,
+        `platform.system()` is called to get the current system.
+    """
 
     super().__init__(config)
     self._config: config_classes.EmulatorConfig = config
+    self._platform_system: str = platform_system or platform.system()
 
     # If adb_port, console_port and grpc_port are all already provided,
     # we assume the emulator already exists and there's no need to launch.
@@ -439,8 +463,10 @@ def send_key(self, keycode: np.int32, event_type: str) -> None:
     """Sends a key event to the emulator.
 
     Args:
-      keycode: Code representing the desired key press in XKB format.
-        See the emulator_controller_pb2 for details.
+      keycode: Code representing the desired key press. Format depends on the
+        host platform (using the `KeyCodeType` enum in
+        `emulator_controller_pb2`: `XKB` for Linux, `Win` for Windows, `Mac` for
+        macOS). See the emulator_controller_pb2 for details.
       event_type: Type of key event to be sent.
     """
 
@@ -452,9 +478,19 @@ def send_key(self, keycode: np.int32, event_type: str) -> None:
     assert (
         self._emulator_stub is not None
     ), 'Emulator stub has not been initialized yet.'
+
+    system_code_type = _KEYCODE_TYPE_BY_SYSTEM.get(self._platform_system)
+    if system_code_type is not None:
+      code_type = system_code_type
+    else:
+      logging.warning(
+          'Unknown system %r, falling back to XKB', self._platform_system
+      )
+      code_type = emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB
+
     self._emulator_stub.sendKey(
         emulator_controller_pb2.KeyboardEvent(
-            codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
+            codeType=code_type,
             eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.Value(
                 event_type
             ),

android_env/components/simulators/emulator/emulator_simulator_test.py

@@ -21,6 +21,7 @@
 from unittest import mock
 
 from absl.testing import absltest
+from absl.testing import parameterized
 from android_env.components import adb_call_parser
 from android_env.components import adb_controller
 from android_env.components import config_classes
@@ -36,28 +37,46 @@
 from android_env.proto import snapshot_service_pb2
 
 
-class EmulatorSimulatorTest(absltest.TestCase):
+class EmulatorSimulatorTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
     self.addCleanup(mock.patch.stopall)  # Disable previous patches.
 
-    self._adb_controller = mock.create_autospec(adb_controller.AdbController)
-    self._adb_call_parser = mock.create_autospec(adb_call_parser.AdbCallParser)
-    self._launcher = mock.create_autospec(emulator_launcher.EmulatorLauncher)
+    self._adb_controller = mock.create_autospec(
+        adb_controller.AdbController, spec_set=True, instance=True
+    )
+    self._adb_call_parser = mock.create_autospec(
+        adb_call_parser.AdbCallParser, spec_set=True, instance=True
+    )
+    self._launcher = mock.create_autospec(
+        emulator_launcher.EmulatorLauncher, spec_set=True, instance=True
+    )
     self._launcher.logfile_path.return_value = 'logfile_path'
     self._emulator_stub = mock.create_autospec(
-        emulator_controller_pb2_grpc.EmulatorControllerStub)
+        emulator_controller_pb2_grpc.EmulatorControllerStub,
+        spec_set=True,
+        instance=True,
+    )
 
-    self._grpc_channel = mock.create_autospec(grpc.Channel)
+    self._grpc_channel = mock.create_autospec(
+        grpc.Channel, spec_set=True, instance=True
+    )
+    self._grpc_channel.unary_unary.side_effect = (
+        lambda *args, **kwargs: mock.create_autospec(
+            grpc.UnaryUnaryMultiCallable, instance=True, spec_set=True
+        )
+    )
     mock.patch.object(
         grpc.aio, 'secure_channel', return_value=self._grpc_channel).start()
     mock.patch.object(
         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)
+    self._mock_future = mock.create_autospec(
+        grpc.Future, spec_set=True, instance=True
+    )
     mock.patch.object(
         grpc, 'channel_ready_future', return_value=self._mock_future).start()
     mock.patch.object(time, 'time', return_value=12345).start()
@@ -521,7 +540,29 @@ def test_send_touch(self):
             }])),
     ])
 
-  def test_send_key(self):
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='Linux',
+          os_name='Linux',
+          expected_code_type=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
+      ),
+      dict(
+          testcase_name='Windows',
+          os_name='Windows',
+          expected_code_type=emulator_controller_pb2.KeyboardEvent.KeyCodeType.Win,
+      ),
+      dict(
+          testcase_name='Darwin',
+          os_name='Darwin',
+          expected_code_type=emulator_controller_pb2.KeyboardEvent.KeyCodeType.Mac,
+      ),
+      dict(
+          testcase_name='UnknownOS',
+          os_name='UnknownOS',
+          expected_code_type=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
+      ),
+  )
+  def test_send_key(self, os_name, expected_code_type):
     config = config_classes.EmulatorConfig(
         emulator_launcher=config_classes.EmulatorLauncherConfig(
             grpc_port=1234, tmp_dir=self.create_tempdir().full_path
@@ -531,12 +572,13 @@ def test_send_key(self):
             adb_server_port=5037,
         ),
     )
-    simulator = emulator_simulator.EmulatorSimulator(config)
+    simulator = emulator_simulator.EmulatorSimulator(
+        config, platform_system=os_name
+    )
 
     # The simulator should launch and not crash.
     simulator.launch()
     self.assertIsNotNone(simulator._emulator_stub)
-    simulator._emulator_stub.sendTouch = mock.MagicMock(return_value=None)
 
     simulator.send_key(123, 'keydown')
     simulator.send_key(321, 'keydown')
@@ -545,50 +587,53 @@ def test_send_key(self):
     simulator.send_key(321, 'keypress')
     simulator.send_key(123, 'keypress')
 
-    simulator._emulator_stub.sendKey.assert_has_calls([
+    expected_calls = [
         mock.call(
             emulator_controller_pb2.KeyboardEvent(
-                codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
-                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType
-                .keydown,
+                codeType=expected_code_type,
+                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.keydown,
                 keyCode=123,
-            )),
+            )
+        ),
         mock.call(
             emulator_controller_pb2.KeyboardEvent(
-                codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
-                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType
-                .keydown,
+                codeType=expected_code_type,
+                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.keydown,
                 keyCode=321,
-            )),
+            )
+        ),
         mock.call(
             emulator_controller_pb2.KeyboardEvent(
-                codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
-                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType
-                .keyup,
+                codeType=expected_code_type,
+                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.keyup,
                 keyCode=321,
-            )),
+            )
+        ),
         mock.call(
             emulator_controller_pb2.KeyboardEvent(
-                codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
-                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType
-                .keyup,
+                codeType=expected_code_type,
+                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.keyup,
                 keyCode=123,
-            )),
+            )
+        ),
         mock.call(
             emulator_controller_pb2.KeyboardEvent(
-                codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
-                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType
-                .keypress,
+                codeType=expected_code_type,
+                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.keypress,
                 keyCode=321,
-            )),
+            )
+        ),
         mock.call(
             emulator_controller_pb2.KeyboardEvent(
-                codeType=emulator_controller_pb2.KeyboardEvent.KeyCodeType.XKB,
-                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType
-                .keypress,
+                codeType=expected_code_type,
+                eventType=emulator_controller_pb2.KeyboardEvent.KeyEventType.keypress,
                 keyCode=123,
-            ))
-    ])
+            )
+        ),
+    ]
+    self.assertSequenceEqual(
+        expected_calls, simulator._emulator_stub.sendKey.call_args_list
+    )
 
 
 if __name__ == '__main__':