Allow loading snippets with the same package but different user IDs

mobly/controllers/android_device_lib/services/snippet_management_service.py

@@ -69,31 +69,36 @@ def add_snippet_client(self, name, package, config=None):
         class for supported configurations.
 
     Raises:
-      Error, if a duplicated name or package is passed in.
+      Error: if a duplicated name is passed in, or the same package has
+        already been registered under the same Android user ID.
     """
     # Should not load snippet with the same name more than once.
     if name in self._snippet_clients:
       raise Error(
           self,
-          'Name "%s" is already registered with package "%s", it cannot '
-          'be used again.' % (name, self._snippet_clients[name].client.package),
+          f'Name "{name}" is already registered with package'
+          f' "{self._snippet_clients[name].package}" for user ID'
+          f' {self._snippet_clients[name].user_id}, the same name'
+          ' cannot be used again.',
       )
-    # Should not load the same snippet package more than once.
+    # Should not load snippets with the same identifier more than once.
+    new_client = snippet_client_v2.SnippetClientV2(
+        package=package,
+        ad=self._device,
+        config=config,
+    )
     for snippet_name, client in self._snippet_clients.items():
-      if package == client.package:
+      if new_client.identifier == client.identifier:
+        del new_client
         raise Error(
             self,
-            'Snippet package "%s" has already been loaded under name "%s".'
-            % (package, snippet_name),
+            f'Snippet "{client.package}" has already been registered for user'
+            f' id {client.user_id} under name "{snippet_name}". The same'
+            ' package cannot be registered again for the same user.',
         )
 
-    client = snippet_client_v2.SnippetClientV2(
-        package=package,
-        ad=self._device,
-        config=config,
-    )
-    client.initialize()
-    self._snippet_clients[name] = client
+    new_client.initialize()
+    self._snippet_clients[name] = new_client
 
   def remove_snippet_client(self, name):
     """Removes a snippet client from management.
@@ -113,24 +118,24 @@ def start(self):
     """Starts all the snippet clients under management."""
     for client in self._snippet_clients.values():
       if not client.is_alive:
-        self._device.log.debug('Starting SnippetClient<%s>.', client.package)
+        self._device.log.debug('Starting SnippetClient<%s>.', str(client))
         client.initialize()
       else:
         self._device.log.debug(
             'Not startng SnippetClient<%s> because it is already alive.',
-            client.package,
+            str(client),
         )
 
   def stop(self):
     """Stops all the snippet clients under management."""
     for client in self._snippet_clients.values():
       if client.is_alive:
-        self._device.log.debug('Stopping SnippetClient<%s>.', client.package)
+        self._device.log.debug('Stopping SnippetClient<%s>.', str(client))
         client.stop()
       else:
         self._device.log.debug(
             'Not stopping SnippetClient<%s> because it is not alive.',
-            client.package,
+            str(client),
         )
 
   def pause(self):
@@ -140,19 +145,17 @@ def pause(self):
     allocated in `resume`.
     """
     for client in self._snippet_clients.values():
-      self._device.log.debug('Pausing SnippetClient<%s>.', client.package)
+      self._device.log.debug('Pausing SnippetClient<%s>.', str(client))
       client.close_connection()
 
   def resume(self):
     """Resumes all paused snippet clients."""
     for client in self._snippet_clients.values():
       if not client.is_alive:
-        self._device.log.debug('Resuming SnippetClient<%s>.', client.package)
+        self._device.log.debug('Resuming SnippetClient<%s>.', str(client))
         client.restore_server_connection()
       else:
-        self._device.log.debug(
-            'Not resuming SnippetClient<%s>.', client.package
-        )
+        self._device.log.debug('Not resuming SnippetClient<%s>.', str(client))
 
   def __getattr__(self, name):
     client = self.get_snippet_client(name)

mobly/controllers/android_device_lib/snippet_client_v2.py

@@ -198,11 +198,27 @@ def user_id(self):
       self._user_id = self._adb.current_user_id
     return self._user_id
 
+  @property
+  def identifier(self):
+    """The unique identifier of this snippet client.
+
+    This property serves as the singular key for the snippet client, ensuring
+    that every loaded snippet client possesses a distinct identifier.
+
+    This identifier is constructed by combining the client's package name
+    with the user ID. The user ID is needed since it's allowed to load snippets
+    with the same package for different Android users.
+    """
+    return f'{self.package}@user_id[{self.user_id}]'
+
   @property
   def is_alive(self):
     """Does the client have an active connection to the snippet server."""
     return self._conn is not None
 
+  def __repr__(self):
+    return self.identifier
+
   def before_starting_server(self):
     """Performs the preparation steps before starting the remote server.
 
@@ -286,8 +302,8 @@ def start_server(self):
     """
     persists_shell_cmd = self._get_persisting_command()
     self.log.debug(
-        'Snippet server for package %s is using protocol %d.%d',
-        self.package,
+        'Snippet server for %s is using protocol %d.%d',
+        str(self),
         _PROTOCOL_MAJOR_VERSION,
         _PROTOCOL_MINOR_VERSION,
     )
@@ -345,8 +361,8 @@ def _get_persisting_command(self):
 
   def _get_instrument_options_str(self):
     self.log.debug(
-        'Got am instrument options in snippet client for package %s: %s',
-        self.package,
+        'Got am instrument options in snippet client "%s": %s',
+        str(self),
         self._config.am_instrument_options,
     )
     if not self._config.am_instrument_options:
@@ -457,7 +473,7 @@ def create_socket_connection(self):
       self.log.debug(
           'Snippet client is creating socket connection to the snippet server '
           'of %s through host port %d.',
-          self.package,
+          str(self),
           self.host_port,
       )
       self._conn = socket.create_connection(
@@ -699,11 +715,11 @@ def stop(self):
       android_device_lib_errors.DeviceError: if the server exited with errors on
         the device side.
     """
-    self.log.debug('Stopping snippet package %s.', self.package)
+    self.log.debug('Stopping snippet client %s.', str(self))
     self.close_connection()
     self._stop_server()
     self._destroy_event_client()
-    self.log.debug('Snippet package %s stopped.', self.package)
+    self.log.debug('Snippet client %s stopped.', str(self))
 
   def close_connection(self):
     """Closes the connection to the snippet server on the device.
@@ -810,8 +826,9 @@ def restore_server_connection(self, port=None):
       raise errors.ServerRestoreConnectionError(
           self._device,
           (
-              f'Failed to restore server connection for {self.package} at '
-              f'host port {self.host_port}, device port {self.device_port}.'
+              'Failed to restore server connection of the snippet package'
+              f' {self.package} for user id {self.user_id} at host port'
+              f' {self.host_port}, device port {self.device_port}.'
           ),
       ) from e
 

tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py

@@ -12,18 +12,62 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import unittest
 from unittest import mock
 
+from mobly.snippet import client_base
 from mobly.controllers.android_device_lib import snippet_client_v2
 from mobly.controllers.android_device_lib.services import snippet_management_service
 
 MOCK_PACKAGE = 'com.mock.package'
+MOCK_PACKAGE2 = 'com.mock.package2'
 SNIPPET_CLIENT_V2_CLASS_PATH = (
     'mobly.controllers.android_device_lib.snippet_client_v2.SnippetClientV2'
 )
 
 
+class MockSnippetClientV2(client_base.ClientBase):
+
+  def __init__(self, package, ad, config=None):
+    self.user_id = (
+        config.user_id
+        if config is not None and config.user_id is not None
+        else ad.adb.current_user_id
+    )
+    self.package = package
+    self.identifier = f'{self.package}@user_id[{self.user_id}]'
+    self.log = logging
+
+  # Override abstract methods so this class can be instantiated.
+  def before_starting_server(self):
+    pass
+
+  def start_server(self):
+    pass
+
+  def make_connection(self):
+    pass
+
+  def restore_server_connection(self, port=None):
+    pass
+
+  def check_server_proc_running(self):
+    pass
+
+  def send_rpc_request(self, request):
+    pass
+
+  def handle_callback(self, callback_id, ret_value, rpc_func_name):
+    pass
+
+  def stop(self):
+    pass
+
+  def close_connection(self):
+    pass
+
+
 class SnippetManagementServiceTest(unittest.TestCase):
   """Tests for the snippet management service."""
 
@@ -99,27 +143,74 @@ def test_add_snippet_client_dup_name(self, _):
     )
     manager.add_snippet_client('foo', MOCK_PACKAGE)
     msg = (
-        '.* Name "foo" is already registered with package ".*", it '
-        'cannot be used again.'
+        '.* Name "foo" is already registered with package ".*" for user ID .*,'
+        ' the same name cannot be used again.'
     )
     with self.assertRaisesRegex(snippet_management_service.Error, msg):
       manager.add_snippet_client('foo', MOCK_PACKAGE + 'ha')
 
-  @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH)
-  def test_add_snippet_client_dup_package(self, mock_class):
-    mock_client = mock_class.return_value
-    mock_client.package = MOCK_PACKAGE
-    manager = snippet_management_service.SnippetManagementService(
-        mock.MagicMock()
-    )
+  @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH, new=MockSnippetClientV2)
+  def test_add_snippet_client_different_package(self):
+    mock_device = mock.MagicMock()
+    manager = snippet_management_service.SnippetManagementService(mock_device)
+    manager.add_snippet_client('foo', MOCK_PACKAGE)
+
+    try:
+      manager.add_snippet_client('bar', MOCK_PACKAGE2)
+    except snippet_management_service.Error as e:
+      self.fail(
+          'Should not fail when loading snippets with the different package'
+      )
+
+  @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH, new=MockSnippetClientV2)
+  def test_add_snippet_client_dup_package_and_none_as_snippet_config(self):
+    user_id = 2
+    mock_adb = mock.MagicMock(current_user_id=user_id)
+    mock_device = mock.MagicMock(adb=mock_adb)
+    manager = snippet_management_service.SnippetManagementService(mock_device)
     manager.add_snippet_client('foo', MOCK_PACKAGE)
     msg = (
-        'Snippet package "com.mock.package" has already been loaded '
-        'under name "foo".'
+        f'Snippet "com.mock.package" has already been registered for user id'
+        f' {user_id} under name "foo". The same package cannot be registered'
+        ' again for the same user.'
     )
     with self.assertRaisesRegex(snippet_management_service.Error, msg):
       manager.add_snippet_client('bar', MOCK_PACKAGE)
 
+  @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH, new=MockSnippetClientV2)
+  def test_add_snippet_client_dup_package_and_user_id(self):
+    user_id = 2
+    config = snippet_client_v2.Config(user_id=user_id)
+    mock_adb = mock.MagicMock(current_user_id=user_id)
+    mock_device = mock.MagicMock(adb=mock_adb)
+    manager = snippet_management_service.SnippetManagementService(mock_device)
+    manager.add_snippet_client('foo', MOCK_PACKAGE, config=config)
+    msg = (
+        f'Snippet "com.mock.package" has already been registered for user id'
+        f' {user_id} under name "foo". The same package cannot be registered'
+        ' again for the same user.'
+    )
+    with self.assertRaisesRegex(snippet_management_service.Error, msg):
+      manager.add_snippet_client('bar', MOCK_PACKAGE, config=config)
+
+  @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH, new=MockSnippetClientV2)
+  def test_add_snippet_client_dup_package_with_different_user_id(self):
+    old_user_id = 2
+    new_user_id = 3
+    old_config = snippet_client_v2.Config(user_id=old_user_id)
+    new_config = snippet_client_v2.Config(user_id=new_user_id)
+    mock_adb = mock.MagicMock(current_user_id=new_user_id)
+    mock_device = mock.MagicMock(adb=mock_adb)
+    manager = snippet_management_service.SnippetManagementService(mock_device)
+    manager.add_snippet_client('foo', MOCK_PACKAGE, old_config)
+    try:
+      manager.add_snippet_client('bar', MOCK_PACKAGE, new_config)
+    except snippet_management_service.Error as e:
+      self.fail(
+          'Should not fail when loading snippets with the same package but'
+          ' different user ID.'
+      )
+
   @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH)
   def test_remove_snippet_client(self, mock_class):
     mock_client = mock.MagicMock()

tests/mobly/controllers/android_device_lib/snippet_client_v2_test.py

@@ -1343,8 +1343,9 @@ def test_restore_event_client(
     with self.assertRaisesRegex(
         errors.ServerRestoreConnectionError,
         (
-            f'Failed to restore server connection for {MOCK_PACKAGE_NAME} at '
-            f'host port {host_port_3}, device port {MOCK_DEVICE_PORT}'
+            'Failed to restore server connection of the snippet package'
+            f' {MOCK_PACKAGE_NAME} for user id {MOCK_USER_ID} at host port'
+            f' {host_port_3}, device port {MOCK_DEVICE_PORT}'
         ),
     ):
       self.client.restore_server_connection()

tests/mobly/controllers/android_device_test.py

@@ -1552,14 +1552,24 @@ def test_AndroidDevice_getattr(
       return_value=MockSnippetClient,
   )
   @mock.patch('mobly.utils.get_available_host_port')
-  def test_AndroidDevice_load_snippet_dup_package(
+  def test_AndroidDevice_load_snippet_dup_identifier(
       self, MockGetPort, MockSnippetClient, MockFastboot, MockAdbProxy
   ):
     ad = android_device.AndroidDevice(serial='1')
+    user_id = 1
+    ad.adb.current_user_id = user_id
+    MockSnippetClient.return_value.user_id = user_id
+    MockSnippetClient.return_value.identifier = (
+        f'{MOCK_SNIPPET_PACKAGE_NAME}@user_id[{user_id}]'
+    )
+    MockSnippetClient.get_identifier.side_effect = (
+        lambda pacakge, user_id: f'{pacakge}@user_id[{user_id}]'
+    )
     ad.load_snippet('snippet', MOCK_SNIPPET_PACKAGE_NAME)
     expected_msg = (
-        'Snippet package "%s" has already been loaded under name "snippet".'
-    ) % MOCK_SNIPPET_PACKAGE_NAME
+        f'Snippet "{MOCK_SNIPPET_PACKAGE_NAME}" has already been registered for'
+        f' user id {user_id}'
+    )
     with self.assertRaisesRegex(android_device.Error, expected_msg):
       ad.load_snippet('snippet2', MOCK_SNIPPET_PACKAGE_NAME)