feat: add env_str to snippet_client_v2.Config to allow am instrument command prefix (fixes #1004)

mobly/controllers/android_device_lib/snippet_client_v2.py

@@ -34,7 +34,7 @@
 
 # The command template to start the snippet server
 _LAUNCH_CMD = (
-    '{shell_cmd} am instrument {user} -w -e action start'
+    '{shell_cmd} {env_str}am instrument {user} -w -e action start'
     ' {instrument_options}'
     f' {{snippet_package}}/{_INSTRUMENTATION_RUNNER_PACKAGE}'
 )
@@ -112,12 +112,16 @@ class Config:
       is because Mobly snippet runner changes the subsequent instrumentation
       process.
     user_id: The user id under which to launch the snippet process.
+    env_str: An optional string of environment variables or command wrappers
+      to prepend to the `am instrument` command. E.g. `'ENV_VAR=value'` or
+      `'wrapper_cmd'`. Defaults to an empty string.
   """
 
   am_instrument_options: Dict[str, str] = dataclasses.field(
       default_factory=dict
   )
   user_id: Union[int, None] = None
+  env_str: str = ''
 
 
 class ConnectionHandshakeCommand(enum.Enum):
@@ -310,6 +314,7 @@ def start_server(self):
     option_str = self._get_instrument_options_str()
     cmd = _LAUNCH_CMD.format(
         shell_cmd=persists_shell_cmd,
+        env_str=self._config.env_str + ' ' if self._config.env_str else '',
         user=self._get_user_command_string(),
         snippet_package=self.package,
         instrument_options=option_str,

tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py

@@ -1901,3 +1901,97 @@ def test_make_conn_with_forwarded_port_continue(
 
 if __name__ == '__main__':
   unittest.main()
+
+  # ---------------------------------------------------------------------------
+  # Tests for Config.env_str (PR #1008)
+  # ---------------------------------------------------------------------------
+
+  def test_config_env_str_default_is_empty_string(self):
+    """Checks that Config.env_str defaults to an empty string."""
+    config = snippet_client_v2.Config()
+    self.assertEqual(config.env_str, '')
+
+  def test_config_env_str_can_be_set(self):
+    """Checks that Config.env_str can be assigned an arbitrary string."""
+    config = snippet_client_v2.Config(env_str='MY_VAR=value')
+    self.assertEqual(config.env_str, 'MY_VAR=value')
+
+  @mock.patch(
+      'mobly.controllers.android_device_lib.snippet_client_v2.'
+      'utils.start_standing_subprocess'
+  )
+  def test_start_server_with_env_str(self, mock_start_subprocess):
+    """Checks that a non-empty env_str is prepended before am instrument."""
+    self._make_client()
+    self.client._config = snippet_client_v2.Config(env_str='MY_VAR=my_value')
+    self._mock_server_process_starting_response(mock_start_subprocess)
+
+    self.client.start_server()
+
+    start_cmd_list = [
+        'adb',
+        'shell',
+        (
+            f'MY_VAR=my_value am instrument --user {MOCK_USER_ID} -w -e action start  '
+            f'{MOCK_SERVER_PATH}'
+        ),
+    ]
+    self.assertListEqual(
+        mock_start_subprocess.call_args_list,
+        [mock.call(start_cmd_list, shell=False)],
+    )
+    self.assertEqual(self.client.device_port, MOCK_DEVICE_PORT)
+
+  @mock.patch(
+      'mobly.controllers.android_device_lib.snippet_client_v2.'
+      'utils.start_standing_subprocess'
+  )
+  def test_start_server_without_env_str_no_extra_space(
+      self, mock_start_subprocess
+  ):
+    """Checks that an empty env_str adds no extra space before am instrument."""
+    self._make_client()
+    self._mock_server_process_starting_response(mock_start_subprocess)
+
+    self.client.start_server()
+
+    # When env_str is empty the shell command must not have a leading space
+    start_cmd = mock_start_subprocess.call_args_list[0][0][0][2]
+    self.assertNotIn('  am instrument', start_cmd)
+    self.assertIn('am instrument', start_cmd)
+    self.assertEqual(self.client.device_port, MOCK_DEVICE_PORT)
+
+  @mock.patch(
+      'mobly.controllers.android_device_lib.snippet_client_v2.'
+      'utils.start_standing_subprocess'
+  )
+  def test_start_server_with_env_str_has_trailing_space(
+      self, mock_start_subprocess
+  ):
+    """Checks env_str is separated from am instrument by exactly one space."""
+    self._make_client()
+    self.client._config = snippet_client_v2.Config(env_str='WRAPPER_CMD')
+    self._mock_server_process_starting_response(mock_start_subprocess)
+
+    self.client.start_server()
+
+    start_cmd = mock_start_subprocess.call_args_list[0][0][0][2]
+    self.assertIn('WRAPPER_CMD am instrument', start_cmd)
+    self.assertNotIn('WRAPPER_CMD  am instrument', start_cmd)
+    self.assertEqual(self.client.device_port, MOCK_DEVICE_PORT)
+
+  @mock.patch(
+      'mobly.controllers.android_device_lib.snippet_client_v2.'
+      'utils.start_standing_subprocess'
+  )
+  def test_start_server_with_multiple_env_vars(self, mock_start_subprocess):
+    """Checks that multiple space-separated env vars in env_str are kept."""
+    self._make_client()
+    self.client._config = snippet_client_v2.Config(env_str='VAR_A=1 VAR_B=2')
+    self._mock_server_process_starting_response(mock_start_subprocess)
+
+    self.client.start_server()
+
+    start_cmd = mock_start_subprocess.call_args_list[0][0][0][2]
+    self.assertIn('VAR_A=1 VAR_B=2 am instrument', start_cmd)
+    self.assertEqual(self.client.device_port, MOCK_DEVICE_PORT)