From 4c3fec11e3a4dacec58d04995b61846675966e76 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Tue, 24 Dec 2024 17:11:08 -0800
Subject: [PATCH 1/4] Install and test LSPosed through test app

---
 app/core/build.gradle.kts                     |  1 +
 .../topjohnwu/magisk/test/AdditionalTest.kt   | 57 +++++++++++++++++
 .../com/topjohnwu/magisk/test/Environment.kt  | 64 ++++++++++++++++---
 app/test/build.gradle.kts                     |  1 +
 .../com/topjohnwu/magisk/test/TestRunner.kt   | 11 ++++
 gradle/libs.versions.toml                     |  1 +
 scripts/avd_test.sh                           | 29 ---------
 scripts/test_common.sh                        | 21 +++---
 8 files changed, 136 insertions(+), 49 deletions(-)
 create mode 100644 app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt

diff --git a/app/core/build.gradle.kts b/app/core/build.gradle.kts
index 0f8f54cb8094..d8569a6ed138 100644
--- a/app/core/build.gradle.kts
+++ b/app/core/build.gradle.kts
@@ -65,4 +65,5 @@ dependencies {
     // However, we don't want to bundle test dependencies.
     // That's why we make it compileOnly.
     compileOnly(libs.test.junit)
+    compileOnly(libs.test.uiautomator)
 }
diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt
new file mode 100644
index 000000000000..f157ebe8c667
--- /dev/null
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt
@@ -0,0 +1,57 @@
+package com.topjohnwu.magisk.test
+
+import android.app.UiAutomation
+import androidx.annotation.Keep
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import androidx.test.uiautomator.By
+import androidx.test.uiautomator.UiDevice
+import androidx.test.uiautomator.Until
+import org.junit.After
+import org.junit.Assert.assertNotNull
+import org.junit.Assume.assumeTrue
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import java.util.concurrent.TimeUnit
+import java.util.regex.Pattern
+
+@Keep
+@RunWith(AndroidJUnit4::class)
+class AdditionalTest {
+
+    companion object {
+        private const val SHELL_PKG = "com.android.shell"
+        private const val LSPOSED_CATEGORY = "org.lsposed.manager.LAUNCH_MANAGER"
+        private const val LSPOSED_PKG = "org.lsposed.manager"
+    }
+
+    private lateinit var uiAutomation: UiAutomation
+    private lateinit var device: UiDevice
+
+    @Before
+    fun setup() {
+        val inst = InstrumentationRegistry.getInstrumentation()
+        uiAutomation = inst.uiAutomation
+        device = UiDevice.getInstance(inst)
+    }
+
+    @After
+    fun teardown() {
+        device.pressHome()
+    }
+
+    @Test
+    fun testLaunchLsposedManager() {
+        assumeTrue(Environment.lsposed())
+
+        uiAutomation.executeShellCommand(
+            "am start -c $LSPOSED_CATEGORY $SHELL_PKG/.BugreportWarningActivity"
+        )
+        val pattern = Pattern.compile("$LSPOSED_PKG:id/.*")
+        assertNotNull(
+            "LSPosed manager launch failed",
+            device.wait(Until.hasObject(By.res(pattern)), TimeUnit.SECONDS.toMillis(10))
+        )
+    }
+}
diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
index 6f304c564159..dccbea056452 100644
--- a/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
@@ -1,17 +1,27 @@
 package com.topjohnwu.magisk.test
 
+import android.app.Notification
+import android.content.Context
+import android.os.Build
 import androidx.annotation.Keep
+import androidx.core.net.toUri
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
 import com.topjohnwu.magisk.core.BuildConfig.APP_PACKAGE_NAME
 import com.topjohnwu.magisk.core.Config
 import com.topjohnwu.magisk.core.di.ServiceLocator
+import com.topjohnwu.magisk.core.download.DownloadNotifier
+import com.topjohnwu.magisk.core.download.DownloadProcessor
+import com.topjohnwu.magisk.core.ktx.cachedFile
 import com.topjohnwu.magisk.core.model.su.SuPolicy
 import com.topjohnwu.magisk.core.tasks.AppMigration
+import com.topjohnwu.magisk.core.tasks.FlashZip
 import com.topjohnwu.magisk.core.tasks.MagiskInstaller
 import com.topjohnwu.superuser.CallbackList
 import kotlinx.coroutines.runBlocking
 import org.junit.Assert.assertTrue
+import org.junit.Assume.assumeTrue
+import org.junit.Before
 import org.junit.BeforeClass
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -25,19 +35,55 @@ class Environment {
         @BeforeClass
         @JvmStatic
         fun before() = MagiskAppTest.before()
+
+        fun lsposed(): Boolean {
+            return Build.VERSION.SDK_INT >= 27 && Build.VERSION.SDK_INT <= 34
+        }
+
+        private const val LSPOSED_URL =
+            "https://github.com/LSPosed/LSPosed/releases/download/v1.9.2/LSPosed-v1.9.2-7024-zygisk-release.zip"
+    }
+
+    object TimberLog : CallbackList<String>(Runnable::run) {
+        override fun onAddElement(e: String) {
+            Timber.i(e)
+        }
+    }
+
+    private lateinit var mContext: Context
+
+    @Before
+    fun setup() {
+        mContext = InstrumentationRegistry.getInstrumentation().targetContext
     }
 
     @Test
     fun setupMagisk() {
-        val log = object : CallbackList<String>(Runnable::run) {
-            override fun onAddElement(e: String) {
-                Timber.i(e)
-            }
-        }
         runBlocking {
             assertTrue(
                 "Magisk setup failed",
-                MagiskInstaller.Emulator(log, log).exec()
+                MagiskInstaller.Emulator(TimberLog, TimberLog).exec()
+            )
+        }
+    }
+
+    @Test
+    fun setupLsposed() {
+        assumeTrue(lsposed())
+
+        val notify = object : DownloadNotifier {
+            override val context = mContext
+            override fun notifyUpdate(id: Int, editor: (Notification.Builder) -> Unit) {}
+        }
+        val processor = DownloadProcessor(notify)
+        val zip = mContext.cachedFile("lsposed.zip")
+        runBlocking {
+            ServiceLocator.networkService.fetchFile(LSPOSED_URL).byteStream().use {
+                processor.handleModule(it, zip.toUri())
+            }
+            assertTrue(
+                "LSPosed installation failed",
+                FlashZip(zip.toUri(), TimberLog, TimberLog).exec()
             )
         }
     }
@@ -65,7 +111,7 @@ class Environment {
             assertTrue(
                 "App hiding failed",
                 AppMigration.patchAndHide(
-                    context = InstrumentationRegistry.getInstrumentation().targetContext,
+                    context = mContext,
                     label = "Settings",
                     pkg = "repackaged.$APP_PACKAGE_NAME"
                 )
@@ -78,9 +124,7 @@ class Environment {
         runBlocking {
             assertTrue(
                 "App restoration failed",
-                AppMigration.restoreApp(
-                    context = InstrumentationRegistry.getInstrumentation().targetContext
-                )
+                AppMigration.restoreApp(mContext)
             )
         }
     }
diff --git a/app/test/build.gradle.kts b/app/test/build.gradle.kts
index 92ae38c9f613..c3c8ff705ab0 100644
--- a/app/test/build.gradle.kts
+++ b/app/test/build.gradle.kts
@@ -26,4 +26,5 @@ dependencies {
     implementation(libs.test.runner)
     implementation(libs.test.rules)
     implementation(libs.test.junit)
+    implementation(libs.test.uiautomator)
 }
diff --git a/app/test/src/main/java/com/topjohnwu/magisk/test/TestRunner.kt b/app/test/src/main/java/com/topjohnwu/magisk/test/TestRunner.kt
index c92080a96c14..e8dd83edb448 100644
--- a/app/test/src/main/java/com/topjohnwu/magisk/test/TestRunner.kt
+++ b/app/test/src/main/java/com/topjohnwu/magisk/test/TestRunner.kt
@@ -6,6 +6,17 @@ import androidx.test.runner.AndroidJUnitRunner
 
 class TestRunner : AndroidJUnitRunner() {
     override fun onCreate(arguments: Bundle) {
+        // Support short-hand ".ClassName"
+        arguments.getString("class")?.let {
+            val classArg = it.split(",").joinToString(separator = ",") { clz ->
+                if (clz.startsWith(".")) {
+                    "com.topjohnwu.magisk.test$clz"
+                } else {
+                    clz
+                }
+            }
+            arguments.putString("class", classArg)
+        }
         // Force using the target context's classloader to run tests
         arguments.putString("classLoader", TestClassLoader::class.java.name)
         super.onCreate(arguments)
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 0e259f0a73dc..d14c6a11724b 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -47,6 +47,7 @@ jdk-libs = { module = "com.android.tools:desugar_jdk_libs_nio", version = "2.1.3
 test-runner = { module = "androidx.test:runner", version = "1.6.2" }
 test-rules = { module = "androidx.test:rules", version = "1.6.1" }
 test-junit = { module = "androidx.test.ext:junit", version = "1.2.1" }
+test-uiautomator = { module = "androidx.test.uiautomator:uiautomator", version = "2.3.0" }
 
 # topjohnwu
 indeterminate-checkbox = { module = "com.github.topjohnwu:indeterminate-checkbox", version = "1.0.7" }
diff --git a/scripts/avd_test.sh b/scripts/avd_test.sh
index fb6e9f7e98d0..2a9eb62640fe 100755
--- a/scripts/avd_test.sh
+++ b/scripts/avd_test.sh
@@ -4,13 +4,10 @@ set -xe
 . scripts/test_common.sh
 
 emu_args_base="-no-window -no-audio -no-boot-anim -gpu swiftshader_indirect -read-only -no-snapshot -cores $core_count"
-lsposed_url='https://github.com/LSPosed/LSPosed/releases/download/v1.9.2/LSPosed-v1.9.2-7024-zygisk-release.zip'
 emu_pid=
 
 atd_min_api=30
 atd_max_api=35
-lsposed_min_api=27
-lsposed_max_api=34
 huge_ram_min_api=26
 
 cleanup() {
@@ -81,35 +78,10 @@ test_emu() {
 
   run_setup $variant
 
-  local lsposed
-  if [ $api -ge $lsposed_min_api -a $api -le $lsposed_max_api ]; then
-    lsposed=true
-  else
-    lsposed=false
-  fi
-
-  # Install LSPosed
-  if $lsposed; then
-    adb push out/lsposed.zip /data/local/tmp/lsposed.zip
-    echo 'PATH=$PATH:/debug_ramdisk magisk --install-module /data/local/tmp/lsposed.zip' | adb shell /system/xbin/su
-  fi
-
   adb reboot
   wait_emu wait_for_boot
 
   run_tests
-
-  # Try to launch LSPosed
-  if $lsposed; then
-    adb shell rm -f /data/local/tmp/window_dump.xml
-    adb shell am start -c org.lsposed.manager.LAUNCH_MANAGER com.android.shell/.BugreportWarningActivity
-    while adb shell '[ ! -f /data/local/tmp/window_dump.xml ]'; do
-      sleep 10
-      adb shell uiautomator dump /data/local/tmp/window_dump.xml
-    done
-    adb shell grep -q org.lsposed.manager /data/local/tmp/window_dump.xml
-    adb pull /data/local/tmp/window_dump.xml
-  fi
 }
 
 test_main() {
@@ -220,7 +192,6 @@ if [ -n "$FORCE_32_BIT" ]; then
 fi
 
 yes | "$sdk" --licenses > /dev/null
-curl -L $lsposed_url -o out/lsposed.zip
 "$sdk" --channel=3 platform-tools emulator
 
 adb kill-server
diff --git a/scripts/test_common.sh b/scripts/test_common.sh
index a9eb4e8e0bb8..341de8e4ae34 100644
--- a/scripts/test_common.sh
+++ b/scripts/test_common.sh
@@ -36,8 +36,7 @@ am_instrument() {
   else
     test_pkg=com.topjohnwu.magisk.test
   fi
-  local out=$(adb shell am instrument -w --user 0 \
-    -e class "com.topjohnwu.magisk.test.$1" \
+  local out=$(adb shell am instrument -w --user 0 -e class "$1" \
     "$test_pkg/com.topjohnwu.magisk.test.TestRunner")
   grep -q 'OK (' <<< "$out"
 }
@@ -59,34 +58,36 @@ run_setup() {
   adb install -r -g out/test.apk
 
   # Run setup through the test app
-  am_instrument 'Environment#setupMagisk'
+  am_instrument '.Environment#setupMagisk'
+  # Install LSPosed
+  am_instrument '.Environment#setupLsposed'
 }
 
 run_tests() {
   # Run app tests
-  am_instrument 'MagiskAppTest'
+  am_instrument '.MagiskAppTest,.AdditionalTest'
 
   # Test shell su request
-  am_instrument 'Environment#setupShellGrantTest'
+  am_instrument '.Environment#setupShellGrantTest'
   adb shell /system/xbin/su 2000 su -c id | tee /dev/fd/2 | grep -q 'uid=0'
   adb shell am force-stop com.topjohnwu.magisk
 
   # Test app hiding
-  am_instrument 'Environment#setupAppHide'
+  am_instrument '.Environment#setupAppHide'
   wait_for_pm com.topjohnwu.magisk
 
   # Make sure it still works
-  am_instrument 'MagiskAppTest' true
+  am_instrument '.MagiskAppTest' true
 
   # Test shell su request
-  am_instrument 'Environment#setupShellGrantTest' true
+  am_instrument '.Environment#setupShellGrantTest' true
   adb shell /system/xbin/su 2000 su -c id | tee /dev/fd/2 | grep -q 'uid=0'
   adb shell am force-stop repackaged.com.topjohnwu.magisk
 
   # Test app restore
-  am_instrument 'Environment#setupAppRestore' true
+  am_instrument '.Environment#setupAppRestore' true
   wait_for_pm repackaged.com.topjohnwu.magisk
 
   # Make sure it still works
-  am_instrument 'MagiskAppTest'
+  am_instrument '.MagiskAppTest'
 }

From aa13d115eed4c0f836ab7f8b2a9ef7cc83e4df1f Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Tue, 24 Dec 2024 23:11:53 -0800
Subject: [PATCH 2/4] Test su request via instrumentation

---
 .../com/topjohnwu/magisk/test/Environment.kt  | 27 ++-----
 .../topjohnwu/magisk/test/MagiskAppTest.kt    | 80 +++++++++++++++++--
 scripts/test_common.sh                        | 10 ---
 3 files changed, 80 insertions(+), 37 deletions(-)

diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
index dccbea056452..c5597161b663 100644
--- a/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
@@ -8,16 +8,16 @@ import androidx.core.net.toUri
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
 import com.topjohnwu.magisk.core.BuildConfig.APP_PACKAGE_NAME
-import com.topjohnwu.magisk.core.Config
 import com.topjohnwu.magisk.core.di.ServiceLocator
 import com.topjohnwu.magisk.core.download.DownloadNotifier
 import com.topjohnwu.magisk.core.download.DownloadProcessor
 import com.topjohnwu.magisk.core.ktx.cachedFile
-import com.topjohnwu.magisk.core.model.su.SuPolicy
 import com.topjohnwu.magisk.core.tasks.AppMigration
 import com.topjohnwu.magisk.core.tasks.FlashZip
 import com.topjohnwu.magisk.core.tasks.MagiskInstaller
+import com.topjohnwu.magisk.core.utils.RootUtils
 import com.topjohnwu.superuser.CallbackList
+import com.topjohnwu.superuser.Shell
 import kotlinx.coroutines.runBlocking
 import org.junit.Assert.assertTrue
 import org.junit.Assume.assumeTrue
@@ -34,7 +34,11 @@ class Environment {
     companion object {
         @BeforeClass
         @JvmStatic
-        fun before() = MagiskAppTest.before()
+        fun before() {
+            assertTrue("Should have root access", Shell.getShell().isRoot)
+            // Make sure the root service is running
+            RootUtils.Connection.await()
+        }
 
         fun lsposed(): Boolean {
             return Build.VERSION.SDK_INT >= 27 && Build.VERSION.SDK_INT <= 34
@@ -88,23 +92,6 @@ class Environment {
         }
     }
 
-    @Test
-    fun setupShellGrantTest() {
-        runBlocking {
-            // Inject an undetermined + mute logging policy for ADB shell
-            val policy = SuPolicy(
-                uid = 2000,
-                logging = false,
-                notification = false,
-                until = 0L
-            )
-            ServiceLocator.policyDB.update(policy)
-            // Bypass the need to actually show a dialog
-            Config.suAutoResponse = Config.Value.SU_AUTO_ALLOW
-            Config.prefs.edit().commit()
-        }
-    }
-
     @Test
     fun setupAppHide() {
         runBlocking {
diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt
index 61d23e1c9801..a06f9c977ca1 100644
--- a/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt
@@ -1,14 +1,26 @@
 package com.topjohnwu.magisk.test
 
+import android.app.Instrumentation
+import android.content.Intent
+import android.content.IntentFilter
+import android.os.Build
 import androidx.annotation.Keep
 import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import com.topjohnwu.magisk.core.Config
 import com.topjohnwu.magisk.core.Info
-import com.topjohnwu.magisk.core.utils.RootUtils
-import com.topjohnwu.superuser.Shell
+import com.topjohnwu.magisk.core.di.ServiceLocator
+import com.topjohnwu.magisk.core.model.su.SuPolicy
+import kotlinx.coroutines.runBlocking
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertNotNull
 import org.junit.Assert.assertTrue
+import org.junit.Before
 import org.junit.BeforeClass
 import org.junit.Test
 import org.junit.runner.RunWith
+import java.io.FileInputStream
+import java.util.concurrent.TimeUnit
 
 @Keep
 @RunWith(AndroidJUnit4::class)
@@ -17,15 +29,69 @@ class MagiskAppTest {
     companion object {
         @BeforeClass
         @JvmStatic
-        fun before() {
-            assertTrue("Should have root access", Shell.getShell().isRoot)
-            // Make sure the root service is running
-            RootUtils.Connection.await()
-        }
+        fun before() = Environment.before()
+    }
+
+    private lateinit var inst: Instrumentation
+    private val uiAutomation get() = inst.uiAutomation
+
+    @Before
+    fun setup() {
+        inst = InstrumentationRegistry.getInstrumentation()
     }
 
     @Test
     fun testZygisk() {
         assertTrue("Zygisk should be enabled", Info.isZygiskEnabled)
     }
+
+    @Test
+    fun testSuRequest() {
+        // Bypass the need to actually show a dialog
+        Config.suAutoResponse = Config.Value.SU_AUTO_ALLOW
+        Config.prefs.edit().commit()
+
+        // Inject an undetermined + mute logging policy for ADB shell
+        val policy = SuPolicy(
+            uid = 2000,
+            logging = false,
+            notification = false,
+            until = 0L
+        )
+        runBlocking {
+            ServiceLocator.policyDB.update(policy)
+        }
+
+        val filter = IntentFilter(Intent.ACTION_VIEW)
+        filter.addCategory(Intent.CATEGORY_DEFAULT)
+        val monitor = inst.addMonitor(filter, null, false)
+
+        // Try to call su from ADB shell
+        val cmd = if (Build.VERSION.SDK_INT < 24) {
+            // API 23 runs executeShellCommand as root
+            "/system/xbin/su 2000 su -c id"
+        } else {
+            "su -c id"
+        }
+        val pfd = uiAutomation.executeShellCommand(cmd)
+
+        // Make sure SuRequestActivity is launched
+        val suRequest = monitor.waitForActivityWithTimeout(TimeUnit.SECONDS.toMillis(10))
+        assertNotNull("SuRequestActivity is not launched", suRequest)
+
+        // Check that the request went through
+        FileInputStream(pfd.fileDescriptor).use {
+            assertTrue(
+                "Cannot grant root permission from shell",
+                it.reader().readText().contains("uid=0")
+            )
+        }
+
+        // Check that the database is updated
+        runBlocking {
+            val policy = ServiceLocator.policyDB.fetch(2000)
+                ?: throw AssertionError("PolicyDB is invalid")
+            assertEquals("Policy for shell is incorrect", SuPolicy.ALLOW, policy.policy)
+        }
+    }
 }
diff --git a/scripts/test_common.sh b/scripts/test_common.sh
index 341de8e4ae34..310541b07d84 100644
--- a/scripts/test_common.sh
+++ b/scripts/test_common.sh
@@ -67,11 +67,6 @@ run_tests() {
   # Run app tests
   am_instrument '.MagiskAppTest,.AdditionalTest'
 
-  # Test shell su request
-  am_instrument '.Environment#setupShellGrantTest'
-  adb shell /system/xbin/su 2000 su -c id | tee /dev/fd/2 | grep -q 'uid=0'
-  adb shell am force-stop com.topjohnwu.magisk
-
   # Test app hiding
   am_instrument '.Environment#setupAppHide'
   wait_for_pm com.topjohnwu.magisk
@@ -79,11 +74,6 @@ run_tests() {
   # Make sure it still works
   am_instrument '.MagiskAppTest' true
 
-  # Test shell su request
-  am_instrument '.Environment#setupShellGrantTest' true
-  adb shell /system/xbin/su 2000 su -c id | tee /dev/fd/2 | grep -q 'uid=0'
-  adb shell am force-stop repackaged.com.topjohnwu.magisk
-
   # Test app restore
   am_instrument '.Environment#setupAppRestore' true
   wait_for_pm repackaged.com.topjohnwu.magisk

From 88ead5fd8b7010b0e1596fb2475e0683bebb7d01 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Wed, 25 Dec 2024 04:29:02 -0800
Subject: [PATCH 3/4] Drive app migration tests through instrumentation

Make tests less flaky
---
 app/test/src/main/AndroidManifest.xml         | 13 ++-
 .../topjohnwu/magisk/test/AppMigrationTest.kt | 88 +++++++++++++++++++
 .../magisk/test/{TestRunner.kt => Runners.kt} |  8 +-
 scripts/test_common.sh                        | 33 ++++---
 4 files changed, 121 insertions(+), 21 deletions(-)
 create mode 100644 app/test/src/main/java/com/topjohnwu/magisk/test/AppMigrationTest.kt
 rename app/test/src/main/java/com/topjohnwu/magisk/test/{TestRunner.kt => Runners.kt} (84%)

diff --git a/app/test/src/main/AndroidManifest.xml b/app/test/src/main/AndroidManifest.xml
index fa7d5c44daa2..f6db926739e1 100644
--- a/app/test/src/main/AndroidManifest.xml
+++ b/app/test/src/main/AndroidManifest.xml
@@ -2,13 +2,22 @@
 <manifest xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:tools="http://schemas.android.com/tools">
 
+    <uses-permission
+        android:name="android.permission.QUERY_ALL_PACKAGES"
+        tools:ignore="QueryAllPackagesPermission" />
+
+    <queries tools:node="removeAll" />
+
     <application tools:node="replace">
         <uses-library android:name="android.test.runner" />
     </application>
 
+    <instrumentation
+        android:name="com.topjohnwu.magisk.test.AppTestRunner"
+        android:targetPackage="com.topjohnwu.magisk" />
+
     <instrumentation
         android:name="com.topjohnwu.magisk.test.TestRunner"
-        android:targetPackage="com.topjohnwu.magisk"
-        android:label="Tests for Magisk" />
+        android:targetPackage="com.topjohnwu.magisk.test" />
 
 </manifest>
diff --git a/app/test/src/main/java/com/topjohnwu/magisk/test/AppMigrationTest.kt b/app/test/src/main/java/com/topjohnwu/magisk/test/AppMigrationTest.kt
new file mode 100644
index 000000000000..7123d47eb726
--- /dev/null
+++ b/app/test/src/main/java/com/topjohnwu/magisk/test/AppMigrationTest.kt
@@ -0,0 +1,88 @@
+package com.topjohnwu.magisk.test
+
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import android.os.ParcelFileDescriptor.AutoCloseInputStream
+import androidx.annotation.Keep
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import org.junit.After
+import org.junit.Assert.assertTrue
+import org.junit.Test
+import org.junit.runner.RunWith
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+
+@Keep
+@RunWith(AndroidJUnit4::class)
+class AppMigrationTest {
+
+    companion object {
+        private const val APP_PKG = "com.topjohnwu.magisk"
+        private const val STUB_PKG = "repackaged.$APP_PKG"
+        private const val RECEIVER_TIMEOUT = 20L
+    }
+
+    private val instrumentation get() = InstrumentationRegistry.getInstrumentation()
+    private val context get() = instrumentation.context
+    private val uiAutomation get() = instrumentation.uiAutomation
+    private val registeredReceivers = mutableListOf<BroadcastReceiver>()
+
+    class PackageRemoveMonitor(
+        context: Context,
+        private val packageName: String
+    ) : BroadcastReceiver() {
+
+        val latch = CountDownLatch(1)
+
+        init {
+            val filter = IntentFilter(Intent.ACTION_PACKAGE_REMOVED)
+            filter.addDataScheme("package")
+            context.registerReceiver(this, filter)
+        }
+
+        override fun onReceive(context: Context, intent: Intent) {
+            if (intent.action != Intent.ACTION_PACKAGE_REMOVED)
+                return
+            val data = intent.data ?: return
+            val pkg = data.schemeSpecificPart
+            if (pkg == packageName) latch.countDown()
+        }
+    }
+
+    @After
+    fun tearDown() {
+        registeredReceivers.forEach(context::unregisterReceiver)
+    }
+
+    private fun testAppMigration(pkg: String, method: String) {
+        val receiver = PackageRemoveMonitor(context, pkg)
+        registeredReceivers.add(receiver)
+
+        // Trigger the test to run migration
+        val pfd = uiAutomation.executeShellCommand(
+            "am instrument -w --user 0 -e class .Environment#$method " +
+                "$pkg.test/${AppTestRunner::class.java.name}"
+        )
+        val output = AutoCloseInputStream(pfd).reader().use { it.readText() }
+        assertTrue("$method failed, inst out: $output", output.contains("OK ("))
+
+        // Wait for migration to complete
+        assertTrue(
+            "$pkg uninstallation failed",
+            receiver.latch.await(RECEIVER_TIMEOUT, TimeUnit.SECONDS)
+        )
+    }
+
+    @Test
+    fun testAppHide() {
+        testAppMigration(APP_PKG, "setupAppHide")
+    }
+
+    @Test
+    fun testAppRestore() {
+        testAppMigration(STUB_PKG, "setupAppRestore")
+    }
+}
diff --git a/app/test/src/main/java/com/topjohnwu/magisk/test/TestRunner.kt b/app/test/src/main/java/com/topjohnwu/magisk/test/Runners.kt
similarity index 84%
rename from app/test/src/main/java/com/topjohnwu/magisk/test/TestRunner.kt
rename to app/test/src/main/java/com/topjohnwu/magisk/test/Runners.kt
index e8dd83edb448..29a0523da5a5 100644
--- a/app/test/src/main/java/com/topjohnwu/magisk/test/TestRunner.kt
+++ b/app/test/src/main/java/com/topjohnwu/magisk/test/Runners.kt
@@ -4,7 +4,7 @@ import android.os.Bundle
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.runner.AndroidJUnitRunner
 
-class TestRunner : AndroidJUnitRunner() {
+open class TestRunner : AndroidJUnitRunner() {
     override fun onCreate(arguments: Bundle) {
         // Support short-hand ".ClassName"
         arguments.getString("class")?.let {
@@ -17,6 +17,12 @@ class TestRunner : AndroidJUnitRunner() {
             }
             arguments.putString("class", classArg)
         }
+        super.onCreate(arguments)
+    }
+}
+
+class AppTestRunner : TestRunner() {
+    override fun onCreate(arguments: Bundle) {
         // Force using the target context's classloader to run tests
         arguments.putString("classLoader", TestClassLoader::class.java.name)
         super.onCreate(arguments)
diff --git a/scripts/test_common.sh b/scripts/test_common.sh
index 310541b07d84..c7a5216f2577 100644
--- a/scripts/test_common.sh
+++ b/scripts/test_common.sh
@@ -28,16 +28,9 @@ print_error() {
 }
 
 # $1 = TestClass#method
-# $2: boolean = isRepackaged
+# $2 = component
 am_instrument() {
-  local test_pkg
-  if [ -n "$2" -a "$2" ]; then
-    test_pkg="repackaged.com.topjohnwu.magisk.test"
-  else
-    test_pkg=com.topjohnwu.magisk.test
-  fi
-  local out=$(adb shell am instrument -w --user 0 -e class "$1" \
-    "$test_pkg/com.topjohnwu.magisk.test.TestRunner")
+  local out=$(adb shell am instrument -w --user 0 -e class "$1" "$2")
   grep -q 'OK (' <<< "$out"
 }
 
@@ -57,27 +50,31 @@ run_setup() {
   # Install the test app
   adb install -r -g out/test.apk
 
+  local app='com.topjohnwu.magisk.test/com.topjohnwu.magisk.test.AppTestRunner'
+
   # Run setup through the test app
-  am_instrument '.Environment#setupMagisk'
+  am_instrument '.Environment#setupMagisk' $app
   # Install LSPosed
-  am_instrument '.Environment#setupLsposed'
+  am_instrument '.Environment#setupLsposed' $app
 }
 
 run_tests() {
+  local self='com.topjohnwu.magisk.test/com.topjohnwu.magisk.test.TestRunner'
+  local app='com.topjohnwu.magisk.test/com.topjohnwu.magisk.test.AppTestRunner'
+  local stub='repackaged.com.topjohnwu.magisk.test/com.topjohnwu.magisk.test.AppTestRunner'
+
   # Run app tests
-  am_instrument '.MagiskAppTest,.AdditionalTest'
+  am_instrument '.MagiskAppTest,.AdditionalTest' $app
 
   # Test app hiding
-  am_instrument '.Environment#setupAppHide'
-  wait_for_pm com.topjohnwu.magisk
+  am_instrument '.AppMigrationTest#testAppHide' $self
 
   # Make sure it still works
-  am_instrument '.MagiskAppTest' true
+  am_instrument '.MagiskAppTest' $stub
 
   # Test app restore
-  am_instrument '.Environment#setupAppRestore' true
-  wait_for_pm repackaged.com.topjohnwu.magisk
+  am_instrument '.AppMigrationTest#testAppRestore' $self
 
   # Make sure it still works
-  am_instrument '.MagiskAppTest'
+  am_instrument '.MagiskAppTest' $app
 }

From db5973096b778c167803ff60172204d6365d9921 Mon Sep 17 00:00:00 2001
From: topjohnwu <topjohnwu@gmail.com>
Date: Wed, 25 Dec 2024 04:50:08 -0800
Subject: [PATCH 4/4] Cleanup test code

---
 .../topjohnwu/magisk/test/AdditionalTest.kt   | 20 +++----------
 .../com/topjohnwu/magisk/test/BaseTest.kt     | 26 +++++++++++++++++
 .../com/topjohnwu/magisk/test/Environment.kt  | 28 ++++---------------
 .../topjohnwu/magisk/test/MagiskAppTest.kt    | 23 ++++-----------
 scripts/test_common.sh                        |  7 +++--
 5 files changed, 46 insertions(+), 58 deletions(-)
 create mode 100644 app/core/src/main/java/com/topjohnwu/magisk/test/BaseTest.kt

diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt
index f157ebe8c667..dcba012f502d 100644
--- a/app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/AdditionalTest.kt
@@ -1,16 +1,13 @@
 package com.topjohnwu.magisk.test
 
-import android.app.UiAutomation
+import android.os.ParcelFileDescriptor.AutoCloseInputStream
 import androidx.annotation.Keep
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.uiautomator.By
-import androidx.test.uiautomator.UiDevice
 import androidx.test.uiautomator.Until
 import org.junit.After
 import org.junit.Assert.assertNotNull
 import org.junit.Assume.assumeTrue
-import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import java.util.concurrent.TimeUnit
@@ -18,7 +15,7 @@ import java.util.regex.Pattern
 
 @Keep
 @RunWith(AndroidJUnit4::class)
-class AdditionalTest {
+class AdditionalTest : BaseTest {
 
     companion object {
         private const val SHELL_PKG = "com.android.shell"
@@ -26,16 +23,6 @@ class AdditionalTest {
         private const val LSPOSED_PKG = "org.lsposed.manager"
     }
 
-    private lateinit var uiAutomation: UiAutomation
-    private lateinit var device: UiDevice
-
-    @Before
-    fun setup() {
-        val inst = InstrumentationRegistry.getInstrumentation()
-        uiAutomation = inst.uiAutomation
-        device = UiDevice.getInstance(inst)
-    }
-
     @After
     fun teardown() {
         device.pressHome()
@@ -47,7 +34,8 @@ class AdditionalTest {
 
         uiAutomation.executeShellCommand(
             "am start -c $LSPOSED_CATEGORY $SHELL_PKG/.BugreportWarningActivity"
-        )
+        ).let { pfd -> AutoCloseInputStream(pfd).use { it.readBytes() } }
+
         val pattern = Pattern.compile("$LSPOSED_PKG:id/.*")
         assertNotNull(
             "LSPosed manager launch failed",
diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/BaseTest.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/BaseTest.kt
new file mode 100644
index 000000000000..73835f34fb81
--- /dev/null
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/BaseTest.kt
@@ -0,0 +1,26 @@
+package com.topjohnwu.magisk.test
+
+import android.app.Instrumentation
+import android.app.UiAutomation
+import android.content.Context
+import androidx.test.platform.app.InstrumentationRegistry
+import androidx.test.uiautomator.UiDevice
+import com.topjohnwu.magisk.core.utils.RootUtils
+import com.topjohnwu.superuser.Shell
+import org.junit.Assert.assertTrue
+
+interface BaseTest {
+    val instrumentation: Instrumentation
+        get() = InstrumentationRegistry.getInstrumentation()
+    val context: Context get() = instrumentation.targetContext
+    val uiAutomation: UiAutomation get() = instrumentation.uiAutomation
+    val device: UiDevice get() = UiDevice.getInstance(instrumentation)
+
+    companion object {
+        fun prerequisite() {
+            assertTrue("Should have root access", Shell.getShell().isRoot)
+            // Make sure the root service is running
+            RootUtils.Connection.await()
+        }
+    }
+}
diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
index c5597161b663..5112c1bf80c7 100644
--- a/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/Environment.kt
@@ -1,12 +1,10 @@
 package com.topjohnwu.magisk.test
 
 import android.app.Notification
-import android.content.Context
 import android.os.Build
 import androidx.annotation.Keep
 import androidx.core.net.toUri
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.platform.app.InstrumentationRegistry
 import com.topjohnwu.magisk.core.BuildConfig.APP_PACKAGE_NAME
 import com.topjohnwu.magisk.core.di.ServiceLocator
 import com.topjohnwu.magisk.core.download.DownloadNotifier
@@ -15,13 +13,10 @@ import com.topjohnwu.magisk.core.ktx.cachedFile
 import com.topjohnwu.magisk.core.tasks.AppMigration
 import com.topjohnwu.magisk.core.tasks.FlashZip
 import com.topjohnwu.magisk.core.tasks.MagiskInstaller
-import com.topjohnwu.magisk.core.utils.RootUtils
 import com.topjohnwu.superuser.CallbackList
-import com.topjohnwu.superuser.Shell
 import kotlinx.coroutines.runBlocking
 import org.junit.Assert.assertTrue
 import org.junit.Assume.assumeTrue
-import org.junit.Before
 import org.junit.BeforeClass
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -29,16 +24,12 @@ import timber.log.Timber
 
 @Keep
 @RunWith(AndroidJUnit4::class)
-class Environment {
+class Environment : BaseTest {
 
     companion object {
         @BeforeClass
         @JvmStatic
-        fun before() {
-            assertTrue("Should have root access", Shell.getShell().isRoot)
-            // Make sure the root service is running
-            RootUtils.Connection.await()
-        }
+        fun before() = BaseTest.prerequisite()
 
         fun lsposed(): Boolean {
             return Build.VERSION.SDK_INT >= 27 && Build.VERSION.SDK_INT <= 34
@@ -54,13 +45,6 @@ class Environment {
         }
     }
 
-    private lateinit var mContext: Context
-
-    @Before
-    fun setup() {
-        mContext = InstrumentationRegistry.getInstrumentation().targetContext
-    }
-
     @Test
     fun setupMagisk() {
         runBlocking {
@@ -76,11 +60,11 @@ class Environment {
         assumeTrue(lsposed())
 
         val notify = object : DownloadNotifier {
-            override val context = mContext
+            override val context = this@Environment.context
             override fun notifyUpdate(id: Int, editor: (Notification.Builder) -> Unit) {}
         }
         val processor = DownloadProcessor(notify)
-        val zip = mContext.cachedFile("lsposed.zip")
+        val zip = context.cachedFile("lsposed.zip")
         runBlocking {
             ServiceLocator.networkService.fetchFile(LSPOSED_URL).byteStream().use {
                 processor.handleModule(it, zip.toUri())
@@ -98,7 +82,7 @@ class Environment {
             assertTrue(
                 "App hiding failed",
                 AppMigration.patchAndHide(
-                    context = mContext,
+                    context = context,
                     label = "Settings",
                     pkg = "repackaged.$APP_PACKAGE_NAME"
                 )
@@ -111,7 +95,7 @@ class Environment {
         runBlocking {
             assertTrue(
                 "App restoration failed",
-                AppMigration.restoreApp(mContext)
+                AppMigration.restoreApp(context)
             )
         }
     }
diff --git a/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt b/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt
index a06f9c977ca1..085659d1a86e 100644
--- a/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt
+++ b/app/core/src/main/java/com/topjohnwu/magisk/test/MagiskAppTest.kt
@@ -1,12 +1,11 @@
 package com.topjohnwu.magisk.test
 
-import android.app.Instrumentation
 import android.content.Intent
 import android.content.IntentFilter
 import android.os.Build
+import android.os.ParcelFileDescriptor.AutoCloseInputStream
 import androidx.annotation.Keep
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import androidx.test.platform.app.InstrumentationRegistry
 import com.topjohnwu.magisk.core.Config
 import com.topjohnwu.magisk.core.Info
 import com.topjohnwu.magisk.core.di.ServiceLocator
@@ -15,29 +14,19 @@ import kotlinx.coroutines.runBlocking
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotNull
 import org.junit.Assert.assertTrue
-import org.junit.Before
 import org.junit.BeforeClass
 import org.junit.Test
 import org.junit.runner.RunWith
-import java.io.FileInputStream
 import java.util.concurrent.TimeUnit
 
 @Keep
 @RunWith(AndroidJUnit4::class)
-class MagiskAppTest {
+class MagiskAppTest : BaseTest {
 
     companion object {
         @BeforeClass
         @JvmStatic
-        fun before() = Environment.before()
-    }
-
-    private lateinit var inst: Instrumentation
-    private val uiAutomation get() = inst.uiAutomation
-
-    @Before
-    fun setup() {
-        inst = InstrumentationRegistry.getInstrumentation()
+        fun before() = BaseTest.prerequisite()
     }
 
     @Test
@@ -64,7 +53,7 @@ class MagiskAppTest {
 
         val filter = IntentFilter(Intent.ACTION_VIEW)
         filter.addCategory(Intent.CATEGORY_DEFAULT)
-        val monitor = inst.addMonitor(filter, null, false)
+        val monitor = instrumentation.addMonitor(filter, null, false)
 
         // Try to call su from ADB shell
         val cmd = if (Build.VERSION.SDK_INT < 24) {
@@ -80,10 +69,10 @@ class MagiskAppTest {
         assertNotNull("SuRequestActivity is not launched", suRequest)
 
         // Check that the request went through
-        FileInputStream(pfd.fileDescriptor).use {
+        AutoCloseInputStream(pfd).reader().use {
             assertTrue(
                 "Cannot grant root permission from shell",
-                it.reader().readText().contains("uid=0")
+                it.readText().contains("uid=0")
             )
         }
 
diff --git a/scripts/test_common.sh b/scripts/test_common.sh
index c7a5216f2577..2887943b6dcf 100644
--- a/scripts/test_common.sh
+++ b/scripts/test_common.sh
@@ -59,9 +59,10 @@ run_setup() {
 }
 
 run_tests() {
-  local self='com.topjohnwu.magisk.test/com.topjohnwu.magisk.test.TestRunner'
-  local app='com.topjohnwu.magisk.test/com.topjohnwu.magisk.test.AppTestRunner'
-  local stub='repackaged.com.topjohnwu.magisk.test/com.topjohnwu.magisk.test.AppTestRunner'
+  local pkg='com.topjohnwu.magisk.test'
+  local self="$pkg/$pkg.TestRunner"
+  local app="$pkg/$pkg.AppTestRunner"
+  local stub="repackaged.$pkg/$pkg.AppTestRunner"
 
   # Run app tests
   am_instrument '.MagiskAppTest,.AdditionalTest' $app