diff --git a/mobly/controllers/android_device_lib/services/snippet_management_service.py b/mobly/controllers/android_device_lib/services/snippet_management_service.py index 7da1b23d..2e1a604b 100644 --- a/mobly/controllers/android_device_lib/services/snippet_management_service.py +++ b/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) diff --git a/mobly/controllers/android_device_lib/snippet_client_v2.py b/mobly/controllers/android_device_lib/snippet_client_v2.py index ee0b2ecd..7963fd50 100644 --- a/mobly/controllers/android_device_lib/snippet_client_v2.py +++ b/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 diff --git a/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py b/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py index b6431d54..f2af7325 100755 --- a/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py +++ b/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() 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 13f57821..7dc8a976 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 @@ -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() diff --git a/tests/mobly/controllers/android_device_test.py b/tests/mobly/controllers/android_device_test.py index 640f28c4..9a03473e 100755 --- a/tests/mobly/controllers/android_device_test.py +++ b/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)