Don't restart adb server if last attempt to execute command fails

PiperOrigin-RevId: 582783136
google-deepmind · Nov 15, 2023 · 5e9ac3b · 5e9ac3b
1 parent 93062aa
commit 5e9ac3b
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 2 deletions.
diff --git a/android_env/components/adb_controller.py b/android_env/components/adb_controller.py
@@ -122,9 +122,10 @@ def execute_command(
     command = self.command_prefix(include_device_name=device_specific) + args
     command_str = 'adb ' + ' '.join(command[1:])
 
+    n_retries = 2
     n_tries = 1
     latest_error = None
-    while n_tries < 3:
+    while n_tries <= n_retries:
       try:
         logging.info('Executing ADB command: [%s]', command_str)
         cmd_output = subprocess.check_output(
@@ -149,7 +150,7 @@ def execute_command(
             logging.error('    %s', line)
         n_tries += 1
         latest_error = e
-        if device_specific:
+        if device_specific and n_tries <= n_retries:
           self._restart_server(timeout=timeout)
 
     raise errors.AdbControllerError(

diff --git a/android_env/components/adb_controller_test.py b/android_env/components/adb_controller_test.py
@@ -112,6 +112,69 @@ def test_restart_server(self, mock_sleep, mock_check_output):
     mock_sleep.assert_has_calls(
         [mock.call(0.2), mock.call(2.0), mock.call(0.2)])
 
+  @mock.patch.object(subprocess, 'check_output', autospec=True)
+  @mock.patch.object(time, 'sleep', autospec=True)
+  def test_invalid_command(self, mock_sleep, mock_check_output):
+    # Arrange.
+    restart_sequence = ['fake_output'.encode('utf-8')] * 3
+    mock_check_output.side_effect = (
+        [
+            subprocess.CalledProcessError(returncode=1, cmd='blah'),
+        ]
+        + restart_sequence
+        + [subprocess.CalledProcessError(returncode=1, cmd='blah')]
+        # Don't restart if last call fails.
+    )
+    adb_controller = adb_controller_lib.AdbController(
+        adb_path='my_adb', device_name='awesome_device', adb_server_port=9999
+    )
+
+    # Act.
+    with self.assertRaises(errors.AdbControllerError):
+      adb_controller.execute_command(['my_command'], timeout=_TIMEOUT)
+
+    # Assert.
+    expected_env = self._env_before
+    expected_env['HOME'] = '/some/path/'
+    mock_check_output.assert_has_calls(
+        [
+            mock.call(
+                ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'kill-server'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'start-server'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', 'devices'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+            mock.call(
+                ['my_adb', '-P', '9999', '-s', 'awesome_device', 'my_command'],
+                stderr=subprocess.STDOUT,
+                timeout=_TIMEOUT,
+                env=expected_env,
+            ),
+        ],
+        any_order=False,
+    )
+    mock_sleep.assert_has_calls(
+        [mock.call(0.2), mock.call(2.0), mock.call(0.2)]
+    )
+
   @mock.patch.object(subprocess, 'check_output', autospec=True)
   @mock.patch.object(time, 'sleep', autospec=True)
   def test_avoid_infinite_recursion(self, mock_sleep, mock_check_output):