1 files changed, 15 insertions, 5 deletions
diff --git a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/UserControlDisabledPackagesTest.java b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/UserControlDisabledPackagesTest.java
index 4da5fff7f0b..9c8a81d05f7 100644
--- a/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/UserControlDisabledPackagesTest.java
+++ b/hostsidetests/devicepolicy/app/DeviceOwner/src/com/android/cts/deviceowner/UserControlDisabledPackagesTest.java
@@ -25,6 +25,8 @@ import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.util.Log;
 
+import androidx.test.InstrumentationRegistry;
+
 import java.util.ArrayList;
 
 /**
@@ -41,6 +43,7 @@ public class UserControlDisabledPackagesTest extends BaseDeviceOwnerTest {
     private static final String SIMPLE_APP_PKG = "com.android.cts.launcherapps.simpleapp";
     private static final String SIMPLE_APP_ACTIVITY =
             "com.android.cts.launcherapps.simpleapp.SimpleActivityImmediateExit";
+    private static final String ARG_PID_BEFORE_STOP = "pidOfSimpleapp";
 
     public void testSetUserControlDisabledPackages() throws Exception {
         ArrayList<String> protectedPackages = new ArrayList<>();
@@ -86,14 +89,15 @@ public class UserControlDisabledPackagesTest extends BaseDeviceOwnerTest {
         // Check if package is part of UserControlDisabledPackages before checking if
         // package is stopped since it is a necessary condition to prevent stopping of
         // package
-
         assertThat(mDevicePolicyManager.getUserControlDisabledPackages(getWho()))
                 .containsExactly(SIMPLE_APP_PKG);
-        assertPackageRunningState(/* running= */ true);
+        assertPackageRunningState(/* running= */ true,
+                InstrumentationRegistry.getArguments().getString(ARG_PID_BEFORE_STOP, "-1"));
     }
 
     public void testFgsStopWithUserControlEnabled() throws Exception {
-        assertPackageRunningState(/* running= */ false);
+        assertPackageRunningState(/* running= */ false,
+                InstrumentationRegistry.getArguments().getString(ARG_PID_BEFORE_STOP, "-1"));
         assertThat(mDevicePolicyManager.getUserControlDisabledPackages(getWho())).isEmpty();
     }
 
@@ -120,9 +124,15 @@ public class UserControlDisabledPackagesTest extends BaseDeviceOwnerTest {
         return pid.length() > 0;
     }
 
-    private void assertPackageRunningState(boolean shouldBeRunning) throws Exception {
+    private void assertPackageRunningState(boolean shouldBeRunning, String argPid)
+            throws Exception {
+        String pid = executeShellCommand(String.format("pidof %s", SIMPLE_APP_PKG)).trim();
+
+        final boolean samePid = pid.equals(argPid);
+        final boolean stillRunning = samePid && isPackageRunning(SIMPLE_APP_PKG);
+
         assertWithMessage("Package %s running for user %s", SIMPLE_APP_PKG,
                 getCurrentUser().getIdentifier())
-                .that(isPackageRunning(SIMPLE_APP_PKG)).isEqualTo(shouldBeRunning);
+                .that(stillRunning).isEqualTo(shouldBeRunning);
     }
 }