commit c057886579b19ba342de7214f753f4142ebd8653
author Kate Hsieh <yuchihsieh@google.com> Wed Dec 27 18:39:31 2023 -0800
committer Copybara-Service <copybara-worker@google.com> Wed Dec 27 18:40:02 2023 -0800
tree e64caccec94999f1a9578c730210ca52518c12fc
parent 3c95ae1daf556ef93acbac829549bc1b7d3b87fa

Unload snippet after actions have been done

PiperOrigin-RevId: 594135259

ui_automator/ui_automator.py

diff --git a/ui_automator/ui_automator.py b/ui_automator/ui_automator.py
index f351812..a9fab62 100644
--- a/ui_automator/ui_automator.py
+++ b/ui_automator/ui_automator.py
@@ -139,14 +139,15 @@
     try:
       self.load_device()
       self.load_snippet()
+      return func(self, *args, **kwargs)
     except (
         errors.AdbError,
         errors.AndroidDeviceNotReadyError,
         errors.NoAndroidDeviceError,
     ) as e:
       raise errors.AndroidDeviceNotReadyError(f'{func.__name__} failed.') from e
-
-    return func(self, *args, **kwargs)
+    finally:
+      self.unload_snippet()
 
   return wrapped_func
 
@@ -235,6 +236,13 @@
       # has been loaded.
       self._logger.debug(str(e))
 
+  def unload_snippet(self) -> None:
+    """Unloads loaded mobly snippet on connected android device."""
+    if not self._connected_device:
+      return
+
+    self._connected_device.unload_snippet(_MOBLY_SNIPPET_APK_NAME)
+
   @get_android_device_ready
   def commission_device(
       self,

ui_automator/ui_automator_test.py

diff --git a/ui_automator/ui_automator_test.py b/ui_automator/ui_automator_test.py
index 7e5cdef..c03b179 100644
--- a/ui_automator/ui_automator_test.py
+++ b/ui_automator/ui_automator_test.py
@@ -350,6 +350,22 @@
           _FAKE_MATTER_DEVICE_NAME, _FAKE_PAIRING_CODE, _FAKE_GHA_ROOM
       )
 
+  @mock.patch.object(ui_automator.UIAutomator, 'unload_snippet', autospec=True)
+  @mock.patch.object(android_device, 'get_all_instances', autospec=True)
+  def test_get_android_device_ready_calls_unload_snippet_on_finish(
+      self, mock_get_all_instances, mock_unload_snippet
+  ):
+    mock_get_all_instances.return_value = [self.mock_android_device]
+
+    @ui_automator.get_android_device_ready
+    def decorated_function(self):
+      del self
+      pass
+
+    decorated_function(self.ui_automator)
+
+    mock_unload_snippet.assert_called_once()
+
   @mock.patch.object(android_device, 'get_all_instances', autospec=True)
   def test_commission_device_calls_a_method_in_snippet_with_desired_args(
       self, mock_get_all_instances