Merge "Refactor preference controllers to receive and mutate ZenMode/ZenPolicy" into main
diff --git a/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java b/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java
index b5e5dc7..587f640 100644
--- a/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java
+++ b/src/com/android/settings/notification/modes/AbstractZenModePreferenceController.java
@@ -16,9 +16,10 @@
 
 package com.android.settings.notification.modes;
 
-import android.app.AutomaticZenRule;
 import android.app.Flags;
 import android.content.Context;
+import android.service.notification.ZenPolicy;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -26,11 +27,17 @@
 
 import com.android.settingslib.core.AbstractPreferenceController;
 
+import com.google.common.base.Preconditions;
+
+import java.util.function.Function;
+
 /**
  * Base class for any preference controllers pertaining to any single Zen mode.
  */
 abstract class AbstractZenModePreferenceController extends AbstractPreferenceController {
 
+    private static final String TAG = "AbstractZenModePreferenceController";
+
     @Nullable
     protected ZenModesBackend mBackend;
 
@@ -38,7 +45,7 @@
     private ZenMode mZenMode;
 
     @NonNull
-    final String mKey;
+    private final String mKey;
 
     // ZenModesBackend should only be passed in if the preference controller may set the user's
     // policy for this zen mode. Otherwise, if the preference controller is essentially read-only
@@ -67,20 +74,56 @@
         updateState(preference);
     }
 
-    @Nullable
-    public ZenMode getMode() {
-        return mZenMode;
+    @Override
+    public final void updateState(Preference preference) {
+        super.updateState(preference);
+        if (mZenMode != null) {
+            updateState(preference, mZenMode);
+        }
     }
 
-    @Nullable
-    public AutomaticZenRule getAZR() {
-        if (mZenMode == null || mZenMode.getRule() == null) {
+    abstract void updateState(Preference preference, @NonNull ZenMode zenMode);
+
+    @Override
+    public final CharSequence getSummary() {
+        if (mZenMode != null) {
+            return getSummary(mZenMode);
+        } else {
             return null;
         }
-        return mZenMode.getRule();
     }
 
-    /** Implementations of this class should override
-     *  {@link AbstractPreferenceController#updateState(Preference)} to specify what should
-     *  happen when the preference is updated */
+    @Nullable
+    protected CharSequence getSummary(@NonNull ZenMode zenMode) {
+        return null;
+    }
+
+    /**
+     * Subclasses should call this method (or a more specific one, like {@link #savePolicy} from
+     * their {@code onPreferenceChange()} or similar, in order to apply changes to the mode being
+     * edited (e.g. {@code saveMode(mode -> { mode.setX(value); return mode; } }.
+     *
+     * @param updater Function to update the {@link ZenMode}. Modifying and returning the same
+     *                instance is ok.
+     */
+    protected final boolean saveMode(Function<ZenMode, ZenMode> updater) {
+        Preconditions.checkState(mBackend != null);
+        ZenMode mode = mZenMode;
+        if (mode == null) {
+            Log.wtf(TAG, "Cannot save mode, it hasn't been loaded (" + getClass() + ")");
+            return false;
+        }
+        mode = updater.apply(mode);
+        mBackend.updateMode(mode);
+        return true;
+    }
+
+    protected final boolean savePolicy(Function<ZenPolicy.Builder, ZenPolicy.Builder> updater) {
+        return saveMode(mode -> {
+            ZenPolicy.Builder policyBuilder = new ZenPolicy.Builder(mode.getPolicy());
+            policyBuilder = updater.apply(policyBuilder);
+            mode.setPolicy(policyBuilder.build());
+            return mode;
+        });
+    }
 }
diff --git a/src/com/android/settings/notification/modes/ZenMode.java b/src/com/android/settings/notification/modes/ZenMode.java
index bf10cd4..058799b 100644
--- a/src/com/android/settings/notification/modes/ZenMode.java
+++ b/src/com/android/settings/notification/modes/ZenMode.java
@@ -26,6 +26,7 @@
 import android.app.NotificationManager;
 import android.content.Context;
 import android.graphics.drawable.Drawable;
+import android.service.notification.ZenDeviceEffects;
 import android.service.notification.ZenPolicy;
 import android.util.Log;
 
@@ -151,14 +152,6 @@
     }
 
     /**
-     * Use sparingly. If you're updating a policy field, use
-     * {@link #setPolicy(android.service.notification.ZenPolicy)} instead.
-     */
-    public void setAzr(@NonNull AutomaticZenRule newRule) {
-        mRule = newRule;
-    }
-
-    /**
      * Updates the {@link ZenPolicy} of the associated {@link AutomaticZenRule} based on the
      * supplied policy. In some cases this involves conversions, so that the following call
      * to {@link #getPolicy} might return a different policy from the one supplied here.
@@ -204,6 +197,13 @@
         mRule.setZenPolicy(policy);
     }
 
+    @NonNull
+    public ZenDeviceEffects getDeviceEffects() {
+        return mRule.getDeviceEffects() != null
+                ? mRule.getDeviceEffects()
+                : new ZenDeviceEffects.Builder().build();
+    }
+
     public boolean canBeDeleted() {
         return !mIsManualDnd;
     }
diff --git a/src/com/android/settings/notification/modes/ZenModeCallsLinkPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeCallsLinkPreferenceController.java
index 1d1d750..746af44 100644
--- a/src/com/android/settings/notification/modes/ZenModeCallsLinkPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeCallsLinkPreferenceController.java
@@ -20,12 +20,15 @@
 
 import android.content.Context;
 import android.os.Bundle;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
+
 import com.android.settings.core.SubSettingLauncher;
 
-public class ZenModeCallsLinkPreferenceController extends AbstractZenModePreferenceController  {
+class ZenModeCallsLinkPreferenceController extends AbstractZenModePreferenceController  {
 
-    private ZenModeSummaryHelper mSummaryHelper;
+    private final ZenModeSummaryHelper mSummaryHelper;
 
     public ZenModeCallsLinkPreferenceController(Context context, String key,
             ZenModesBackend backend) {
@@ -34,16 +37,15 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         Bundle bundle = new Bundle();
-        bundle.putString(MODE_ID, getMode().getId());
+        bundle.putString(MODE_ID, zenMode.getId());
         // TODO(b/332937635): Update metrics category
         preference.setIntent(new SubSettingLauncher(mContext)
                 .setDestination(ZenModeCallsFragment.class.getName())
                 .setSourceMetricsCategory(0)
                 .setArguments(bundle)
                 .toIntent());
-        preference.setSummary(mSummaryHelper.getCallsSettingSummary(getMode()));
+        preference.setSummary(mSummaryHelper.getCallsSettingSummary(zenMode));
     }
 }
\ No newline at end of file
diff --git a/src/com/android/settings/notification/modes/ZenModeDisplayEffectPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeDisplayEffectPreferenceController.java
index f69d0d6..bca7b55 100644
--- a/src/com/android/settings/notification/modes/ZenModeDisplayEffectPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeDisplayEffectPreferenceController.java
@@ -16,16 +16,10 @@
 
 package com.android.settings.notification.modes;
 
-import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_ALARMS;
-import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_EVENTS;
-import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_MEDIA;
-import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_REMINDERS;
-import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_SYSTEM;
-
-import android.app.AutomaticZenRule;
 import android.content.Context;
 import android.service.notification.ZenDeviceEffects;
-import android.service.notification.ZenPolicy;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
 import androidx.preference.TwoStatePreference;
 
@@ -38,9 +32,9 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         TwoStatePreference pref = (TwoStatePreference) preference;
-        ZenDeviceEffects effects =  getMode().getRule().getDeviceEffects();
+        ZenDeviceEffects effects =  zenMode.getRule().getDeviceEffects();
         if (effects == null) {
             pref.setChecked(false);
         } else {
@@ -62,33 +56,27 @@
     }
 
     @Override
-    public boolean onPreferenceChange(Preference preference, Object newValue) {
+    public boolean onPreferenceChange(@NonNull Preference preference, Object newValue) {
         final boolean allow = (Boolean) newValue;
-
-        ZenDeviceEffects currEffects =  getMode().getRule().getDeviceEffects();
-        ZenDeviceEffects.Builder updatedEffects = currEffects == null
-                ? new ZenDeviceEffects.Builder()
-                : new ZenDeviceEffects.Builder(getMode().getRule().getDeviceEffects());
-        switch (getPreferenceKey()) {
-            case "effect_greyscale":
-                updatedEffects.setShouldDisplayGrayscale(allow);
-                break;
-            case "effect_aod":
-                updatedEffects.setShouldSuppressAmbientDisplay(allow);
-                break;
-            case "effect_wallpaper":
-                updatedEffects.setShouldDimWallpaper(allow);
-                break;
-            case "effect_dark_theme":
-                updatedEffects.setShouldUseNightMode(allow);
-                break;
-        }
-        AutomaticZenRule updatedAzr = new AutomaticZenRule.Builder(getMode().getRule())
-                .setDeviceEffects(updatedEffects.build())
-                .build();
-        getMode().setAzr(updatedAzr);
-        mBackend.updateMode(getMode());
-
-        return true;
+        return saveMode(zenMode -> {
+            ZenDeviceEffects.Builder updatedEffects = new ZenDeviceEffects.Builder(
+                    zenMode.getDeviceEffects());
+            switch (getPreferenceKey()) {
+                case "effect_greyscale":
+                    updatedEffects.setShouldDisplayGrayscale(allow);
+                    break;
+                case "effect_aod":
+                    updatedEffects.setShouldSuppressAmbientDisplay(allow);
+                    break;
+                case "effect_wallpaper":
+                    updatedEffects.setShouldDimWallpaper(allow);
+                    break;
+                case "effect_dark_theme":
+                    updatedEffects.setShouldUseNightMode(allow);
+                    break;
+            }
+            zenMode.getRule().setDeviceEffects(updatedEffects.build());
+            return zenMode;
+        });
     }
 }
diff --git a/src/com/android/settings/notification/modes/ZenModeDisplayLinkPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeDisplayLinkPreferenceController.java
index 943874a..8720a4b 100644
--- a/src/com/android/settings/notification/modes/ZenModeDisplayLinkPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeDisplayLinkPreferenceController.java
@@ -20,12 +20,15 @@
 
 import android.content.Context;
 import android.os.Bundle;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
+
 import com.android.settings.core.SubSettingLauncher;
 
-public class ZenModeDisplayLinkPreferenceController extends AbstractZenModePreferenceController  {
+class ZenModeDisplayLinkPreferenceController extends AbstractZenModePreferenceController  {
 
-    ZenModeSummaryHelper mSummaryHelper;
+    private final ZenModeSummaryHelper mSummaryHelper;
 
     public ZenModeDisplayLinkPreferenceController(Context context, String key,
             ZenModesBackend backend) {
@@ -34,10 +37,9 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
+    void updateState(Preference preference, @NonNull ZenMode zenMode) {
         Bundle bundle = new Bundle();
-        bundle.putString(MODE_ID, getMode().getId());
+        bundle.putString(MODE_ID, zenMode.getId());
         // TODO(b/332937635): Update metrics category
         preference.setIntent(new SubSettingLauncher(mContext)
                 .setDestination(ZenModeDisplayFragment.class.getName())
@@ -47,7 +49,7 @@
     }
 
     @Override
-    public CharSequence getSummary() {
-        return mSummaryHelper.getDisplayEffectsSummary(getMode());
+    public CharSequence getSummary(@NonNull ZenMode zenMode) {
+        return mSummaryHelper.getDisplayEffectsSummary(zenMode);
     }
 }
\ No newline at end of file
diff --git a/src/com/android/settings/notification/modes/ZenModeHeaderController.java b/src/com/android/settings/notification/modes/ZenModeHeaderController.java
index f55c02d..f791519 100644
--- a/src/com/android/settings/notification/modes/ZenModeHeaderController.java
+++ b/src/com/android/settings/notification/modes/ZenModeHeaderController.java
@@ -47,9 +47,8 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        ZenMode mode = getMode();
-        if (mode == null || mFragment == null) {
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
+        if (mFragment == null) {
             return;
         }
 
@@ -62,9 +61,9 @@
         }
 
         FutureUtil.whenDone(
-                mode.getIcon(IconLoader.getInstance(mContext)),
+                zenMode.getIcon(IconLoader.getInstance(mContext)),
                 icon -> mHeaderController.setIcon(icon)
-                        .setLabel(mode.getRule().getName())
+                        .setLabel(zenMode.getRule().getName())
                         .done(false /* rebindActions */),
                 mContext.getMainExecutor());
     }
diff --git a/src/com/android/settings/notification/modes/ZenModeMessagesLinkPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeMessagesLinkPreferenceController.java
index 8261008..300ebbc 100644
--- a/src/com/android/settings/notification/modes/ZenModeMessagesLinkPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeMessagesLinkPreferenceController.java
@@ -20,10 +20,13 @@
 
 import android.content.Context;
 import android.os.Bundle;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
+
 import com.android.settings.core.SubSettingLauncher;
 
-public class ZenModeMessagesLinkPreferenceController extends AbstractZenModePreferenceController {
+class ZenModeMessagesLinkPreferenceController extends AbstractZenModePreferenceController {
     private final ZenModeSummaryHelper mSummaryHelper;
 
     public ZenModeMessagesLinkPreferenceController(Context context, String key,
@@ -33,11 +36,9 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
-
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         Bundle bundle = new Bundle();
-        bundle.putString(MODE_ID, getMode().getId());
+        bundle.putString(MODE_ID, zenMode.getId());
         // TODO(b/332937635): Update metrics category
         preference.setIntent(new SubSettingLauncher(mContext)
                 .setDestination(ZenModeMessagesFragment.class.getName())
@@ -46,6 +47,6 @@
                 .toIntent());
 
         preference.setEnabled(true);
-        preference.setSummary(mSummaryHelper.getMessagesSettingSummary(getMode().getPolicy()));
+        preference.setSummary(mSummaryHelper.getMessagesSettingSummary(zenMode.getPolicy()));
     }
 }
diff --git a/src/com/android/settings/notification/modes/ZenModeNotifVisLinkPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeNotifVisLinkPreferenceController.java
index 9807431..da3b3be 100644
--- a/src/com/android/settings/notification/modes/ZenModeNotifVisLinkPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeNotifVisLinkPreferenceController.java
@@ -20,12 +20,16 @@
 
 import android.content.Context;
 import android.os.Bundle;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
+
 import com.android.settings.core.SubSettingLauncher;
 
-public class ZenModeNotifVisLinkPreferenceController extends AbstractZenModePreferenceController  {
+class ZenModeNotifVisLinkPreferenceController extends AbstractZenModePreferenceController  {
 
-    ZenModeSummaryHelper mSummaryBuilder;
+    private final ZenModeSummaryHelper mSummaryBuilder;
+
     public ZenModeNotifVisLinkPreferenceController(Context context, String key,
             ZenModesBackend backend) {
         super(context, key, backend);
@@ -33,10 +37,9 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         Bundle bundle = new Bundle();
-        bundle.putString(MODE_ID, getMode().getId());
+        bundle.putString(MODE_ID, zenMode.getId());
         // TODO(b/332937635): Update metrics category
         preference.setIntent(new SubSettingLauncher(mContext)
                 .setDestination(ZenModeNotifVisFragment.class.getName())
@@ -46,7 +49,7 @@
     }
 
     @Override
-    public CharSequence getSummary() {
-        return mSummaryBuilder.getBlockedEffectsSummary(getMode());
+    public CharSequence getSummary(@NonNull ZenMode zenMode) {
+        return mSummaryBuilder.getBlockedEffectsSummary(zenMode);
     }
 }
\ No newline at end of file
diff --git a/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceController.java
index 9b89a64..39f0d3c 100644
--- a/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceController.java
@@ -18,9 +18,12 @@
 
 import android.content.Context;
 import android.service.notification.ZenPolicy;
+
+import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 import androidx.preference.CheckBoxPreference;
 import androidx.preference.Preference;
+
 import com.android.settings.widget.DisabledCheckBoxPreference;
 
 public class ZenModeNotifVisPreferenceController extends AbstractZenModePreferenceController
@@ -53,13 +56,13 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
 
-        boolean suppressed = !getMode().getPolicy().isVisualEffectAllowed(mEffect, false);
+        boolean suppressed = !zenMode.getPolicy().isVisualEffectAllowed(mEffect, false);
         boolean parentSuppressed = false;
         if (mParentSuppressedEffects != null) {
             for (@ZenPolicy.VisualEffect int parentEffect : mParentSuppressedEffects) {
-                if (!getMode().getPolicy().isVisualEffectAllowed(parentEffect, true)) {
+                if (!zenMode.getPolicy().isVisualEffectAllowed(parentEffect, true)) {
                     parentSuppressed = true;
                 }
             }
@@ -77,15 +80,6 @@
     @Override
     public boolean onPreferenceChange(Preference preference, Object newValue) {
         final boolean allowEffect = !((Boolean) newValue);
-
-        if (getMode().getPolicy().isVisualEffectAllowed(mEffect, true) != allowEffect) {
-            ZenPolicy diffPolicy = new ZenPolicy.Builder()
-                    .showVisualEffect(mEffect, allowEffect)
-                    .build();
-            getMode().setPolicy(diffPolicy);
-            mBackend.updateMode(getMode());
-        }
-
-        return true;
+        return savePolicy(policy -> policy.showVisualEffect(mEffect, allowEffect));
     }
 }
diff --git a/src/com/android/settings/notification/modes/ZenModeOtherLinkPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeOtherLinkPreferenceController.java
index a43f8b0..1a00207 100644
--- a/src/com/android/settings/notification/modes/ZenModeOtherLinkPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeOtherLinkPreferenceController.java
@@ -16,20 +16,22 @@
 
 package com.android.settings.notification.modes;
 
-
 import static com.android.settings.notification.modes.ZenModeFragmentBase.MODE_ID;
 
 import android.content.Context;
 import android.os.Bundle;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
+
 import com.android.settings.core.SubSettingLauncher;
 
 /**
  * Preference with a link and summary about what other sounds can break through the mode
  */
-public class ZenModeOtherLinkPreferenceController extends AbstractZenModePreferenceController {
+class ZenModeOtherLinkPreferenceController extends AbstractZenModePreferenceController {
 
-    ZenModeSummaryHelper mSummaryHelper;
+    private final ZenModeSummaryHelper mSummaryHelper;
 
     public ZenModeOtherLinkPreferenceController(Context context, String key,
             ZenModesBackend backend) {
@@ -38,15 +40,14 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         Bundle bundle = new Bundle();
-        bundle.putString(MODE_ID, getMode().getId());
+        bundle.putString(MODE_ID, zenMode.getId());
         preference.setIntent(new SubSettingLauncher(mContext)
                 .setDestination(ZenModeOtherFragment.class.getName())
                 .setSourceMetricsCategory(0)
                 .setArguments(bundle)
                 .toIntent());
-        preference.setSummary(mSummaryHelper.getOtherSoundCategoriesSummary(getMode()));
+        preference.setSummary(mSummaryHelper.getOtherSoundCategoriesSummary(zenMode));
     }
 }
diff --git a/src/com/android/settings/notification/modes/ZenModeOtherPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeOtherPreferenceController.java
index e31fa0f..a770164e 100644
--- a/src/com/android/settings/notification/modes/ZenModeOtherPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeOtherPreferenceController.java
@@ -23,11 +23,12 @@
 import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_SYSTEM;
 
 import android.content.Context;
-import android.service.notification.ZenPolicy;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
 import androidx.preference.TwoStatePreference;
 
-public class ZenModeOtherPreferenceController extends AbstractZenModePreferenceController
+class ZenModeOtherPreferenceController extends AbstractZenModePreferenceController
         implements Preference.OnPreferenceChangeListener {
 
     public ZenModeOtherPreferenceController(Context context, String key,
@@ -36,24 +37,15 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
-
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         TwoStatePreference pref = (TwoStatePreference) preference;
-        pref.setChecked(getMode().getPolicy().isCategoryAllowed(getCategory(), true));
+        pref.setChecked(zenMode.getPolicy().isCategoryAllowed(getCategory(), true));
     }
 
     @Override
     public boolean onPreferenceChange(Preference preference, Object newValue) {
         final boolean allow = (Boolean) newValue;
-
-        ZenPolicy diffPolicy = new ZenPolicy.Builder()
-                .allowCategory(getCategory(), allow)
-                .build();
-        getMode().setPolicy(diffPolicy);
-        mBackend.updateMode(getMode());
-
-        return true;
+        return savePolicy(policy -> policy.allowCategory(getCategory(), allow));
     }
 
     private int getCategory() {
diff --git a/src/com/android/settings/notification/modes/ZenModePeopleLinkPreferenceController.java b/src/com/android/settings/notification/modes/ZenModePeopleLinkPreferenceController.java
index f122006..55a83d6 100644
--- a/src/com/android/settings/notification/modes/ZenModePeopleLinkPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModePeopleLinkPreferenceController.java
@@ -16,21 +16,22 @@
 
 package com.android.settings.notification.modes;
 
-
 import static com.android.settings.notification.modes.ZenModeFragmentBase.MODE_ID;
 
 import android.content.Context;
 import android.os.Bundle;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
-import com.android.settings.R;
+
 import com.android.settings.core.SubSettingLauncher;
 
 /**
  * Preference with a link and summary about what calls and messages can break through the mode
  */
-public class ZenModePeopleLinkPreferenceController extends AbstractZenModePreferenceController {
+class ZenModePeopleLinkPreferenceController extends AbstractZenModePreferenceController {
 
-    ZenModeSummaryHelper mSummaryHelper;
+    private final ZenModeSummaryHelper mSummaryHelper;
 
     public ZenModePeopleLinkPreferenceController(Context context, String key,
             ZenModesBackend backend) {
@@ -39,16 +40,15 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         Bundle bundle = new Bundle();
-        bundle.putString(MODE_ID, getMode().getId());
+        bundle.putString(MODE_ID, zenMode.getId());
         // TODO(b/332937635): Update metrics category
         preference.setIntent(new SubSettingLauncher(mContext)
                 .setDestination(ZenModePeopleFragment.class.getName())
                 .setSourceMetricsCategory(0)
                 .setArguments(bundle)
                 .toIntent());
-        preference.setSummary(mSummaryHelper.getPeopleSummary(getMode()));
+        preference.setSummary(mSummaryHelper.getPeopleSummary(zenMode));
     }
 }
diff --git a/src/com/android/settings/notification/modes/ZenModePrioritySendersPreferenceController.java b/src/com/android/settings/notification/modes/ZenModePrioritySendersPreferenceController.java
index a71bbe8..16e8858 100644
--- a/src/com/android/settings/notification/modes/ZenModePrioritySendersPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModePrioritySendersPreferenceController.java
@@ -36,10 +36,13 @@
 import android.service.notification.ConversationChannelWrapper;
 import android.service.notification.ZenPolicy;
 import android.view.View;
+
+import androidx.annotation.NonNull;
 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.core.SubSettingLauncher;
 import com.android.settings.notification.app.ConversationListSettings;
@@ -59,7 +62,7 @@
  * bypass DND for calls or messages, which may be one of the following values: starred contacts, all
  * contacts, priority conversations (for messages only), anyone, or no one.
  */
-public class ZenModePrioritySendersPreferenceController
+class ZenModePrioritySendersPreferenceController
         extends AbstractZenModePreferenceController {
     private final boolean mIsMessages; // if this is false, then this preference is for calls
 
@@ -124,12 +127,12 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         if (mIsMessages) {
             updateChannelCounts();
         }
-        final int currContactsSetting = getPrioritySenders();
-        final int currConversationsSetting = getPriorityConversationSenders();
+        final int currContactsSetting = getPrioritySenders(zenMode.getPolicy());
+        final int currConversationsSetting = getPriorityConversationSenders(zenMode.getPolicy());
         for (SelectorWithWidgetPreference pref : mSelectorPreferences) {
             // for each preference, check whether the current state matches what this state
             // would look like if the button were checked.
@@ -173,17 +176,17 @@
         mNumImportantConversations = numImportantConversations;
     }
 
-    private int getPrioritySenders() {
+    private int getPrioritySenders(ZenPolicy policy) {
         if (mIsMessages) {
-            return getMode().getPolicy().getPriorityMessageSenders();
+            return policy.getPriorityMessageSenders();
         } else {
-            return getMode().getPolicy().getPriorityCallSenders();
+            return policy.getPriorityCallSenders();
         }
     }
 
-    private int getPriorityConversationSenders() {
+    private int getPriorityConversationSenders(ZenPolicy policy) {
         if (mIsMessages) {
-            return getMode().getPolicy().getPriorityConversationSenders();
+            return policy.getPriorityConversationSenders();
         }
         return CONVERSATION_SENDERS_UNSET;
     }
@@ -419,29 +422,31 @@
     @VisibleForTesting
     SelectorWithWidgetPreference.OnClickListener mSelectorClickListener =
             new SelectorWithWidgetPreference.OnClickListener() {
-        @Override
-        public void onRadioButtonClicked(SelectorWithWidgetPreference preference) {
-            // The settingsToSaveOnClick function takes whether the preference is a
-            // checkbox into account to determine whether this selection is checked or unchecked.
-            final int[] settingsToSave = settingsToSaveOnClick(preference,
-                    getPrioritySenders(), getPriorityConversationSenders());
-            final int prioritySendersSetting = settingsToSave[0];
-            final int priorityConvosSetting = settingsToSave[1];
+                @Override
+                public void onRadioButtonClicked(SelectorWithWidgetPreference preference) {
+                    savePolicy(policy -> {
+                        ZenPolicy previousPolicy = policy.build();
+                        // The settingsToSaveOnClick function takes whether the preference is a
+                        // checkbox into account to determine whether this selection is checked or
+                        // unchecked.
+                        final int[] settingsToSave = settingsToSaveOnClick(preference,
+                                getPrioritySenders(previousPolicy),
+                                getPriorityConversationSenders(previousPolicy));
+                        final int prioritySendersSetting = settingsToSave[0];
+                        final int priorityConvosSetting = settingsToSave[1];
 
-            ZenPolicy.Builder diffPolicy = new ZenPolicy.Builder();
-            if (prioritySendersSetting != PEOPLE_TYPE_UNSET) {
-                if (mIsMessages) {
-                    diffPolicy.allowMessages(prioritySendersSetting);
-
-                } else {
-                    diffPolicy.allowCalls(prioritySendersSetting);
+                        if (prioritySendersSetting != PEOPLE_TYPE_UNSET) {
+                            if (mIsMessages) {
+                                policy.allowMessages(prioritySendersSetting);
+                            } else {
+                                policy.allowCalls(prioritySendersSetting);
+                            }
+                        }
+                        if (mIsMessages && priorityConvosSetting != CONVERSATION_SENDERS_UNSET) {
+                            policy.allowConversations(priorityConvosSetting);
+                        }
+                        return policy;
+                    });
                 }
-            }
-            if (mIsMessages && priorityConvosSetting != CONVERSATION_SENDERS_UNSET) {
-                diffPolicy.allowConversations(priorityConvosSetting);
-            }
-            getMode().setPolicy(diffPolicy.build());
-            mBackend.updateMode(getMode());
-        }
-    };
+            };
 }
diff --git a/src/com/android/settings/notification/modes/ZenModeRepeatCallersPreferenceController.java b/src/com/android/settings/notification/modes/ZenModeRepeatCallersPreferenceController.java
index d6de9c2..7569051 100644
--- a/src/com/android/settings/notification/modes/ZenModeRepeatCallersPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModeRepeatCallersPreferenceController.java
@@ -17,20 +17,17 @@
 package com.android.settings.notification.modes;
 
 import static android.service.notification.ZenPolicy.PEOPLE_TYPE_ANYONE;
-import static android.service.notification.ZenPolicy.PEOPLE_TYPE_NONE;
 import static android.service.notification.ZenPolicy.STATE_ALLOW;
 
-import android.app.settings.SettingsEnums;
 import android.content.Context;
-import android.provider.Settings;
-import android.service.notification.ZenPolicy;
-import android.util.Log;
+
+import androidx.annotation.NonNull;
 import androidx.preference.Preference;
-import androidx.preference.PreferenceScreen;
 import androidx.preference.TwoStatePreference;
+
 import com.android.settings.R;
 
-public class ZenModeRepeatCallersPreferenceController extends AbstractZenModePreferenceController
+class ZenModeRepeatCallersPreferenceController extends AbstractZenModePreferenceController
         implements Preference.OnPreferenceChangeListener {
 
     private final int mRepeatCallersThreshold;
@@ -43,14 +40,12 @@
     }
 
     @Override
-    public void updateState(Preference preference) {
-        super.updateState(preference);
-
+    public void updateState(Preference preference, @NonNull ZenMode zenMode) {
         TwoStatePreference pref = (TwoStatePreference) preference;
 
         boolean anyCallersCanBypassDnd =
-                getMode().getPolicy().getPriorityCategoryCalls() == STATE_ALLOW
-                && getMode().getPolicy().getPriorityCallSenders() == PEOPLE_TYPE_ANYONE;
+                zenMode.getPolicy().getPriorityCategoryCalls() == STATE_ALLOW
+                && zenMode.getPolicy().getPriorityCallSenders() == PEOPLE_TYPE_ANYONE;
         // if any caller can bypass dnd then repeat callers preference is disabled
         if (anyCallersCanBypassDnd) {
             pref.setEnabled(false);
@@ -58,21 +53,16 @@
         } else {
             pref.setEnabled(true);
             pref.setChecked(
-                    getMode().getPolicy().getPriorityCategoryRepeatCallers() == STATE_ALLOW);
+                    zenMode.getPolicy().getPriorityCategoryRepeatCallers() == STATE_ALLOW);
         }
 
         setRepeatCallerSummary(preference);
     }
 
     @Override
-    public boolean onPreferenceChange(Preference preference, Object newValue) {
+    public boolean onPreferenceChange(@NonNull Preference preference, Object newValue) {
         final boolean allowRepeatCallers = (Boolean) newValue;
-        ZenPolicy diffPolicy = new ZenPolicy.Builder()
-                .allowRepeatCallers(allowRepeatCallers)
-                .build();
-        getMode().setPolicy(diffPolicy);
-        mBackend.updateMode(getMode());
-        return true;
+        return savePolicy(policy -> policy.allowRepeatCallers(allowRepeatCallers));
     }
 
     private void setRepeatCallerSummary(Preference preference) {
diff --git a/src/com/android/settings/notification/modes/ZenModeSummaryHelper.java b/src/com/android/settings/notification/modes/ZenModeSummaryHelper.java
index e1a7caf..41a3d20 100644
--- a/src/com/android/settings/notification/modes/ZenModeSummaryHelper.java
+++ b/src/com/android/settings/notification/modes/ZenModeSummaryHelper.java
@@ -30,7 +30,6 @@
 import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_REMINDERS;
 import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_REPEAT_CALLERS;
 import static android.service.notification.ZenPolicy.PRIORITY_CATEGORY_SYSTEM;
-import static android.service.notification.ZenPolicy.STATE_DISALLOW;
 import static android.service.notification.ZenPolicy.VISUAL_EFFECT_AMBIENT;
 import static android.service.notification.ZenPolicy.VISUAL_EFFECT_BADGE;
 import static android.service.notification.ZenPolicy.VISUAL_EFFECT_FULL_SCREEN_INTENT;
@@ -44,7 +43,6 @@
 import android.service.notification.ZenDeviceEffects;
 import android.service.notification.ZenPolicy;
 
-import android.util.SparseArray;
 import com.android.settings.R;
 
 import java.util.ArrayList;
@@ -54,10 +52,10 @@
 import java.util.Map;
 import java.util.function.Predicate;
 
-public class ZenModeSummaryHelper {
+class ZenModeSummaryHelper {
 
-    private Context mContext;
-    private ZenModesBackend mBackend;
+    private final Context mContext;
+    private final ZenModesBackend mBackend;
 
     public ZenModeSummaryHelper(Context context, ZenModesBackend backend) {
         mContext = context;
diff --git a/src/com/android/settings/notification/modes/ZenModesListPreferenceController.java b/src/com/android/settings/notification/modes/ZenModesListPreferenceController.java
index ff5c111..ca8fe05 100644
--- a/src/com/android/settings/notification/modes/ZenModesListPreferenceController.java
+++ b/src/com/android/settings/notification/modes/ZenModesListPreferenceController.java
@@ -39,7 +39,7 @@
  * containing links to each individual mode. This is a central controller that populates and updates
  * all the preferences that then lead to a mode configuration page.
  */
-public class ZenModesListPreferenceController extends BasePreferenceController {
+class ZenModesListPreferenceController extends BasePreferenceController {
     protected static final String KEY = "zen_modes_list";
 
     @Nullable
diff --git a/tests/robotests/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceControllerTest.java
index 7424ae6..05b4848 100644
--- a/tests/robotests/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/notification/modes/ZenModeNotifVisPreferenceControllerTest.java
@@ -19,18 +19,16 @@
 import static android.app.NotificationManager.INTERRUPTION_FILTER_PRIORITY;
 import static android.service.notification.ZenPolicy.STATE_ALLOW;
 import static android.service.notification.ZenPolicy.STATE_DISALLOW;
-import static android.service.notification.ZenPolicy.STATE_UNSET;
 import static android.service.notification.ZenPolicy.VISUAL_EFFECT_LIGHTS;
 import static android.service.notification.ZenPolicy.VISUAL_EFFECT_NOTIFICATION_LIST;
 import static android.service.notification.ZenPolicy.VISUAL_EFFECT_PEEK;
 import static android.service.notification.ZenPolicy.VISUAL_EFFECT_STATUS_BAR;
+
 import static com.google.common.truth.Truth.assertThat;
+
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.eq;
-import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -41,12 +39,10 @@
 import android.net.Uri;
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
-import android.service.notification.ZenDeviceEffects;
 import android.service.notification.ZenPolicy;
-import androidx.preference.TwoStatePreference;
-import com.android.settings.notification.zen.ZenModeVisEffectPreferenceController;
+
 import com.android.settings.widget.DisabledCheckBoxPreference;
-import com.android.settingslib.core.lifecycle.Lifecycle;
+
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -162,7 +158,7 @@
         assertThat(captor.getValue().getPolicy().getVisualEffectStatusBar())
                 .isEqualTo(STATE_DISALLOW);
         assertThat(captor.getValue().getPolicy().getVisualEffectNotificationList())
-                .isEqualTo(STATE_UNSET);
+                .isEqualTo(STATE_DISALLOW); // Untouched
     }
 
     @Test
@@ -211,7 +207,7 @@
         assertThat(captor.getValue().getPolicy().getVisualEffectPeek())
                 .isEqualTo(STATE_ALLOW);
         assertThat(captor.getValue().getPolicy().getVisualEffectNotificationList())
-                .isEqualTo(STATE_UNSET);
+                .isEqualTo(STATE_DISALLOW); // Untouched
     }
 
     @Test
@@ -236,6 +232,6 @@
         assertThat(captor.getValue().getPolicy().getVisualEffectPeek())
                 .isEqualTo(STATE_DISALLOW);
         assertThat(captor.getValue().getPolicy().getVisualEffectNotificationList())
-                .isEqualTo(STATE_UNSET);
+                .isEqualTo(STATE_ALLOW); // Untouched
     }
 }
\ No newline at end of file