Merge "Fix to hide Display over other apps permission setting for the Scone application." into main
diff --git a/res/values/strings.xml b/res/values/strings.xml
index a490856..af7473b 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -8076,6 +8076,9 @@
     <!-- Duration in hours and minutes for the length of a Do Not Disturb schedule. For example "1 hr, 22 min" -->
     <string name="zen_mode_schedule_duration"><xliff:g example="10" id="hours">%1$d</xliff:g> hr, <xliff:g example="20" id="minutes">%2$d</xliff:g> min</string>
 
+    <!-- Priority Modes: Label for switch to enable/disable a rule turning on automatically; links to an app-provided configuration page [CHAR LIMIT=40] -->
+    <string name="zen_mode_configuration_link_title">Turn on automatically</string>
+
     <!--  Do not disturb: Title do not disturb settings representing automatic (scheduled) do not disturb rules. [CHAR LIMIT=30] -->
     <string name="zen_mode_schedule_category_title">Schedule</string>
 
@@ -12383,7 +12386,7 @@
     <!-- Toast on failure to reformat data to ext4 -->
     <string name="format_ext4_failure_toast">Failed to reformat and wipe the data partition to ext4.</string>
     <!-- Dialog to OEM unlock the device before using 16K developer option -->
-    <string name="confirm_oem_unlock_for_16k_title">Bootloader Unlock Required for 16KB Mode</string>
+    <string name="confirm_oem_unlock_for_16k_title">Bootloader Unlock Required</string>
     <string name="confirm_oem_unlock_for_16k_text">This device needs to have the bootloader unlocked before using the 16KB developer option.
         Software integrity cannot be guaranteed in this mode, and any data stored on the phone while the bootloader is unlocked may be at risk.
         All user data and settings will be wiped when activating 16KB mode. Once the bootloader is unlocked, activating the 16KB option will require two reboots.
@@ -12412,10 +12415,10 @@
     <!-- persistent notification 16k page agnostic mode text -->
     <string name="page_agnostic_16k_pages_text_short">You are in the 16KB mode of the page-agnostic mode. Software integrity cannot be guaranteed in this mode,
         and any data stored on the phone while the bootloader is unlocked may be at risk. Some features will be disabled in these modes, so some applications may not work.
-        In order to re-enter the production mode, you must, switch back to 4K mode and then lock the bootloader of the device. Tap to read more.</string>
+        In order to re-enter the production mode, you must, switch back to 4KB mode and then lock the bootloader of the device. Tap to read more.</string>
     <string name="page_agnostic_16k_pages_text">You are in the 16KB mode of the page-agnostic mode. Software integrity cannot be guaranteed in this mode,
         and any data stored on the phone while the bootloader is unlocked may be at risk. Some features will be disabled in these modes, so some applications may not work.
-        In order to re-enter the production mode, you must, switch back to 4K mode and then lock the bootloader of the device. This would factory reset the device again and
+        In order to re-enter the production mode, you must, switch back to 4KB mode and then lock the bootloader of the device. This would factory reset the device again and
         restore it to production settings. After the device successfully boots into Android, disable OEM unlocking in Developer options.
         If the device fails to boot into Android or is unstable, re-flash the device with the latest factory images from
         &lt;a href=\"https://developers.google.com/android/images\"&gt;https://developers.google.com/android/images&lt;/a&gt;
diff --git a/src/com/android/settings/development/PageAgnosticNotificationService.java b/src/com/android/settings/development/PageAgnosticNotificationService.java
index bce1dd9..40dff0c 100644
--- a/src/com/android/settings/development/PageAgnosticNotificationService.java
+++ b/src/com/android/settings/development/PageAgnosticNotificationService.java
@@ -96,13 +96,13 @@
                                 notifyPendingIntent)
                         .build();
 
+        // TODO:(b/349860833) Change text style to BigTextStyle once the ellipsis issue is fixed.
         Notification.Builder builder =
                 new Notification.Builder(this, NOTIFICATION_CHANNEL_ID)
                         .setContentTitle(title)
                         .setContentText(text)
                         .setOngoing(true)
                         .setSmallIcon(R.drawable.ic_settings_24dp)
-                        .setStyle(new Notification.BigTextStyle().bigText(text))
                         .setContentIntent(notifyPendingIntent)
                         .addAction(action);
 
diff --git a/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java b/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java
index 98e1a6e..e12cc9f 100644
--- a/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java
+++ b/src/com/android/settings/fuelgauge/PowerUsageFeatureProvider.java
@@ -26,6 +26,7 @@
 import androidx.annotation.NonNull;
 
 import com.android.settings.fuelgauge.batteryusage.BatteryDiffData;
+import com.android.settings.fuelgauge.batteryusage.BatteryEvent;
 import com.android.settings.fuelgauge.batteryusage.DetectRequestSourceType;
 import com.android.settings.fuelgauge.batteryusage.PowerAnomalyEventList;
 import com.android.settingslib.fuelgauge.Estimate;
@@ -43,6 +44,9 @@
     /** Check whether the battery tips card is enabled in the battery usage page */
     boolean isBatteryTipsEnabled();
 
+    /** Check whether overwrite the app optimization mode to restricted mode is enabled */
+    boolean isRestrictedModeOverwriteEnabled();
+
     /** Check whether force expire the app optimization mode. */
     boolean isForceExpireAppOptimizationModeEnabled();
 
@@ -166,5 +170,6 @@
     boolean processBatteryReattributeData(
             @NonNull Context context,
             @NonNull Map<Long, BatteryDiffData> batteryDiffDataMap,
+            @NonNull List<BatteryEvent> batteryEventList,
             final boolean isFromPeriodJob);
 }
diff --git a/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java b/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
index 8ba6374..4f163a2 100644
--- a/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
+++ b/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImpl.java
@@ -18,7 +18,6 @@
 
 import static com.android.settings.Utils.SYSTEMUI_PACKAGE_NAME;
 
-import android.annotation.Nullable;
 import android.content.Context;
 import android.content.Intent;
 import android.content.pm.PackageManager;
@@ -28,9 +27,11 @@
 import android.util.SparseIntArray;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import com.android.internal.util.ArrayUtils;
 import com.android.settings.fuelgauge.batteryusage.BatteryDiffData;
+import com.android.settings.fuelgauge.batteryusage.BatteryEvent;
 import com.android.settings.fuelgauge.batteryusage.DetectRequestSourceType;
 import com.android.settings.fuelgauge.batteryusage.PowerAnomalyEventList;
 import com.android.settingslib.fuelgauge.Estimate;
@@ -88,6 +89,11 @@
     }
 
     @Override
+    public boolean isRestrictedModeOverwriteEnabled() {
+        return false;
+    }
+
+    @Override
     public boolean isForceExpireAppOptimizationModeEnabled() {
         return false;
     }
@@ -254,6 +260,7 @@
     public boolean processBatteryReattributeData(
             @NonNull Context context,
             @NonNull Map<Long, BatteryDiffData> batteryDiffDataMap,
+            @NonNull List<BatteryEvent> batteryEventList,
             final boolean isFromPeriodJob) {
         return false;
     }
diff --git a/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt b/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt
index afff1c9..508c4df 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt
+++ b/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtils.kt
@@ -76,8 +76,7 @@
     fun resetExpiredAppOptModeBeforeTimestamp(context: Context, queryTimestampMs: Long) =
         synchronized(appOptimizationModeLock) {
             val forceExpireEnabled =
-                featureFactory
-                    .powerUsageFeatureProvider.isForceExpireAppOptimizationModeEnabled
+                featureFactory.powerUsageFeatureProvider.isForceExpireAppOptimizationModeEnabled
             val eventsMap = getAppOptModeEventsMap(context)
             val expirationUids = ArrayList<Int>(eventsMap.size)
             for ((uid, event) in eventsMap) {
@@ -113,12 +112,22 @@
         getBatteryOptimizeUtils: (Int, String) -> BatteryOptimizeUtils,
     ) =
         synchronized(appOptimizationModeLock) {
+            val restrictedModeOverwriteEnabled =
+                featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled
             val eventsMap = getAppOptModeEventsMap(context)
             val expirationEvents: MutableMap<Int, AppOptimizationModeEvent> = ArrayMap()
             for (i in uids.indices) {
                 val uid = uids[i]
                 val packageName = packageNames[i]
                 val optimizationMode = optimizationModes[i]
+                if (
+                    !restrictedModeOverwriteEnabled &&
+                        optimizationMode == BatteryOptimizeUtils.MODE_RESTRICTED
+                ) {
+                    // Unable to set restricted mode due to flag protection.
+                    Log.w(TAG, "setOptimizationMode($packageName) into restricted ignored")
+                    continue
+                }
                 val originalOptMode: Int =
                     updateBatteryOptimizationMode(
                         context,
diff --git a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
index 2b88d34..0a1a547 100644
--- a/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
+++ b/src/com/android/settings/fuelgauge/batteryusage/DataProcessManager.java
@@ -522,7 +522,7 @@
                         FeatureFactory.getFeatureFactory()
                                 .getPowerUsageFeatureProvider();
                 featureProvider.processBatteryReattributeData(
-                        mContext, batteryDiffDataMap, mIsFromPeriodJob);
+                        mContext, batteryDiffDataMap, mBatteryEventList, mIsFromPeriodJob);
 
                 Log.d(
                         TAG,
diff --git a/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java b/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java
index 9f819d1..a4bd2aa 100644
--- a/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java
+++ b/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java
@@ -24,6 +24,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
 
 import com.android.settingslib.core.AbstractPreferenceController;
 import com.android.settingslib.notification.modes.ZenMode;
@@ -92,6 +93,16 @@
     }
 
     @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        if (mZenMode != null) {
+            displayPreference(screen, mZenMode);
+        }
+    }
+
+    public void displayPreference(PreferenceScreen screen, @NonNull ZenMode zenMode) {}
+
+    @Override
     public final void updateState(Preference preference) {
         super.updateState(preference);
         if (mZenMode != null) {
diff --git a/src/com/android/settings/notification/modes/ZenModeFragment.java b/src/com/android/settings/notification/modes/ZenModeFragment.java
index 67815b1..3a64fb2 100644
--- a/src/com/android/settings/notification/modes/ZenModeFragment.java
+++ b/src/com/android/settings/notification/modes/ZenModeFragment.java
@@ -35,6 +35,7 @@
 import java.util.List;
 
 public class ZenModeFragment extends ZenModeFragmentBase {
+
     // for mode deletion menu
     private static final int DELETE_MODE = 1;
 
@@ -61,7 +62,8 @@
         prefControllers.add(new ZenModeDisplayLinkPreferenceController(
                 context, "mode_display_settings", mBackend, mHelperBackend));
         prefControllers.add(new ZenModeSetTriggerLinkPreferenceController(context,
-                "zen_automatic_trigger_category", this, mBackend));
+                "zen_automatic_trigger_category", this, mBackend,
+                context.getPackageManager()));
         prefControllers.add(new InterruptionFilterPreferenceController(
                 context, "allow_filtering", mBackend));
         prefControllers.add(new ManualDurationPreferenceController(
diff --git a/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceController.java
index 1c96fee..7328d91 100644
--- a/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceController.java
@@ -18,34 +18,62 @@
 
 import static android.app.AutomaticZenRule.TYPE_SCHEDULE_CALENDAR;
 import static android.app.AutomaticZenRule.TYPE_SCHEDULE_TIME;
+import static android.app.NotificationManager.EXTRA_AUTOMATIC_RULE_ID;
 
+import android.content.ComponentName;
 import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ActivityInfo;
+import android.content.pm.ComponentInfo;
+import android.content.pm.PackageManager;
+import android.content.pm.ResolveInfo;
+import android.service.notification.ConditionProviderService;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import androidx.preference.Preference;
 import androidx.preference.PreferenceCategory;
+import androidx.preference.PreferenceScreen;
 
 import com.android.settings.R;
 import com.android.settings.dashboard.DashboardFragment;
+import com.android.settings.utils.ManagedServiceSettings;
+import com.android.settings.utils.ZenServiceListing;
 import com.android.settingslib.PrimarySwitchPreference;
 import com.android.settingslib.notification.modes.ZenMode;
 import com.android.settingslib.notification.modes.ZenModesBackend;
 
+import java.util.List;
+
 /**
  * Preference controller for the link to an individual mode's configuration page.
  */
 class ZenModeSetTriggerLinkPreferenceController extends AbstractZenModePreferenceController {
+    private static final String TAG = "ZenModeSetTriggerLink";
+
     @VisibleForTesting
     protected static final String AUTOMATIC_TRIGGER_PREF_KEY = "zen_automatic_trigger_settings";
 
+    private static final ManagedServiceSettings.Config CONFIG =
+            ZenModesListFragment.getConditionProviderConfig();
+
+    private ZenServiceListing mServiceListing;
+    private final PackageManager mPm;
     private final DashboardFragment mFragment;
 
     ZenModeSetTriggerLinkPreferenceController(Context context, String key,
-            DashboardFragment fragment,
-            ZenModesBackend backend) {
+            DashboardFragment fragment, ZenModesBackend backend,
+            PackageManager packageManager) {
         super(context, key, backend);
         mFragment = fragment;
+        mPm = packageManager;
+    }
+
+    @VisibleForTesting
+    protected void setServiceListing(ZenServiceListing serviceListing) {
+        mServiceListing = serviceListing;
     }
 
     @Override
@@ -54,6 +82,15 @@
     }
 
     @Override
+    public void displayPreference(PreferenceScreen screen, @NonNull ZenMode zenMode) {
+        if (mServiceListing == null) {
+            mServiceListing = new ZenServiceListing(
+                    mContext, CONFIG, zenMode.getRule().getPackageName());
+        }
+        mServiceListing.reloadApprovedServices();
+    }
+
+    @Override
     public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         // This controller is expected to govern a preference category so that it controls the
         // availability of the entire preference category if the mode doesn't have a way to
@@ -70,29 +107,39 @@
         switchPref.setOnPreferenceClickListener(null);
         switchPref.setIntent(null);
 
-        if (zenMode.isSystemOwned() && zenMode.getType() == TYPE_SCHEDULE_TIME) {
-            switchPref.setTitle(R.string.zen_mode_set_schedule_link);
-            // TODO: b/332937635 - set correct metrics category
-            switchPref.setIntent(ZenSubSettingLauncher.forModeFragment(mContext,
-                    ZenModeSetScheduleFragment.class, zenMode.getId(), 0).toIntent());
-        } else if (zenMode.isSystemOwned() && zenMode.getType() == TYPE_SCHEDULE_CALENDAR) {
-            switchPref.setTitle(R.string.zen_mode_set_calendar_link);
-            switchPref.setIcon(null);
-            // TODO: b/332937635 - set correct metrics category
-            switchPref.setIntent(ZenSubSettingLauncher.forModeFragment(mContext,
-                    ZenModeSetCalendarFragment.class, zenMode.getId(), 0).toIntent());
-        } else if (zenMode.isSystemOwned()) {
-            switchPref.setTitle(R.string.zen_mode_select_schedule);
-            switchPref.setIcon(R.drawable.ic_add_24dp);
-            switchPref.setSummary("");
-            // TODO: b/342156843 - Hide the switch (needs support in SettingsLib).
-            switchPref.setOnPreferenceClickListener(clickedPreference -> {
-                ZenModeScheduleChooserDialog.show(mFragment, mOnScheduleOptionListener);
-                return true;
-            });
+        if (zenMode.isSystemOwned()) {
+            if (zenMode.getType() == TYPE_SCHEDULE_TIME) {
+                switchPref.setTitle(R.string.zen_mode_set_schedule_link);
+                // TODO: b/332937635 - set correct metrics category
+                switchPref.setIntent(ZenSubSettingLauncher.forModeFragment(mContext,
+                        ZenModeSetScheduleFragment.class, zenMode.getId(), 0).toIntent());
+            } else if (zenMode.getType() == TYPE_SCHEDULE_CALENDAR) {
+                switchPref.setTitle(R.string.zen_mode_set_calendar_link);
+                switchPref.setIcon(null);
+                // TODO: b/332937635 - set correct metrics category
+                switchPref.setIntent(ZenSubSettingLauncher.forModeFragment(mContext,
+                        ZenModeSetCalendarFragment.class, zenMode.getId(), 0).toIntent());
+            } else {
+                switchPref.setTitle(R.string.zen_mode_select_schedule);
+                switchPref.setIcon(R.drawable.ic_add_24dp);
+                switchPref.setSummary("");
+                // TODO: b/342156843 - Hide the switch (needs support in SettingsLib).
+                switchPref.setOnPreferenceClickListener(clickedPreference -> {
+                    ZenModeScheduleChooserDialog.show(mFragment, mOnScheduleOptionListener);
+                    return true;
+                });
+            }
         } else {
-            // TODO: b/341961712 - direct preference to app-owned intent if available
-            switchPref.setTitle("not implemented");
+            Intent intent = getAppRuleIntent(zenMode);
+            if (intent != null && isValidIntent(intent)) {
+                preference.setVisible(true);
+                switchPref.setTitle(R.string.zen_mode_configuration_link_title);
+                switchPref.setSummary(zenMode.getRule().getTriggerDescription());
+                switchPref.setIntent(intent);
+            } else {
+                Log.i(TAG, "No intent found for " + zenMode.getRule().getName());
+                preference.setVisible(false);
+            }
         }
     }
 
@@ -114,4 +161,68 @@
         });
         // TODO: b/342156843 - Do we want to jump to the corresponding schedule editing screen?
     };
+
+    @VisibleForTesting
+    protected @Nullable Intent getAppRuleIntent(ZenMode zenMode) {
+        Intent intent = new Intent().addFlags(Intent.FLAG_ACTIVITY_CLEAR_TOP)
+                .putExtra(ConditionProviderService.EXTRA_RULE_ID, zenMode.getId())
+                .putExtra(EXTRA_AUTOMATIC_RULE_ID, zenMode.getId());
+        String owner = zenMode.getRule().getPackageName();
+        ComponentName configActivity = null;
+        if (zenMode.getRule().getConfigurationActivity() != null) {
+            // If a configuration activity is present, use that directly in the intent
+            configActivity = zenMode.getRule().getConfigurationActivity();
+        } else {
+            // Otherwise, look for a condition provider service for the rule's package
+            ComponentInfo ci = mServiceListing.findService(zenMode.getRule().getOwner());
+            if (ci == null) {
+                // do nothing
+            } else if (ci instanceof ActivityInfo) {
+                // new activity backed rule
+                intent.setComponent(new ComponentName(ci.packageName, ci.name));
+                return intent;
+            } else if (ci.metaData != null) {
+                // old service backed rule
+                final String configurationActivity = ci.metaData.getString(
+                        ConditionProviderService.META_DATA_CONFIGURATION_ACTIVITY);
+                if (configurationActivity != null) {
+                    configActivity = ComponentName.unflattenFromString(configurationActivity);
+                }
+            }
+        }
+
+        if (configActivity != null) {
+            // verify that the owner of the rule owns the configuration activity, but only if
+            // owner exists
+            intent.setComponent(configActivity);
+            if (owner == null) {
+                return intent;
+            }
+            try {
+                int ownerUid = mPm.getPackageUid(owner, 0);
+                int configActivityOwnerUid = mPm.getPackageUid(configActivity.getPackageName(), 0);
+                if (ownerUid == configActivityOwnerUid) {
+                    return intent;
+                } else {
+                    Log.w(TAG, "Config activity not in owner package for "
+                            + zenMode.getRule().getName());
+                    return null;
+                }
+            } catch (PackageManager.NameNotFoundException e) {
+                Log.e(TAG, "Failed to find config activity");
+                return null;
+            }
+        }
+        return null;
+    }
+
+    private boolean isValidIntent(Intent intent) {
+        List<ResolveInfo> results = mPm.queryIntentActivities(
+                intent, PackageManager.ResolveInfoFlags.of(0));
+        if (intent.resolveActivity(mPm) == null || results.size() == 0) {
+            Log.w(TAG, "intent for zen rule invalid: " + intent);
+            return false;
+        }
+        return true;
+    }
 }
diff --git a/src/com/android/settings/notification/modes/ZenModesListFragment.java b/src/com/android/settings/notification/modes/ZenModesListFragment.java
index 77107f8..1883945 100644
--- a/src/com/android/settings/notification/modes/ZenModesListFragment.java
+++ b/src/com/android/settings/notification/modes/ZenModesListFragment.java
@@ -78,7 +78,7 @@
         return SettingsEnums.NOTIFICATION_ZEN_MODE_AUTOMATION;
     }
 
-    private static ManagedServiceSettings.Config getConditionProviderConfig() {
+    static ManagedServiceSettings.Config getConditionProviderConfig() {
         return new ManagedServiceSettings.Config.Builder()
                 .setTag(TAG)
                 .setIntentAction(ConditionProviderService.SERVICE_INTERFACE)
diff --git a/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt b/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt
index 05a8f6a..873a2c3 100644
--- a/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt
+++ b/src/com/android/settings/spa/network/NetworkCellularGroupProvider.kt
@@ -279,14 +279,14 @@
 ) {
     CreatePrimarySimListPreference(
         stringResource(id = R.string.primary_sim_calls_title),
-        primarySimInfo.callsAndSmsList,
+        primarySimInfo.callsList,
         callsSelectedId,
         ImageVector.vectorResource(R.drawable.ic_phone),
         actionSetCalls
     )
     CreatePrimarySimListPreference(
         stringResource(id = R.string.primary_sim_texts_title),
-        primarySimInfo.callsAndSmsList,
+        primarySimInfo.smsList,
         textsSelectedId,
         Icons.AutoMirrored.Outlined.Message,
         actionSetTexts
diff --git a/src/com/android/settings/spa/network/PrimarySimRepository.kt b/src/com/android/settings/spa/network/PrimarySimRepository.kt
index e34a29e..420f870 100644
--- a/src/com/android/settings/spa/network/PrimarySimRepository.kt
+++ b/src/com/android/settings/spa/network/PrimarySimRepository.kt
@@ -22,12 +22,14 @@
 import android.util.Log
 import com.android.settings.R
 import com.android.settings.network.SubscriptionUtil
+import com.android.settings.network.telephony.CarrierConfigRepository
 import com.android.settingslib.spa.widget.preference.ListPreferenceOption
 
 class PrimarySimRepository(private val context: Context) {
 
     data class PrimarySimInfo(
-        val callsAndSmsList: List<ListPreferenceOption>,
+        val callsList: List<ListPreferenceOption>,
+        val smsList: List<ListPreferenceOption>,
         val dataList: List<ListPreferenceOption>,
     )
 
@@ -37,7 +39,8 @@
             return null
         }
 
-        val callsAndSmsList = mutableListOf<ListPreferenceOption>()
+        val callsList = mutableListOf<ListPreferenceOption>()
+        val smsList = mutableListOf<ListPreferenceOption>()
         val dataList = mutableListOf<ListPreferenceOption>()
         for (info in selectableSubscriptionInfoList) {
             val item = ListPreferenceOption(
@@ -45,15 +48,22 @@
                 text = "${info.displayName}",
                 summary = SubscriptionUtil.getBidiFormattedPhoneNumber(context, info) ?: "",
             )
-            callsAndSmsList += item
+            callsList += item
+            smsList += item
             dataList += item
         }
-        callsAndSmsList += ListPreferenceOption(
+
+        val askEveryTime = ListPreferenceOption(
             id = SubscriptionManager.INVALID_SUBSCRIPTION_ID,
             text = context.getString(R.string.sim_calls_ask_first_prefs_title),
         )
+        callsList += askEveryTime
+        if (context.resources
+                .getBoolean(com.android.internal.R.bool.config_sms_ask_every_time_support)) {
+            smsList += askEveryTime
+        }
 
-        return PrimarySimInfo(callsAndSmsList, dataList)
+        return PrimarySimInfo(callsList, smsList, dataList)
     }
 
     private companion object {
diff --git a/src/com/android/settings/utils/ZenServiceListing.java b/src/com/android/settings/utils/ZenServiceListing.java
index 99f56f6..96d700b 100644
--- a/src/com/android/settings/utils/ZenServiceListing.java
+++ b/src/com/android/settings/utils/ZenServiceListing.java
@@ -28,11 +28,14 @@
 import android.util.ArraySet;
 import android.util.Slog;
 
+import androidx.annotation.Nullable;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Set;
 
 public class ZenServiceListing {
+    private static final String TAG = "ZenServiceListing";
 
     private final Context mContext;
     private final ManagedServiceSettings.Config mConfig;
@@ -40,9 +43,18 @@
     private final List<Callback> mZenCallbacks = new ArrayList<>();
     private final NotificationManager mNm;
 
+    // only used when android.app.modes_ui flag is true
+    @Nullable
+    private String mPkg = null;
+
     public ZenServiceListing(Context context, ManagedServiceSettings.Config config) {
+        this(context, config, null);
+    }
+
+    public ZenServiceListing(Context context, ManagedServiceSettings.Config config, @Nullable String pkg) {
         mContext = context;
         mConfig = config;
+        mPkg = pkg;
         mNm = (NotificationManager) context.getSystemService(Context.NOTIFICATION_SERVICE);
     }
 
@@ -89,12 +101,16 @@
         }
     }
 
-    private static void getServices(ManagedServiceSettings.Config c, List<ComponentInfo> list,
+    private void getServices(ManagedServiceSettings.Config c, List<ComponentInfo> list,
             PackageManager pm) {
         final int user = ActivityManager.getCurrentUser();
 
+        Intent queryIntent = new Intent(c.intentAction);
+        if (mPkg != null) {
+            queryIntent.setPackage(mPkg);
+        }
         List<ResolveInfo> installedServices = pm.queryIntentServicesAsUser(
-                new Intent(c.intentAction),
+                queryIntent,
                 PackageManager.GET_SERVICES | PackageManager.GET_META_DATA,
                 user);
 
@@ -115,12 +131,16 @@
         }
     }
 
-    private static void getActivities(ManagedServiceSettings.Config c, List<ComponentInfo> list,
+    private void getActivities(ManagedServiceSettings.Config c, List<ComponentInfo> list,
             PackageManager pm) {
         final int user = ActivityManager.getCurrentUser();
 
+        Intent queryIntent = new Intent(c.configIntentAction);
+        if (mPkg != null) {
+            queryIntent.setPackage(mPkg);
+        }
         List<ResolveInfo> resolveInfos = pm.queryIntentActivitiesAsUser(
-                new Intent(c.configIntentAction),
+                queryIntent,
                 PackageManager.GET_ACTIVITIES | PackageManager.GET_META_DATA,
                 user);
 
diff --git a/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java b/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
index cecf8f0..962c40f 100644
--- a/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
+++ b/tests/robotests/src/com/android/settings/fuelgauge/PowerUsageFeatureProviderImplTest.java
@@ -78,6 +78,11 @@
     }
 
     @Test
+    public void isRestrictedModeOverwriteEnabled_returnFalse() {
+        assertThat(mPowerFeatureProvider.isRestrictedModeOverwriteEnabled()).isFalse();
+    }
+
+    @Test
     public void isForceExpireAppOptimizationModeEnabled_returnFalse() {
         assertThat(mPowerFeatureProvider.isForceExpireAppOptimizationModeEnabled()).isFalse();
     }
diff --git a/tests/robotests/src/com/android/settings/notification/modes/TestModeBuilder.java b/tests/robotests/src/com/android/settings/notification/modes/TestModeBuilder.java
index 26c7fe1..6a1f474 100644
--- a/tests/robotests/src/com/android/settings/notification/modes/TestModeBuilder.java
+++ b/tests/robotests/src/com/android/settings/notification/modes/TestModeBuilder.java
@@ -18,6 +18,7 @@
 
 import android.app.AutomaticZenRule;
 import android.app.NotificationManager;
+import android.content.ComponentName;
 import android.net.Uri;
 import android.service.notification.Condition;
 import android.service.notification.ZenDeviceEffects;
@@ -149,6 +150,18 @@
         return this;
     }
 
+    TestModeBuilder setConfigurationActivity(ComponentName configActivity) {
+        mRule.setConfigurationActivity(configActivity);
+        mConfigZenRule.configurationActivity = configActivity;
+        return this;
+    }
+
+    TestModeBuilder setOwner(ComponentName owner) {
+        mRule.setOwner(owner);
+        mConfigZenRule.component = owner;
+        return this;
+    }
+
     ZenMode build() {
         return new ZenMode(mId, mRule, mConfigZenRule);
     }
diff --git a/tests/robotests/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceControllerTest.java
index 31959e5..4ba2146 100644
--- a/tests/robotests/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/modes/ZenModeSetTriggerLinkPreferenceControllerTest.java
@@ -19,8 +19,10 @@
 import static android.app.AutomaticZenRule.TYPE_OTHER;
 import static android.app.AutomaticZenRule.TYPE_SCHEDULE_CALENDAR;
 import static android.app.AutomaticZenRule.TYPE_SCHEDULE_TIME;
+import static android.app.NotificationManager.EXTRA_AUTOMATIC_RULE_ID;
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
 import static android.platform.test.flag.junit.SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT;
+import static android.service.notification.ConditionProviderService.EXTRA_RULE_ID;
 
 import static com.android.settings.notification.modes.ZenModeSetTriggerLinkPreferenceController.AUTOMATIC_TRIGGER_PREF_KEY;
 
@@ -31,10 +33,16 @@
 
 import android.app.AutomaticZenRule;
 import android.app.Flags;
+import android.content.ComponentName;
 import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ComponentInfo;
+import android.content.pm.PackageManager;
 import android.net.Uri;
+import android.os.Bundle;
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
+import android.service.notification.ConditionProviderService;
 import android.service.notification.SystemZenRules;
 import android.service.notification.ZenModeConfig;
 
@@ -44,6 +52,7 @@
 import com.android.settings.R;
 import com.android.settings.SettingsActivity;
 import com.android.settings.dashboard.DashboardFragment;
+import com.android.settings.utils.ZenServiceListing;
 import com.android.settingslib.PrimarySwitchPreference;
 import com.android.settingslib.notification.modes.ZenMode;
 import com.android.settingslib.notification.modes.ZenModesBackend;
@@ -72,6 +81,11 @@
     private PrimarySwitchPreference mPreference;
 
     @Mock
+    private ZenServiceListing mServiceListing;
+    @Mock
+    private PackageManager mPm;
+
+    @Mock
     private PreferenceCategory mPrefCategory;
     @Mock
     private DashboardFragment mFragment;
@@ -84,8 +98,10 @@
         mContext = ApplicationProvider.getApplicationContext();
 
         mPrefController = new ZenModeSetTriggerLinkPreferenceController(mContext,
-                "zen_automatic_trigger_category", mFragment, mBackend);
+                "zen_automatic_trigger_category", mFragment, mBackend, mPm);
+        mPrefController.setServiceListing(mServiceListing);
         mPreference = new PrimarySwitchPreference(mContext);
+
         when(mPrefCategory.findPreference(AUTOMATIC_TRIGGER_PREF_KEY)).thenReturn(mPreference);
     }
 
@@ -93,9 +109,9 @@
     public void testIsAvailable() {
         // should not be available for manual DND
         ZenMode manualMode = ZenMode.manualDndMode(new AutomaticZenRule.Builder("Do Not Disturb",
-                        Uri.parse("manual"))
-                        .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
-                        .build(), true);
+                Uri.parse("manual"))
+                .setInterruptionFilter(INTERRUPTION_FILTER_PRIORITY)
+                .build(), true);
 
         mPrefController.updateZenMode(mPrefCategory, manualMode);
         assertThat(mPrefController.isAvailable()).isFalse();
@@ -164,7 +180,7 @@
     @Test
     public void testRuleLink_schedule() {
         ZenModeConfig.ScheduleInfo scheduleInfo = new ZenModeConfig.ScheduleInfo();
-        scheduleInfo.days = new int[] { Calendar.MONDAY, Calendar.TUESDAY, Calendar.THURSDAY };
+        scheduleInfo.days = new int[]{Calendar.MONDAY, Calendar.TUESDAY, Calendar.THURSDAY};
         scheduleInfo.startHour = 1;
         scheduleInfo.endHour = 15;
         ZenMode mode = new TestModeBuilder()
@@ -237,4 +253,109 @@
         assertThat(updatedMode.getRule().getOwner()).isEqualTo(
                 ZenModeConfig.getScheduleConditionProvider());
     }
+
+    @Test
+    public void testGetAppRuleIntent_configActivity() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setId("id")
+                .setPackage(mContext.getPackageName())
+                .setConfigurationActivity(new ComponentName(mContext.getPackageName(), "test"))
+                .setType(TYPE_OTHER)
+                .setTriggerDescription("some rule")
+                .build();
+
+        when(mPm.getPackageUid(null, 0)).thenReturn(-1);
+        when(mPm.getPackageUid(mContext.getPackageName(), 0)).thenReturn(1);
+
+        Intent res = mPrefController.getAppRuleIntent(mode);
+        assertThat(res).isNotNull();
+        assertThat(res.getStringExtra(EXTRA_RULE_ID)).isEqualTo("id");
+        assertThat(res.getStringExtra(EXTRA_AUTOMATIC_RULE_ID)).isEqualTo("id");
+        assertThat(res.getComponent()).isEqualTo(
+                new ComponentName(mContext.getPackageName(), "test"));
+    }
+
+    @Test
+    public void testGetAppRuleIntent_configActivity_wrongPackage() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setPackage(mContext.getPackageName())
+                .setConfigurationActivity(new ComponentName("another", "test"))
+                .setType(TYPE_OTHER)
+                .build();
+
+        when(mPm.getPackageUid(null, 0)).thenReturn(-1);
+        when(mPm.getPackageUid(mContext.getPackageName(), 0)).thenReturn(1);
+
+        Intent res = mPrefController.getAppRuleIntent(mode);
+        assertThat(res).isNull();
+    }
+
+    @Test
+    public void testGetAppRuleIntent_configActivity_unspecifiedOwner() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setId("id")
+                .setPackage(null)
+                .setConfigurationActivity(new ComponentName("another", "test"))
+                .setType(TYPE_OTHER)
+                .build();
+
+        when(mPm.getPackageUid(null, 0)).thenReturn(-1);
+        when(mPm.getPackageUid(mContext.getPackageName(), 0)).thenReturn(1);
+
+        Intent res = mPrefController.getAppRuleIntent(mode);
+        assertThat(res).isNotNull();
+        assertThat(res.getStringExtra(EXTRA_RULE_ID)).isEqualTo("id");
+        assertThat(res.getStringExtra(EXTRA_AUTOMATIC_RULE_ID)).isEqualTo("id");
+        assertThat(res.getComponent()).isEqualTo(new ComponentName("another", "test"));
+    }
+
+    @Test
+    public void testGetAppRuleIntent_cps() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setId("id")
+                .setPackage(mContext.getPackageName())
+                .setOwner(new ComponentName(mContext.getPackageName(), "service"))
+                .build();
+
+        ComponentInfo ci = new ComponentInfo();
+        ci.packageName = mContext.getPackageName();
+        ci.metaData = new Bundle();
+        ci.metaData.putString(ConditionProviderService.META_DATA_CONFIGURATION_ACTIVITY,
+                ComponentName.flattenToShortString(
+                        new ComponentName(mContext.getPackageName(), "activity")));
+
+        when(mServiceListing.findService(new ComponentName(mContext.getPackageName(), "service")))
+                .thenReturn(ci);
+        when(mPm.getPackageUid(null, 0)).thenReturn(-1);
+        when(mPm.getPackageUid(mContext.getPackageName(), 0)).thenReturn(1);
+
+        Intent res = mPrefController.getAppRuleIntent(mode);
+        assertThat(res).isNotNull();
+        assertThat(res.getStringExtra(EXTRA_RULE_ID)).isEqualTo("id");
+        assertThat(res.getStringExtra(EXTRA_AUTOMATIC_RULE_ID)).isEqualTo("id");
+        assertThat(res.getComponent()).isEqualTo(
+                new ComponentName(mContext.getPackageName(), "activity"));
+    }
+
+    @Test
+    public void testGetAppRuleIntent_cps_wrongPackage() throws Exception {
+        ZenMode mode = new TestModeBuilder()
+                .setPackage("other")
+                .setOwner(new ComponentName(mContext.getPackageName(), "service"))
+                .setType(TYPE_OTHER)
+                .build();
+
+        ComponentInfo ci = new ComponentInfo();
+        ci.packageName = mContext.getPackageName();
+        ci.metaData = new Bundle();
+        ci.metaData.putString(ConditionProviderService.META_DATA_CONFIGURATION_ACTIVITY,
+                ComponentName.flattenToShortString(
+                        new ComponentName(mContext.getPackageName(), "activity")));
+
+        when(mPm.getPackageUid(null, 0)).thenReturn(-1);
+        when(mPm.getPackageUid(mContext.getPackageName(), 0)).thenReturn(1);
+
+        Intent res = mPrefController.getAppRuleIntent(mode);
+        assertThat(res).isNull();
+    }
 }
diff --git a/tests/spa_unit/src/com/android/settings/spa/network/PrimarySimRepositoryTest.kt b/tests/spa_unit/src/com/android/settings/spa/network/PrimarySimRepositoryTest.kt
new file mode 100644
index 0000000..459afa9
--- /dev/null
+++ b/tests/spa_unit/src/com/android/settings/spa/network/PrimarySimRepositoryTest.kt
@@ -0,0 +1,206 @@
+/*
+ * Copyright (C) 2024 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.spa.network
+
+import android.content.Context
+import android.content.res.Resources
+import android.telephony.SubscriptionInfo
+import android.telephony.SubscriptionManager
+import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.settings.R
+import com.android.settingslib.spa.widget.preference.ListPreferenceOption
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.stub
+
+@RunWith(AndroidJUnit4::class)
+class PrimarySimRepositoryTest {
+    @get:Rule
+    val composeTestRule = createComposeRule()
+
+    private val mockSubscriptionManager = mock<SubscriptionManager> {
+        on { addOnSubscriptionsChangedListener(any(), any()) } doAnswer {
+            val listener = it.arguments[1] as SubscriptionManager.OnSubscriptionsChangedListener
+            listener.onSubscriptionsChanged()
+        }
+        on { getPhoneNumber(SUB_ID_1) } doReturn NUMBER_1
+        on { getPhoneNumber(SUB_ID_2) } doReturn NUMBER_2
+    }
+
+    private val context: Context = spy(ApplicationProvider.getApplicationContext()) {
+        on { getSystemService(SubscriptionManager::class.java) } doReturn mockSubscriptionManager
+    }
+    private val spyResources: Resources = spy(context.resources)
+
+    @Test
+    fun getPrimarySimInfo_oneSim_returnNull() {
+        val simList = listOf(
+            SUB_INFO_1,
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNull()
+    }
+
+    @Test
+    fun getPrimarySimInfo_verifyCallsList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+            ListPreferenceOption(
+                id = SubscriptionManager.INVALID_SUBSCRIPTION_ID,
+                text = context.getString(R.string.sim_calls_ask_first_prefs_title),
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.callsList).isEqualTo(expectedList)
+    }
+
+    @Test
+    fun getPrimarySimInfo_verifySmsList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+            ListPreferenceOption(
+                id = SubscriptionManager.INVALID_SUBSCRIPTION_ID,
+                text = context.getString(R.string.sim_calls_ask_first_prefs_title),
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.smsList).isEqualTo(expectedList)
+    }
+
+    @Test
+    fun getPrimarySimInfo_noAskEveryTime_verifySmsList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        context.stub {
+            on { resources } doReturn spyResources
+        }
+        spyResources.stub {
+            on {
+                getBoolean(com.android.internal.R.bool.config_sms_ask_every_time_support)
+            } doReturn false
+        }
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.smsList).isEqualTo(expectedList)
+    }
+
+    @Test
+    fun getPrimarySimInfo_verifyDataList() {
+        val simList = listOf(
+            SUB_INFO_1,
+            SUB_INFO_2
+        )
+        val expectedList = listOf(
+            ListPreferenceOption(
+                id = SUB_INFO_1.subscriptionId,
+                text = "${SUB_INFO_1.displayName}",
+                summary = NUMBER_1
+            ),
+            ListPreferenceOption(
+                id = SUB_INFO_2.subscriptionId,
+                text = "${SUB_INFO_2.displayName}",
+                summary = NUMBER_2
+            ),
+        )
+
+        val primarySimInfo = PrimarySimRepository(context).getPrimarySimInfo(simList)
+
+        assertThat(primarySimInfo).isNotNull()
+        assertThat(primarySimInfo?.dataList).isEqualTo(expectedList)
+    }
+
+    private companion object {
+        const val SUB_ID_1 = 1
+        const val SUB_ID_2 = 2
+        const val DISPLAY_NAME_1 = "Sub 1"
+        const val DISPLAY_NAME_2 = "Sub 2"
+        const val NUMBER_1 = "000000001"
+        const val NUMBER_2 = "000000002"
+        const val MCC = "310"
+
+        val SUB_INFO_1: SubscriptionInfo = SubscriptionInfo.Builder().apply {
+            setId(SUB_ID_1)
+            setDisplayName(DISPLAY_NAME_1)
+            setMcc(MCC)
+        }.build()
+
+        val SUB_INFO_2: SubscriptionInfo = SubscriptionInfo.Builder().apply {
+            setId(SUB_ID_2)
+            setDisplayName(DISPLAY_NAME_2)
+            setMcc(MCC)
+        }.build()
+    }
+}
diff --git a/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt b/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt
index 94f082d..95f69da 100644
--- a/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt
+++ b/tests/unit/src/com/android/settings/fuelgauge/batteryusage/AppOptModeSharedPreferencesUtilsTest.kt
@@ -83,8 +83,17 @@
     }
 
     @Test
-    fun updateAppOptModeExpirationInternal_withExpirationTime_verifyData() {
-        insertAppOptModeEventForTest(expirationTime = 1000L)
+    fun updateAppOptModeExpirationInternal_withoutExpirationTime_verifyEmptyList() {
+        insertAppOptModeEventForTest(expirationTime = UNLIMITED_EXPIRE_TIME)
+
+        assertThat(AppOptModeSharedPreferencesUtils.getAllEvents(context)).isEmpty()
+    }
+
+    @Test
+    fun updateAppOptModeExpirationInternal_setOptimizedModeWithFlagEnabled_verifyData() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(true)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_OPTIMIZED)
 
         val events = AppOptModeSharedPreferencesUtils.getAllEvents(context)
 
@@ -99,8 +108,46 @@
     }
 
     @Test
-    fun updateAppOptModeExpirationInternal_withoutExpirationTime_verifyEmptyList() {
-        insertAppOptModeEventForTest(expirationTime = UNLIMITED_EXPIRE_TIME)
+    fun updateAppOptModeExpirationInternal_setOptimizedModeWithFlagDisabled_verifyData() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(false)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_OPTIMIZED)
+
+        val events = AppOptModeSharedPreferencesUtils.getAllEvents(context)
+
+        assertThat(events).hasSize(1)
+        assertAppOptimizationModeEventInfo(
+            events[0],
+            UID,
+            PACKAGE_NAME,
+            MODE_OPTIMIZED,
+            expirationTime = 1000L
+        )
+    }
+
+    @Test
+    fun updateAppOptModeExpirationInternal_setRestrictedModeWithFlagEnabled_verifyData() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(true)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_RESTRICTED)
+
+        val events = AppOptModeSharedPreferencesUtils.getAllEvents(context)
+
+        assertThat(events).hasSize(1)
+        assertAppOptimizationModeEventInfo(
+            events[0],
+            UID,
+            PACKAGE_NAME,
+            MODE_RESTRICTED,
+            expirationTime = 1000L
+        )
+    }
+
+    @Test
+    fun updateAppOptModeExpirationInternal_setRestrictedModeWithFlagDisabled_verifyEmptyList() {
+        whenever(featureFactory.powerUsageFeatureProvider.isRestrictedModeOverwriteEnabled)
+            .thenReturn(false)
+        insertAppOptModeEventForTest(expirationTime = 1000L, mode = MODE_RESTRICTED)
 
         assertThat(AppOptModeSharedPreferencesUtils.getAllEvents(context)).isEmpty()
     }
@@ -237,14 +284,14 @@
         assertThat(currentOptMode).isEqualTo(MODE_RESTRICTED)
     }
 
-    private fun insertAppOptModeEventForTest(expirationTime: Long) {
+    private fun insertAppOptModeEventForTest(expirationTime: Long, mode: Int = MODE_OPTIMIZED) {
         whenever(testBatteryOptimizeUtils?.isOptimizeModeMutable).thenReturn(true)
-        whenever(testBatteryOptimizeUtils?.getAppOptimizationMode(true)).thenReturn(MODE_OPTIMIZED)
+        whenever(testBatteryOptimizeUtils?.getAppOptimizationMode(true)).thenReturn(mode)
         AppOptModeSharedPreferencesUtils.updateAppOptModeExpirationInternal(
             context,
             mutableListOf(UID),
             mutableListOf(PACKAGE_NAME),
-            mutableListOf(MODE_OPTIMIZED),
+            mutableListOf(mode),
             longArrayOf(expirationTime),
         ) { _: Int, _: String ->
             testBatteryOptimizeUtils