Repository: google/mobly

Issue:
Snippet client reference is not removed from ad object after teardown
'load_snippet' sets a snippet client reference attribute on the AndroidDevice object.
We should remove the attribute after snippet is stopped.

Discussion:

Patch:
diff --git a/mobly/controllers/android_device.py b/mobly/controllers/android_device.py
index af00375..f42cf40 100644
--- a/mobly/controllers/android_device.py
+++ b/mobly/controllers/android_device.py
@@ -49,6 +49,10 @@ class Error(signals.ControllerError):
     pass
 
 
+class SnippetError(signals.ControllerError):
+    """Raised when somethig wrong with snippet happens."""
+
+
 class DoesNotExistError(Error):
     """Raised when something that does not exist is referenced."""
 
@@ -366,7 +370,9 @@ class AndroidDevice(object):
         self.fastboot = fastboot.FastbootProxy(serial)
         if not self.is_bootloader and self.is_rootable:
             self.root_adb()
-        self._snippet_clients = []
+        # A dict for tracking snippet clients. Keys are clients' attribute
+        # names, values are the clients: {<attr name string>: <client object>}.
+        self._snippet_clients = {}
 
     def set_logger_prefix_tag(self, tag):
         """Set a tag for the log line prefix of this instance.
@@ -413,9 +419,10 @@ class AndroidDevice(object):
         if self._adb_logcat_process:
             self.stop_adb_logcat()
         self._terminate_sl4a()
-        for client in self._snippet_clients:
+        for name, client in self._snippet_clients.items():
             self._terminate_jsonrpc_client(client)
-        self._snippet_clients = []
+            delattr(self, name)
+        self._snippet_clients = {}
 
     @property
     def build_info(self):
@@ -492,9 +499,8 @@ class AndroidDevice(object):
         """
         for k, v in config.items():
             if hasattr(self, k):
-                raise Error(
-                    "Attempting to set existing attribute %s on %s" %
-                    (k, self.serial))
+                raise Error("Attempting to set existing attribute %s on %s" %
+                            (k, self.serial))
             setattr(self, k, v)
 
     def root_adb(self):
@@ -509,18 +515,38 @@ class AndroidDevice(object):
     def load_snippet(self, name, package):
         """Starts the snippet apk with the given package name and connects.
 
-        Args:
-          name: The attribute name to which to attach the snippet server.
-            e.g. name='maps' will attach the snippet server to ad.maps.
-          package: The package name defined in AndroidManifest.xml of the
-            snippet apk.
-
         Examples:
             >>> ad = AndroidDevice()
             >>> ad.load_snippet(
                     name='maps', package='com.google.maps.snippets')
             >>> ad.maps.activateZoom('3')
+
+        Args:
+            name: The attribute name to which to attach the snippet server.
+                  e.g. name='maps' will attach the snippet server to ad.maps.
+            package: The package name defined in AndroidManifest.xml of the
+                     snippet apk.
+
+        Raises:
+            SnippetError is raised if illegal load operations are attempted.
         """
+        # Should not load snippet with the same attribute more than once.
+        if name in self._snippet_clients:
+            raise SnippetError(
+                ('Attribute "%s" is already registered with '
+                 'package "%s", it cannot be used again.') % (
+                     name, self._snippet_clients[name].package))
+        # Should not load snippet with an existing attribute.
+        if hasattr(self, name):
+            raise SnippetError(
+                'Attribute "%s" already exists, please use a different name.' %
+                name)
+        # Should not load the same snippet package more than once.
+        for client_name, client in self._snippet_clients.items():
+            if package == client.package:
+                raise SnippetError(
+                    'Snippet package "%s" has already been loaded under name "%s".'
+                    % (package, client_name))
         host_port = utils.get_available_host_port()
         # TODO(adorokhine): Don't assume that a free host-side port is free on
         # the device as well. Both sides should allocate a unique port.
@@ -528,7 +554,7 @@ class AndroidDevice(object):
         client = snippet_client.SnippetClient(
             package=package, port=host_port, adb_proxy=self.adb)
         self._start_jsonrpc_client(client, host_port, device_port)
-        self._snippet_clients.append(client)
+        self._snippet_clients[name] = client
         setattr(self, name, client)
 
     def _start_sl4a(self):
@@ -548,7 +574,8 @@ class AndroidDevice(object):
         # Start an EventDispatcher for the current sl4a session
         event_client = sl4a_client.Sl4aClient(self.adb)
         event_client.connect(
-            port=host_port, uid=self.sl4a.uid,
+            port=host_port,
+            uid=self.sl4a.uid,
             cmd=jsonrpc_client_base.JsonRpcCommand.CONTINUE)
         self.ed = event_dispatcher.EventDispatcher(event_client)
         self.ed.start()
@@ -577,7 +604,8 @@ class AndroidDevice(object):
         self.adb.forward("--remove tcp:%d" % client.port)
 
     def _is_timestamp_in_range(self, target, begin_time, end_time):
-        low = mobly_logger.logline_timestamp_comparator(begin_time, target) <= 0
+        low = mobly_logger.logline_timestamp_comparator(begin_time,
+                                                        target) <= 0
         high = mobly_logger.logline_timestamp_comparator(end_time, target) >= 0
         return low and high
 
@@ -636,8 +664,8 @@ class AndroidDevice(object):
         """
         if self._adb_logcat_process:
             raise Error(
-              'Android device %s already has an adb logcat thread going on. '
-              'Cannot start another one.' % self.serial)
+                'Android device %s already has an adb logcat thread going on. '
+                'Cannot start another one.' % self.serial)
         # Disable adb log spam filter for rootable. Have to stop and clear
         # settings first because 'start' doesn't support --clear option before
         # Android N.
@@ -697,7 +725,7 @@ class AndroidDevice(object):
             out = self.adb.shell("bugreportz").decode("utf-8")
             if not out.startswith("OK"):
                 raise Error("Failed to take bugreport on %s: %s" %
-                                         (self.serial, out))
+                            (self.serial, out))
             br_out_path = out.split(':')[1].strip()
             self.adb.pull("%s %s" % (br_out_path, full_out_path))
         else:
@@ -796,6 +824,7 @@ class AndroidDeviceLoggerAdapter(logging.LoggerAdapter):
         Then each log line added by my_log will have a prefix
         "[AndroidDevice|<tag>]"
     """
+
     def process(self, msg, kwargs):
         msg = "[AndroidDevice|%s] %s" % (self.extra["tag"], msg)
         return (msg, kwargs)
diff --git a/mobly/controllers/android_device_lib/snippet_client.py b/mobly/controllers/android_device_lib/snippet_client.py
index 813d9dd..8cf875b 100644
--- a/mobly/controllers/android_device_lib/snippet_client.py
+++ b/mobly/controllers/android_device_lib/snippet_client.py
@@ -50,6 +50,10 @@ class SnippetClient(jsonrpc_client_base.JsonRpcClientBase):
         self._package = package
         self._port = port
 
+    @property
+    def package(self):
+        return self._package
+
     def _do_start_app(self):
         """Overrides superclass."""
         cmd = _LAUNCH_CMD.format(self._port, self._package)
@@ -62,7 +66,7 @@ class SnippetClient(jsonrpc_client_base.JsonRpcClientBase):
         logging.info('Stopping snippet apk with: %s', cmd)
         out = self._adb.shell(_STOP_CMD.format(self._package)).decode('utf-8')
         if 'OK (0 tests)' not in out:
-            raise Error('Failed to stop existing apk. Unexpected output: ' +
+            raise Error('Failed to stop existing apk. Unexpected output: %s' %
                         out)
 
     def _is_app_installed(self):


Test patch:
diff --git a/tests/mobly/controllers/android_device_test.py b/tests/mobly/controllers/android_device_test.py
index 062de6b..d66b4dc 100755
--- a/tests/mobly/controllers/android_device_test.py
+++ b/tests/mobly/controllers/android_device_test.py
@@ -40,6 +40,11 @@ MOCK_ADB_LOGCAT = ("02-29 14:02:19.123  4454  Nothing\n"
 # Mock start and end time of the adb cat.
 MOCK_ADB_LOGCAT_BEGIN_TIME = "02-29 14:02:20.123"
 MOCK_ADB_LOGCAT_END_TIME = "02-29 14:02:22.000"
+MOCK_SNIPPET_PACKAGE_NAME = "com.my.snippet"
+
+# A mock SnippetClient used for testing snippet management logic.
+MockSnippetClient = mock.MagicMock()
+MockSnippetClient.package = MOCK_SNIPPET_PACKAGE_NAME
 
 
 class AndroidDeviceTest(unittest.TestCase):
@@ -340,6 +345,81 @@ class AndroidDeviceTest(unittest.TestCase):
         # Stops adb logcat.
         ad.stop_adb_logcat()
 
+    @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+                return_value=mock_android_device.MockAdbProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+                return_value=mock_android_device.MockFastbootProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.SnippetClient')
+    @mock.patch('mobly.utils.get_available_host_port')
+    def test_AndroidDevice_load_snippet(
+        self, MockGetPort, MockSnippetClient, MockFastboot, MockAdbProxy):
+        ad = android_device.AndroidDevice(serial=1)
+        ad.load_snippet('snippet', MOCK_SNIPPET_PACKAGE_NAME)
+        self.assertTrue(hasattr(ad, 'snippet'))
+
+    @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+                return_value=mock_android_device.MockAdbProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+                return_value=mock_android_device.MockFastbootProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.SnippetClient',
+                return_value=MockSnippetClient)
+    @mock.patch('mobly.utils.get_available_host_port')
+    def test_AndroidDevice_load_snippet_dup_package(
+        self, MockGetPort, MockSnippetClient, MockFastboot, MockAdbProxy):
+        ad = android_device.AndroidDevice(serial=1)
+        ad.load_snippet('snippet', MOCK_SNIPPET_PACKAGE_NAME)
+        expected_msg = ('Snippet package "%s" has already been loaded under '
+                        'name "snippet".') % MOCK_SNIPPET_PACKAGE_NAME
+        with self.assertRaisesRegexp(android_device.SnippetError,
+                                     expected_msg):
+            ad.load_snippet('snippet2', MOCK_SNIPPET_PACKAGE_NAME)
+
+    @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+                return_value=mock_android_device.MockAdbProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+                return_value=mock_android_device.MockFastbootProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.SnippetClient',
+                return_value=MockSnippetClient)
+    @mock.patch('mobly.utils.get_available_host_port')
+    def test_AndroidDevice_load_snippet_dup_snippet_name(
+        self, MockGetPort, MockSnippetClient, MockFastboot, MockAdbProxy):
+        ad = android_device.AndroidDevice(serial=1)
+        ad.load_snippet('snippet', MOCK_SNIPPET_PACKAGE_NAME)
+        expected_msg = ('Attribute "%s" is already registered with package '
+                        '"%s", it cannot be used again.') % (
+                        'snippet', MOCK_SNIPPET_PACKAGE_NAME)
+        with self.assertRaisesRegexp(android_device.SnippetError,
+                                     expected_msg):
+            ad.load_snippet('snippet', MOCK_SNIPPET_PACKAGE_NAME + 'haha')
+
+    @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+                return_value=mock_android_device.MockAdbProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+                return_value=mock_android_device.MockFastbootProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.SnippetClient')
+    @mock.patch('mobly.utils.get_available_host_port')
+    def test_AndroidDevice_load_snippet_dup_attribute_name(
+        self, MockGetPort, MockSnippetClient, MockFastboot, MockAdbProxy):
+        ad = android_device.AndroidDevice(serial=1)
+        expected_msg = ('Attribute "%s" already exists, please use a different'
+                        ' name') % 'adb'
+        with self.assertRaisesRegexp(android_device.SnippetError,
+                                     expected_msg):
+            ad.load_snippet('adb', MOCK_SNIPPET_PACKAGE_NAME)
+
+    @mock.patch('mobly.controllers.android_device_lib.adb.AdbProxy',
+                return_value=mock_android_device.MockAdbProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.fastboot.FastbootProxy',
+                return_value=mock_android_device.MockFastbootProxy(1))
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.SnippetClient')
+    @mock.patch('mobly.utils.get_available_host_port')
+    def test_AndroidDevice_snippet_cleanup(self, MockGetPort, MockSnippetClient,
+                                           MockFastboot, MockAdbProxy):
+        ad = android_device.AndroidDevice(serial=1)
+        ad.load_snippet('snippet', MOCK_SNIPPET_PACKAGE_NAME)
+        ad.stop_services()
+        self.assertFalse(hasattr(ad, 'snippet'))
+
 
 if __name__ == "__main__":
     unittest.main()