From 04c015a82cb723157abdb2a52ef572fcf6b52f0e Mon Sep 17 00:00:00 2001
From: Pavel Grafov <pgrafov@google.com>
Date: Tue, 14 Feb 2023 16:35:24 +0000
Subject: [PATCH] Allow overriding "keep profiles running" in tests

Bug: 258425381
Test: atest android.devicepolicy.cts.QuietModeTest
Change-Id: I92f3054c3d33284511b36718284ede3146f3cb30
---
 core/api/test-current.txt                       |  1 +
 .../android/app/admin/DevicePolicyManager.java  | 17 +++++++++++++++++
 .../android/app/admin/IDevicePolicyManager.aidl |  2 ++
 .../DevicePolicyManagerService.java             | 10 +++++++++-
 4 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/core/api/test-current.txt b/core/api/test-current.txt
index bea2bfc097410..135ff01d97881 100644
--- a/core/api/test-current.txt
+++ b/core/api/test-current.txt
@@ -546,6 +546,7 @@ package android.app.admin {
     method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public boolean setDeviceOwnerOnly(@NonNull android.content.ComponentName, int);
     method public void setDeviceOwnerType(@NonNull android.content.ComponentName, int);
     method @RequiresPermission(android.Manifest.permission.MANAGE_DEVICE_ADMINS) public void setNextOperationSafety(int, int);
+    method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public void setOverrideKeepProfilesRunning(boolean);
     method @RequiresPermission(anyOf={android.Manifest.permission.MARK_DEVICE_ORGANIZATION_OWNED, android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS}, conditional=true) public void setProfileOwnerOnOrganizationOwnedDevice(@NonNull android.content.ComponentName, boolean);
     method @RequiresPermission(android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS) public boolean triggerDevicePolicyEngineMigration(boolean);
     field public static final String ACTION_DATA_SHARING_RESTRICTION_APPLIED = "android.app.action.DATA_SHARING_RESTRICTION_APPLIED";
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index ac65a6b5c048f..0abc9d9409785 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -16352,6 +16352,23 @@ public class DevicePolicyManager {
         return null;
     }
 
+    /**
+     * Overrides the effective cached value of enable_keep_profiles_running for testing purposes.
+     *
+     * @hide
+     */
+    @TestApi
+    @RequiresPermission(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS)
+    public void setOverrideKeepProfilesRunning(boolean enabled) {
+        if (mService != null) {
+            try {
+                mService.setOverrideKeepProfilesRunning(enabled);
+            } catch (RemoteException e) {
+                throw e.rethrowFromSystemServer();
+            }
+        }
+    }
+
     /**
      * Triggers the data migration of device policies for existing DPCs to the Device Policy Engine.
      * If {@code forceMigration} is set to {@code true} it skips the prerequisite checks before
diff --git a/core/java/android/app/admin/IDevicePolicyManager.aidl b/core/java/android/app/admin/IDevicePolicyManager.aidl
index 3f66b45b99bf5..0169e304de9bc 100644
--- a/core/java/android/app/admin/IDevicePolicyManager.aidl
+++ b/core/java/android/app/admin/IDevicePolicyManager.aidl
@@ -596,5 +596,7 @@ interface IDevicePolicyManager {
 
     DevicePolicyState getDevicePolicyState();
 
+    void setOverrideKeepProfilesRunning(boolean enabled);
+
     boolean triggerDevicePolicyEngineMigration(boolean forceMigration);
 }
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 8e46edecf4d0f..c6a0d1dca99ed 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -804,7 +804,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
      * This feature flag is checked once after boot and this value us used until the next reboot to
      * avoid needing to handle the flag changing on the fly.
      */
-    private final boolean mKeepProfilesRunning = isKeepProfilesRunningFlagEnabled();
+    private boolean mKeepProfilesRunning = isKeepProfilesRunningFlagEnabled();
 
     /**
      * For apps targeting U+
@@ -20790,6 +20790,14 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
                 DEFAULT_WORK_PROFILE_TELEPHONY_FLAG);
     }
 
+    @Override
+    public void setOverrideKeepProfilesRunning(boolean enabled) {
+        Preconditions.checkCallAuthorization(
+                hasCallingOrSelfPermission(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS));
+        mKeepProfilesRunning = enabled;
+        Slog.i(LOG_TAG, "Keep profiles running overridden to: " + enabled);
+    }
+
     @Override
     public void setMtePolicy(int flags) {
         final Set<Integer> allowedModes =