From 0c6eab5a44c7a32ba98748f35a7126b24373a0b3 Mon Sep 17 00:00:00 2001
From: Daniel Toyama <kenjitoyama@google.com>
Date: Tue, 19 May 2026 14:27:53 -0700
Subject: [PATCH] Fix type safety issues in AndroidEnv components and tests. -
 Refactor action_fns_test to wrap enums in np.array to match strict library
 type signatures.

PiperOrigin-RevId: 918030200
---
 android_env/components/action_fns_test.py | 72 +++++------------------
 android_env/components/device_settings.py |  1 -
 2 files changed, 16 insertions(+), 57 deletions(-)

diff --git a/android_env/components/action_fns_test.py b/android_env/components/action_fns_test.py
index c9cf91e9..20b46b0d 100644
--- a/android_env/components/action_fns_test.py
+++ b/android_env/components/action_fns_test.py
@@ -51,7 +51,7 @@ def test_send_action_to_simulator_sendactionerror(self):
     simulator = mock.create_autospec(base_simulator.BaseSimulator)
     simulator.send_touch.side_effect = errors.SendActionError('oops!')
     action = {
-        'action_type': action_type_lib.ActionType.TOUCH,
+        'action_type': np.array(action_type_lib.ActionType.TOUCH),
         'touch_position': np.array([0.3, 0.5], np.float32),
     }
 
@@ -74,7 +74,7 @@ def test_send_action_to_simulator_touch_success_one_finger(self):
     # Arrange.
     simulator = mock.create_autospec(base_simulator.BaseSimulator)
     action = {
-        'action_type': action_type_lib.ActionType.TOUCH,
+        'action_type': np.array(action_type_lib.ActionType.TOUCH),
         'touch_position': np.array([0.2, 0.5], np.float32),
     }
 
@@ -99,11 +99,11 @@ def test_send_action_to_simulator_touch_success_multiple_finger(self):
     # Arrange.
     simulator = mock.create_autospec(base_simulator.BaseSimulator)
     action = {
-        'action_type': action_type_lib.ActionType.TOUCH,
+        'action_type': np.array(action_type_lib.ActionType.TOUCH),
         'touch_position': np.array([0.2, 0.5], np.float32),
-        'action_type_2': action_type_lib.ActionType.LIFT,
+        'action_type_2': np.array(action_type_lib.ActionType.LIFT),
         'touch_position_2': np.array([0.1, 0.2], np.float32),
-        'action_type_3': action_type_lib.ActionType.TOUCH,
+        'action_type_3': np.array(action_type_lib.ActionType.TOUCH),
         'touch_position_3': np.array([0.5, 0.2], np.float32),
     }
 
@@ -124,60 +124,20 @@ def test_send_action_to_simulator_touch_success_multiple_finger(self):
         (np.int32(400), np.int32(120), True, 2),
     ])
 
-  def test_send_action_to_simulator_keydown_success(self):
-    """Returns `True` with a proper keydown action."""
-
-    # Arrange.
-    simulator = mock.create_autospec(base_simulator.BaseSimulator)
-    action = {
-        'action_type': action_type_lib.ActionType.KEYDOWN,
-        'keycode': np.array([21], np.int32),
-    }
-
-    # Act.
-    output = action_fns.send_action_to_simulator(
-        action,
-        simulator,
-        800,
-        600,
-        1,
-    )
-
-    # Assert.
-    self.assertTrue(output)
-    simulator.send_key.assert_called_once_with(21, event_type='keydown')
-
-  def test_send_action_to_simulator_keyup_success(self):
-    """Returns `True` with a proper keyup action."""
-
-    # Arrange.
-    simulator = mock.create_autospec(base_simulator.BaseSimulator)
-    action = {
-        'action_type': action_type_lib.ActionType.KEYUP,
-        'keycode': np.array([42], np.int32),
-    }
-
-    # Act.
-    output = action_fns.send_action_to_simulator(
-        action,
-        simulator,
-        800,
-        600,
-        1,
-    )
-
-    # Assert.
-    self.assertTrue(output)
-    simulator.send_key.assert_called_once_with(42, event_type='keyup')
-
-  def test_send_action_to_simulator_keypress_success(self):
-    """Returns `True` with a proper keypress action."""
+  @parameterized.named_parameters(
+      ('keydown', action_type_lib.ActionType.KEYDOWN, 21, 'keydown'),
+      ('keyup', action_type_lib.ActionType.KEYUP, 42, 'keyup'),
+      ('keypress', action_type_lib.ActionType.KEYPRESS, 96, 'keypress'),
+  )
+  def test_send_action_to_simulator_key_event_success(
+      self, action_type, keycode, event_type
+  ):
 
     # Arrange.
     simulator = mock.create_autospec(base_simulator.BaseSimulator)
     action = {
-        'action_type': action_type_lib.ActionType.KEYPRESS,
-        'keycode': np.array([96], np.int32),
+        'action_type': np.array(action_type),
+        'keycode': np.array([keycode], np.int32),
     }
 
     # Act.
@@ -191,7 +151,7 @@ def test_send_action_to_simulator_keypress_success(self):
 
     # Assert.
     self.assertTrue(output)
-    simulator.send_key.assert_called_once_with(96, event_type='keypress')
+    simulator.send_key.assert_called_once_with(keycode, event_type=event_type)
 
   @parameterized.named_parameters(
       (
diff --git a/android_env/components/device_settings.py b/android_env/components/device_settings.py
index 13335ed4..acd3341b 100644
--- a/android_env/components/device_settings.py
+++ b/android_env/components/device_settings.py
@@ -25,7 +25,6 @@
 from android_env.proto import adb_pb2
 import numpy as np
 
-
 # The internal `AdbCallParser` instance is lazily instantiated within
 # `DeviceSettings`. If we make it optional (i.e. `| None`), pytype will think
 # that it could be `None`, requiring either explicit runtime checks or escape