diff --git a/aconfig/settings_security_flag_declarations.aconfig b/aconfig/settings_security_flag_declarations.aconfig
new file mode 100644
index 0000000..42ef4d0
--- /dev/null
+++ b/aconfig/settings_security_flag_declarations.aconfig
@@ -0,0 +1,8 @@
+package: "com.android.settings.flags"
+
+flag {
+  name: "protect_lock_after_timeout_with_auth"
+  namespace: "safety_center"
+  description: "Require an auth challenge to open Lock after timeout page"
+  bug: "315937886"
+}
diff --git a/res/xml/app_ops_permissions_details.xml b/res/xml/app_ops_permissions_details.xml
index eb8188b..3cbe009 100644
--- a/res/xml/app_ops_permissions_details.xml
+++ b/res/xml/app_ops_permissions_details.xml
@@ -17,7 +17,7 @@
 <PreferenceScreen xmlns:android="http://schemas.android.com/apk/res/android"
     xmlns:settings="http://schemas.android.com/apk/res-auto">
 
-    <com.android.settings.widget.FilterTouchesSwitchPreference
+    <com.android.settingslib.RestrictedSwitchPreference
         android:key="app_ops_settings_switch" />
 
     <com.android.settingslib.widget.FooterPreference
diff --git a/res/xml/screen_lock_settings.xml b/res/xml/screen_lock_settings.xml
index 3f732b9..19061d9 100644
--- a/res/xml/screen_lock_settings.xml
+++ b/res/xml/screen_lock_settings.xml
@@ -40,7 +40,7 @@
 
 
     <!-- available in pin/pattern/password -->
-    <com.android.settings.display.TimeoutListPreference
+    <com.android.settings.security.screenlock.ProtectedTimeoutListPreference
         android:key="lock_after_timeout"
         android:title="@string/lock_after_timeout"
         android:summary="@string/summary_placeholder"
diff --git a/src/com/android/settings/applications/UsageAccessDetails.java b/src/com/android/settings/applications/UsageAccessDetails.java
index b81c719..6976149 100644
--- a/src/com/android/settings/applications/UsageAccessDetails.java
+++ b/src/com/android/settings/applications/UsageAccessDetails.java
@@ -38,11 +38,11 @@
 import androidx.preference.Preference;
 import androidx.preference.Preference.OnPreferenceChangeListener;
 import androidx.preference.Preference.OnPreferenceClickListener;
-import androidx.preference.TwoStatePreference;
 
 import com.android.settings.R;
 import com.android.settings.applications.AppStateUsageBridge.UsageState;
 import com.android.settings.overlay.FeatureFactory;
+import com.android.settingslib.RestrictedSwitchPreference;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 
 public class UsageAccessDetails extends AppInfoWithHeader implements OnPreferenceChangeListener,
@@ -57,7 +57,7 @@
     // TODO: Break out this functionality into its own class.
     private AppStateUsageBridge mUsageBridge;
     private AppOpsManager mAppOpsManager;
-    private TwoStatePreference mSwitchPref;
+    private RestrictedSwitchPreference mSwitchPref;
     private Preference mUsageDesc;
     private Intent mSettingsIntent;
     private UsageState mUsageState;
@@ -78,7 +78,7 @@
         mDpm = context.getSystemService(DevicePolicyManager.class);
 
         addPreferencesFromResource(R.xml.app_ops_permissions_details);
-        mSwitchPref = (TwoStatePreference) findPreference(KEY_APP_OPS_SETTINGS_SWITCH);
+        mSwitchPref = (RestrictedSwitchPreference) findPreference(KEY_APP_OPS_SETTINGS_SWITCH);
         mUsageDesc = findPreference(KEY_APP_OPS_SETTINGS_DESC);
 
         getPreferenceScreen().setTitle(R.string.usage_access);
@@ -170,8 +170,16 @@
                 mPackageInfo.applicationInfo.uid);
 
         boolean hasAccess = mUsageState.isPermissible();
+        boolean shouldEnable = mUsageState.permissionDeclared;
+
+        if (shouldEnable && !hasAccess) {
+            mSwitchPref.checkEcmRestrictionAndSetDisabled(AppOpsManager.OPSTR_GET_USAGE_STATS,
+                    mPackageName, mPackageInfo.applicationInfo.uid);
+            shouldEnable = !mSwitchPref.isDisabledByEcm();
+        }
+
         mSwitchPref.setChecked(hasAccess);
-        mSwitchPref.setEnabled(mUsageState.permissionDeclared);
+        mSwitchPref.setEnabled(shouldEnable);
 
         ResolveInfo resolveInfo = mPm.resolveActivityAsUser(mSettingsIntent,
                 PackageManager.GET_META_DATA, mUserId);
diff --git a/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListItem.java b/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListItem.java
index 370a4df..f7dff2e 100644
--- a/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListItem.java
+++ b/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListItem.java
@@ -86,6 +86,14 @@
         return new UserHandle(getUserIdFromDeviceAdminInfo(mInfo));
     }
 
+    public int getUid() {
+        return mInfo.getActivityInfo().applicationInfo.uid;
+    }
+
+    public String getPackageName() {
+        return mInfo.getPackageName();
+    }
+
     public Intent getLaunchIntent(Context context) {
         return new Intent(context, DeviceAdminAdd.class)
                 .putExtra(DevicePolicyManager.EXTRA_DEVICE_ADMIN, mInfo.getComponent());
diff --git a/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListPreferenceController.java b/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListPreferenceController.java
index 1184d8e..55ba8ac 100644
--- a/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListPreferenceController.java
+++ b/src/com/android/settings/applications/specialaccess/deviceadmin/DeviceAdminListPreferenceController.java
@@ -18,6 +18,7 @@
 
 import static android.app.admin.DevicePolicyManager.ACTION_DEVICE_POLICY_MANAGER_STATE_CHANGED;
 
+import android.Manifest;
 import android.app.AppGlobals;
 import android.app.admin.DeviceAdminInfo;
 import android.app.admin.DeviceAdminReceiver;
@@ -45,12 +46,13 @@
 
 import com.android.settings.core.BasePreferenceController;
 import com.android.settings.overlay.FeatureFactory;
+import com.android.settingslib.RestrictedSwitchPreference;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 import com.android.settingslib.core.lifecycle.LifecycleObserver;
 import com.android.settingslib.core.lifecycle.events.OnStart;
 import com.android.settingslib.core.lifecycle.events.OnStop;
-import com.android.settingslib.widget.AppSwitchPreference;
 import com.android.settingslib.widget.FooterPreference;
+import com.android.settingslib.widget.TwoTargetPreference;
 
 import org.xmlpull.v1.XmlPullParserException;
 
@@ -167,35 +169,35 @@
         if (mFooterPreference != null) {
             mFooterPreference.setVisible(mAdmins.isEmpty());
         }
-        final Map<String, AppSwitchPreference> preferenceCache = new ArrayMap<>();
+        final Map<String, RestrictedSwitchPreference> preferenceCache = new ArrayMap<>();
         final Context prefContext = mPreferenceGroup.getContext();
         final int childrenCount = mPreferenceGroup.getPreferenceCount();
         for (int i = 0; i < childrenCount; i++) {
             final Preference pref = mPreferenceGroup.getPreference(i);
-            if (!(pref instanceof AppSwitchPreference)) {
+            if (!(pref instanceof RestrictedSwitchPreference switchPref)) {
                 continue;
             }
-            final AppSwitchPreference appSwitch = (AppSwitchPreference) pref;
-            preferenceCache.put(appSwitch.getKey(), appSwitch);
+            preferenceCache.put(switchPref.getKey(), switchPref);
         }
         for (DeviceAdminListItem item : mAdmins) {
             final String key = item.getKey();
-            AppSwitchPreference pref = preferenceCache.remove(key);
+            RestrictedSwitchPreference pref = preferenceCache.remove(key);
             if (pref == null) {
-                pref = new AppSwitchPreference(prefContext);
+                pref = new RestrictedSwitchPreference(prefContext);
                 mPreferenceGroup.addPreference(pref);
             }
             bindPreference(item, pref);
         }
-        for (AppSwitchPreference unusedCacheItem : preferenceCache.values()) {
+        for (RestrictedSwitchPreference unusedCacheItem : preferenceCache.values()) {
             mPreferenceGroup.removePreference(unusedCacheItem);
         }
     }
 
-    private void bindPreference(DeviceAdminListItem item, AppSwitchPreference pref) {
+    private void bindPreference(DeviceAdminListItem item, RestrictedSwitchPreference pref) {
         pref.setKey(item.getKey());
         pref.setTitle(item.getName());
         pref.setIcon(item.getIcon());
+        pref.setIconSize(TwoTargetPreference.ICON_SIZE_DEFAULT);
         pref.setChecked(item.isActive());
         pref.setSummary(item.getDescription());
         pref.setEnabled(item.isEnabled());
@@ -207,6 +209,8 @@
         });
         pref.setOnPreferenceChangeListener((preference, newValue) -> false);
         pref.setSingleLineTitle(true);
+        pref.checkEcmRestrictionAndSetDisabled(Manifest.permission.BIND_DEVICE_ADMIN,
+                item.getPackageName(), item.getUid());
     }
 
     /**
diff --git a/src/com/android/settings/datetime/TimeZonePreferenceController.java b/src/com/android/settings/datetime/TimeZonePreferenceController.java
index 913640d..d45173f 100644
--- a/src/com/android/settings/datetime/TimeZonePreferenceController.java
+++ b/src/com/android/settings/datetime/TimeZonePreferenceController.java
@@ -26,6 +26,7 @@
 import androidx.preference.Preference;
 
 import com.android.settings.core.BasePreferenceController;
+import com.android.settingslib.RestrictedPreference;
 import com.android.settingslib.datetime.ZoneGetter;
 
 import java.util.Calendar;
@@ -52,6 +53,12 @@
     @Override
     public void updateState(Preference preference) {
         super.updateState(preference);
+
+        if (preference instanceof RestrictedPreference
+                && ((RestrictedPreference) preference).isDisabledByAdmin()) {
+            return;
+        }
+
         preference.setEnabled(shouldEnableManualTimeZoneSelection());
     }
 
diff --git a/src/com/android/settings/security/screenlock/ProtectedTimeoutListPreference.java b/src/com/android/settings/security/screenlock/ProtectedTimeoutListPreference.java
new file mode 100644
index 0000000..26d914c
--- /dev/null
+++ b/src/com/android/settings/security/screenlock/ProtectedTimeoutListPreference.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.security.screenlock;
+
+import android.content.Context;
+import android.util.AttributeSet;
+
+import com.android.settings.display.TimeoutListPreference;
+import com.android.settings.flags.Flags;
+import com.android.settings.wifi.dpp.WifiDppUtils;
+
+/** Wraps {@link TimeoutListPreference} with an authentication challenge for user. */
+public class ProtectedTimeoutListPreference extends TimeoutListPreference {
+    public ProtectedTimeoutListPreference(Context context, AttributeSet attrs) {
+        super(context, attrs);
+    }
+
+    @Override
+    public void performClick() {
+        if (Flags.protectLockAfterTimeoutWithAuth()) {
+            WifiDppUtils.showLockScreen(getContext(), super::performClick);
+        } else {
+            super.performClick();
+        }
+    }
+}
diff --git a/src/com/android/settings/spa/app/appinfo/AppArchiveButton.kt b/src/com/android/settings/spa/app/appinfo/AppArchiveButton.kt
index 0eef9c3..e4fb1ea 100644
--- a/src/com/android/settings/spa/app/appinfo/AppArchiveButton.kt
+++ b/src/com/android/settings/spa/app/appinfo/AppArchiveButton.kt
@@ -54,6 +54,7 @@
     private val packageName = packageInfoPresenter.packageName
     private val userHandle = UserHandle.of(packageInfoPresenter.userId)
     private var broadcastReceiverIsCreated = false
+    private lateinit var appLabel: CharSequence
 
     @Composable
     fun getActionButton(app: ApplicationInfo): ActionButton {
@@ -61,11 +62,12 @@
             val intentFilter = IntentFilter(INTENT_ACTION)
             DisposableBroadcastReceiverAsUser(intentFilter, userHandle) { intent ->
                 if (app.packageName == intent.getStringExtra(PackageInstaller.EXTRA_PACKAGE_NAME)) {
-                    onReceive(intent, app)
+                    onReceive(intent)
                 }
             }
             broadcastReceiverIsCreated = true
         }
+        appLabel = userPackageManager.getApplicationLabel(app)
         return ActionButton(
             text = context.getString(R.string.archive),
             imageVector = Icons.Outlined.CloudUpload,
@@ -113,10 +115,9 @@
         }
     }
 
-    private fun onReceive(intent: Intent, app: ApplicationInfo) {
+    private fun onReceive(intent: Intent) {
         when (val status = intent.getIntExtra(PackageInstaller.EXTRA_STATUS, Int.MIN_VALUE)) {
             PackageInstaller.STATUS_SUCCESS -> {
-                val appLabel = userPackageManager.getApplicationLabel(app)
                 Toast.makeText(
                     context,
                     context.getString(R.string.archiving_succeeded, appLabel),
diff --git a/src/com/android/settings/spa/app/appinfo/AppForceStopButton.kt b/src/com/android/settings/spa/app/appinfo/AppForceStopButton.kt
index 345d931..c3183a7 100644
--- a/src/com/android/settings/spa/app/appinfo/AppForceStopButton.kt
+++ b/src/com/android/settings/spa/app/appinfo/AppForceStopButton.kt
@@ -19,6 +19,7 @@
 import android.app.settings.SettingsEnums
 import android.content.pm.ApplicationInfo
 import android.os.UserManager
+import androidx.annotation.VisibleForTesting
 import androidx.compose.material.icons.Icons
 import androidx.compose.material.icons.outlined.Report
 import androidx.compose.material3.Text
@@ -87,9 +88,10 @@
         dialogPresenter.open()
     }
 
-    private fun getAdminRestriction(app: ApplicationInfo): EnforcedAdmin? = when {
+    @VisibleForTesting
+    fun getAdminRestriction(app: ApplicationInfo): EnforcedAdmin? = when {
         packageManager.isPackageStateProtected(app.packageName, app.userId) -> {
-            RestrictedLockUtilsInternal.getDeviceOwner(context)
+            RestrictedLockUtilsInternal.getDeviceOwner(context) ?: EnforcedAdmin()
         }
 
         else -> RestrictedLockUtilsInternal.checkIfRestrictionEnforced(
diff --git a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowContentResolver.java b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowContentResolver.java
index 5d9202f..e259cad 100644
--- a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowContentResolver.java
+++ b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowContentResolver.java
@@ -16,16 +16,10 @@
 
 package com.android.settings.testutils.shadow;
 
-import static android.provider.SearchIndexablesContract.INDEXABLES_RAW_COLUMNS;
-
 import android.accounts.Account;
 import android.annotation.UserIdInt;
 import android.content.ContentResolver;
 import android.content.SyncAdapterType;
-import android.database.Cursor;
-import android.database.MatrixCursor;
-import android.net.Uri;
-import android.provider.SearchIndexablesContract;
 import android.text.TextUtils;
 
 import org.robolectric.annotation.Implementation;
@@ -35,7 +29,7 @@
 import java.util.Map;
 
 @Implements(ContentResolver.class)
-public class ShadowContentResolver {
+public class ShadowContentResolver extends org.robolectric.shadows.ShadowContentResolver {
 
     private static SyncAdapterType[] sSyncAdapterTypes = new SyncAdapterType[0];
     private static Map<String, Integer> sSyncable = new HashMap<>();
@@ -48,23 +42,14 @@
     }
 
     @Implementation
-    protected final Cursor query(Uri uri, String[] projection, String selection,
-            String[] selectionArgs, String sortOrder) {
-        MatrixCursor cursor = new MatrixCursor(INDEXABLES_RAW_COLUMNS);
-        MatrixCursor.RowBuilder builder = cursor.newRow()
-                .add(SearchIndexablesContract.NonIndexableKey.COLUMN_KEY_VALUE, "");
-        return cursor;
-    }
-
-    @Implementation
     protected static int getIsSyncableAsUser(Account account, String authority, int userId) {
-        return sSyncable.containsKey(authority) ? sSyncable.get(authority) : 1;
+        return sSyncable.getOrDefault(authority, 1);
     }
 
     @Implementation
     protected static boolean getSyncAutomaticallyAsUser(Account account, String authority,
             int userId) {
-        return sSyncAutomatically.containsKey(authority) ? sSyncAutomatically.get(authority) : true;
+        return sSyncAutomatically.getOrDefault(authority, true);
     }
 
     @Implementation
diff --git a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowRestrictedLockUtilsInternal.java b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowRestrictedLockUtilsInternal.java
index d66aa61..5989d49 100644
--- a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowRestrictedLockUtilsInternal.java
+++ b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowRestrictedLockUtilsInternal.java
@@ -50,7 +50,7 @@
     }
 
     @Implementation
-    protected static EnforcedAdmin checkIfMeteredDataRestricted(Context context,
+    protected static EnforcedAdmin checkIfMeteredDataUsageUserControlDisabled(Context context,
             String packageName, int userId) {
         if (sIsRestricted) {
             return new EnforcedAdmin();
diff --git a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppArchiveButtonTest.kt b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppArchiveButtonTest.kt
index df1f153..6b4cc0d 100644
--- a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppArchiveButtonTest.kt
+++ b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppArchiveButtonTest.kt
@@ -63,8 +63,10 @@
         whenever(packageInfoPresenter.context).thenReturn(context)
         whenever(packageInfoPresenter.userPackageManager).thenReturn(userPackageManager)
         whenever(userPackageManager.packageInstaller).thenReturn(packageInstaller)
+        whenever(userPackageManager.getApplicationLabel(any())).thenReturn(APP_LABEL)
         whenever(packageInfoPresenter.packageName).thenReturn(PACKAGE_NAME)
-        appArchiveButton = AppArchiveButton(packageInfoPresenter, isHibernationSwitchEnabledStateFlow)
+        appArchiveButton =
+            AppArchiveButton(packageInfoPresenter, isHibernationSwitchEnabledStateFlow)
     }
 
     @Test
@@ -150,5 +152,6 @@
 
     private companion object {
         const val PACKAGE_NAME = "package.name"
+        const val APP_LABEL = "App label"
     }
 }
diff --git a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppButtonsTest.kt b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppButtonsTest.kt
index 733e1a4..c742bd7 100644
--- a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppButtonsTest.kt
+++ b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppButtonsTest.kt
@@ -45,6 +45,7 @@
 import org.mockito.Mock
 import org.mockito.MockitoSession
 import org.mockito.Spy
+import org.mockito.kotlin.any
 import org.mockito.quality.Strictness
 import org.mockito.Mockito.`when` as whenever
 
@@ -80,6 +81,7 @@
         whenever(packageInfoPresenter.context).thenReturn(context)
         whenever(packageInfoPresenter.packageName).thenReturn(PACKAGE_NAME)
         whenever(packageInfoPresenter.userPackageManager).thenReturn(packageManager)
+        whenever(packageManager.getApplicationLabel(any())).thenReturn(APP_LABEL)
         whenever(packageManager.packageInstaller).thenReturn(packageInstaller)
         whenever(packageManager.getPackageInfo(PACKAGE_NAME, 0)).thenReturn(PACKAGE_INFO)
         whenever(AppUtils.isMainlineModule(packageManager, PACKAGE_NAME)).thenReturn(false)
@@ -113,7 +115,8 @@
         featureFlags.setFlag(Flags.FLAG_ARCHIVING, false)
         setContent()
 
-        composeTestRule.onNodeWithText(context.getString(R.string.launch_instant_app)).assertIsDisplayed()
+        composeTestRule.onNodeWithText(context.getString(R.string.launch_instant_app))
+            .assertIsDisplayed()
     }
 
     @Test
@@ -122,7 +125,8 @@
         featureFlags.setFlag(Flags.FLAG_ARCHIVING, true)
         setContent()
 
-        composeTestRule.onNodeWithText(context.getString(R.string.launch_instant_app)).assertIsNotDisplayed()
+        composeTestRule.onNodeWithText(context.getString(R.string.launch_instant_app))
+            .assertIsNotDisplayed()
     }
 
     @Test
@@ -184,6 +188,7 @@
 
     private companion object {
         const val PACKAGE_NAME = "package.name"
+        const val APP_LABEL = "App label"
         val PACKAGE_INFO = PackageInfo().apply {
             applicationInfo = ApplicationInfo().apply {
                 packageName = PACKAGE_NAME
diff --git a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppForceStopButtonTest.kt b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppForceStopButtonTest.kt
index c093863..84d6651 100644
--- a/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppForceStopButtonTest.kt
+++ b/tests/spa_unit/src/com/android/settings/spa/app/appinfo/AppForceStopButtonTest.kt
@@ -17,79 +17,82 @@
 package com.android.settings.spa.app.appinfo
 
 import android.app.admin.DevicePolicyManager
+import android.content.ComponentName
 import android.content.Context
 import android.content.pm.ApplicationInfo
 import android.content.pm.PackageManager
+import android.os.UserHandle
+import android.os.UserManager
+import androidx.compose.material3.Button
+import androidx.compose.material3.Text
+import androidx.compose.ui.test.assertIsEnabled
+import androidx.compose.ui.test.assertIsNotEnabled
 import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.onNodeWithText
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.settingslib.spa.widget.button.ActionButton
+import com.android.settings.R
+import com.android.settingslib.spa.testutils.delay
 import com.android.settingslib.spaprivileged.framework.common.devicePolicyManager
 import com.android.settingslib.spaprivileged.model.app.userId
 import com.google.common.truth.Truth.assertThat
-import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Spy
-import org.mockito.junit.MockitoJUnit
-import org.mockito.junit.MockitoRule
-import org.mockito.Mockito.`when` as whenever
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
 
 @RunWith(AndroidJUnit4::class)
 class AppForceStopButtonTest {
     @get:Rule
     val composeTestRule = createComposeRule()
 
-    @get:Rule
-    val mockito: MockitoRule = MockitoJUnit.rule()
+    private val mockPackageManager = mock<PackageManager>()
 
-    @Spy
-    private val context: Context = ApplicationProvider.getApplicationContext()
+    private val mockDevicePolicyManager = mock<DevicePolicyManager>()
 
-    @Mock
-    private lateinit var packageInfoPresenter: PackageInfoPresenter
-
-    @Mock
-    private lateinit var packageManager: PackageManager
-
-    @Mock
-    private lateinit var devicePolicyManager: DevicePolicyManager
-
-    private lateinit var appForceStopButton: AppForceStopButton
-
-    @Before
-    fun setUp() {
-        whenever(packageInfoPresenter.context).thenReturn(context)
-        whenever(context.packageManager).thenReturn(packageManager)
-        whenever(context.devicePolicyManager).thenReturn(devicePolicyManager)
-        appForceStopButton = AppForceStopButton(packageInfoPresenter)
+    private val mockUserManager = mock<UserManager> {
+        on { getUserRestrictionSources(any(), any()) } doReturn emptyList()
     }
 
-    @Test
-    fun getActionButton() {
+    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {
+        on { packageManager } doReturn mockPackageManager
+        on { devicePolicyManager } doReturn mockDevicePolicyManager
+        on { getSystemService(Context.DEVICE_POLICY_SERVICE) } doReturn mockDevicePolicyManager
+        on { getSystemService(Context.USER_SERVICE) } doReturn mockUserManager
     }
 
+    private val packageInfoPresenter = mock<PackageInfoPresenter> {
+        on { context } doReturn context
+    }
+
+    private val appForceStopButton = AppForceStopButton(packageInfoPresenter)
+
     @Test
     fun getActionButton_isActiveAdmin_buttonDisabled() {
         val app = createApp()
-        whenever(devicePolicyManager.packageHasActiveAdmins(PACKAGE_NAME, app.userId))
-            .thenReturn(true)
+        mockDevicePolicyManager.stub {
+            on { packageHasActiveAdmins(PACKAGE_NAME, app.userId) } doReturn true
+        }
 
-        val actionButton = setForceStopButton(app)
+        setForceStopButton(app)
 
-        assertThat(actionButton.enabled).isFalse()
+        composeTestRule.onNodeWithText(context.getString(R.string.force_stop)).assertIsNotEnabled()
     }
 
     @Test
     fun getActionButton_isUninstallInQueue_buttonDisabled() {
         val app = createApp()
-        whenever(devicePolicyManager.isUninstallInQueue(PACKAGE_NAME)).thenReturn(true)
+        mockDevicePolicyManager.stub {
+            on { isUninstallInQueue(PACKAGE_NAME) } doReturn true
+        }
 
-        val actionButton = setForceStopButton(app)
+        setForceStopButton(app)
 
-        assertThat(actionButton.enabled).isFalse()
+        composeTestRule.onNodeWithText(context.getString(R.string.force_stop)).assertIsNotEnabled()
     }
 
     @Test
@@ -98,35 +101,79 @@
             flags = ApplicationInfo.FLAG_STOPPED
         }
 
-        val actionButton = setForceStopButton(app)
+        setForceStopButton(app)
 
-        assertThat(actionButton.enabled).isFalse()
+        composeTestRule.onNodeWithText(context.getString(R.string.force_stop)).assertIsNotEnabled()
     }
 
     @Test
     fun getActionButton_regularApp_buttonEnabled() {
         val app = createApp()
 
-        val actionButton = setForceStopButton(app)
+        setForceStopButton(app)
 
-        assertThat(actionButton.enabled).isTrue()
+        composeTestRule.onNodeWithText(context.getString(R.string.force_stop)).assertIsEnabled()
     }
 
-    private fun setForceStopButton(app: ApplicationInfo): ActionButton {
-        lateinit var actionButton: ActionButton
-        composeTestRule.setContent {
-            actionButton = appForceStopButton.getActionButton(app)
+    @Test
+    fun getAdminRestriction_packageNotProtected() {
+        mockPackageManager.stub {
+            on { isPackageStateProtected(PACKAGE_NAME, UserHandle.getUserId(UID)) } doReturn false
         }
-        return actionButton
+
+        val admin = appForceStopButton.getAdminRestriction(createApp())
+
+        assertThat(admin).isNull()
+    }
+
+    @Test
+    fun getAdminRestriction_packageProtectedAndHaveOwner() {
+        mockPackageManager.stub {
+            on { isPackageStateProtected(PACKAGE_NAME, UserHandle.getUserId(UID)) } doReturn true
+        }
+        mockDevicePolicyManager.stub {
+            on { deviceOwnerComponentOnAnyUser } doReturn DEVICE_OWNER
+        }
+
+        val admin = appForceStopButton.getAdminRestriction(createApp())!!
+
+        assertThat(admin.component).isEqualTo(DEVICE_OWNER)
+    }
+
+    @Test
+    fun getAdminRestriction_packageProtectedAndNotHaveOwner() {
+        mockPackageManager.stub {
+            on { isPackageStateProtected(PACKAGE_NAME, UserHandle.getUserId(UID)) } doReturn true
+        }
+        mockDevicePolicyManager.stub {
+            on { deviceOwnerComponentOnAnyUser } doReturn null
+        }
+
+        val admin = appForceStopButton.getAdminRestriction(createApp())!!
+
+        assertThat(admin.component).isNull()
+    }
+
+    private fun setForceStopButton(app: ApplicationInfo) {
+        composeTestRule.setContent {
+            val actionButton = appForceStopButton.getActionButton(app)
+            Button(onClick = {}, enabled = actionButton.enabled) {
+                Text(actionButton.text)
+            }
+        }
+        composeTestRule.delay()
     }
 
     private fun createApp(builder: ApplicationInfo.() -> Unit = {}) =
         ApplicationInfo().apply {
             packageName = PACKAGE_NAME
+            uid = UID
             enabled = true
         }.apply(builder)
 
     private companion object {
         const val PACKAGE_NAME = "package.name"
+        const val UID = 10000
+        val DEVICE_OWNER = ComponentName("device", "Owner")
     }
-}
\ No newline at end of file
+}