Cleanup UiAutomator connection so it won't block other from using since UiAutomator only support a single connection at a time.

1 files changed, 12 insertions, 3 deletions

diff --git a/src/main/java/com/google/android/mobly/snippet/bundled/WifiManagerSnippet.java b/src/main/java/com/google/android/mobly/snippet/bundled/WifiManagerSnippet.java
index 3c46708..55c7a20 100644
--- a/src/main/java/com/google/android/mobly/snippet/bundled/WifiManagerSnippet.java
+++ b/src/main/java/com/google/android/mobly/snippet/bundled/WifiManagerSnippet.java
@@ -63,9 +63,18 @@ public class WifiManagerSnippet implements Snippet {
                 (WifiManager)
                         mContext.getApplicationContext().getSystemService(Context.WIFI_SERVICE);
         if (Build.VERSION.SDK_INT >= 29) {
-            InstrumentationRegistry.getInstrumentation()
-                    .getUiAutomation()
-                    .adoptShellPermissionIdentity();
+            UiAutomation uia = InstrumentationRegistry.getInstrumentation().getUiAutomation();
+            uia.adoptShellPermissionIdentity();
+            try {
+                Class cls = Class.forName("android.app.UiAutomation");
+                Method destroyMethod = cls.getDeclaredMethod("destroy");
+                destroyMethod.invoke(uia);
+            } catch (NoSuchMethodException
+                  | IllegalAccessException
+                  | ClassNotFoundException
+                  | InvocationTargetException e) {
+                      throw new WifiManagerSnippetException("Failed to cleaup Ui Automation", e);
+            }
         }
     }