Allow configuring snippet user. (#899)

Co-authored-by: Nitin Prasad <nkprasad@google.com>

2 files changed, 23 insertions, 3 deletions

diff --git a/mobly/controllers/android_device_lib/snippet_client_v2.py b/mobly/controllers/android_device_lib/snippet_client_v2.py
index f7494c2..2e62f95 100644
--- a/mobly/controllers/android_device_lib/snippet_client_v2.py
+++ b/mobly/controllers/android_device_lib/snippet_client_v2.py
@@ -18,7 +18,7 @@ import enum
 import json
 import re
 import socket
-from typing import Dict
+from typing import Dict, Union
 
 from mobly import utils
 from mobly.controllers.android_device_lib import adb
@@ -89,11 +89,12 @@ class Config:
       other purposes may not take effect and you should use snippet RPCs. This
       is because Mobly snippet runner changes the subsequent instrumentation
       process.
+    user_id: The user id under which to launch the snippet process.
   """
 
   am_instrument_options: Dict[str, str] = dataclasses.field(
       default_factory=dict)
-
+  user_id: Union[int, None] = None
 
 class ConnectionHandshakeCommand(enum.Enum):
   """Commands to send to the server when sending the handshake request.
@@ -142,7 +143,7 @@ class SnippetClientV2(client_base.ClientBase):
     self.device_port = None
     self.uid = UNKNOWN_UID
     self._adb = ad.adb
-    self._user_id = None
+    self._user_id = None if config is None else config.user_id
     self._proc = None
     self._client = None  # keep it to prevent close errors on connect failure
     self._conn = None
diff --git a/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py b/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py
index 86d889b..3b42c9b 100644
--- a/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py
+++ b/tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py
@@ -423,6 +423,25 @@ class SnippetClientV2Test(unittest.TestCase):
   @mock.patch('mobly.controllers.android_device_lib.snippet_client_v2.'
               'utils.start_standing_subprocess')
   @mock.patch.object(_MockAdbProxy, 'shell', return_value=b'setsid')
+  def test_start_server_with_config_specific_user_id(self, _,
+                                                     mock_start_subprocess):
+    """Tests that the correct `--user` argument is passed."""
+    self._make_client(config=snippet_client_v2.Config(user_id=42))
+    self._mock_server_process_starting_response(mock_start_subprocess)
+    self.assertEqual(self.client.user_id, 42)
+
+    self.client.start_server()
+    start_cmd_list = [
+        'adb', 'shell',
+        (f'setsid am instrument --user 42 -w -e action start  '
+         f'{MOCK_SERVER_PATH}')
+    ]
+    self.assertListEqual(mock_start_subprocess.call_args_list,
+                         [mock.call(start_cmd_list, shell=False)])
+
+  @mock.patch('mobly.controllers.android_device_lib.snippet_client_v2.'
+              'utils.start_standing_subprocess')
+  @mock.patch.object(_MockAdbProxy, 'shell', return_value=b'setsid')
   def test_start_server_without_user_id(self, mock_adb, mock_start_subprocess):
     """Tests that `--user` is not added to starting command on SDK < 24."""
     self._make_client_with_extra_adb_properties({'ro.build.version.sdk': '21'})