Change zen mode schedules page rule handling

The approach of reloading all rules (and recreating all ZenRulePreferences) every time the rule set changed causes the switches on the page to stop working.

This change keeps the ZenRulePreference around as long as the rule itself is around and keeps it updated, while re-adding the preferences to the PreferenceCategory if needed due to rules changing.

Fixes: 229879326
Test: manual, ZenModeAutomaticRulesPreferenceControllerTest
Change-Id: I4eba41e8252cedd87ac866e4b97513970ca2d94a

src/com/android/settings/notification/zen/ZenModeAutomaticRulesPreferenceController.java

@@ -18,6 +18,7 @@ package com.android.settings.notification.zen;
 
 import android.app.AutomaticZenRule;
 import android.content.Context;
+import android.util.ArrayMap;
 
 import androidx.annotation.VisibleForTesting;
 import androidx.fragment.app.Fragment;
@@ -28,7 +29,6 @@ import androidx.preference.PreferenceScreen;
 import com.android.settingslib.core.lifecycle.Lifecycle;
 
 import java.util.Map;
-import java.util.Objects;
 
 public class ZenModeAutomaticRulesPreferenceController extends
         AbstractZenModeAutomaticRulePreferenceController {
@@ -38,6 +38,10 @@ public class ZenModeAutomaticRulesPreferenceController extends
     @VisibleForTesting
     protected PreferenceCategory mPreferenceCategory;
 
+    // Map of rule key -> preference so that we can update each preference as needed
+    @VisibleForTesting
+    protected Map<String, ZenRulePreference> mZenRulePreferences = new ArrayMap<>();
+
     public ZenModeAutomaticRulesPreferenceController(Context context, Fragment parent, Lifecycle
             lifecycle) {
         super(context, KEY, parent, lifecycle);
@@ -58,38 +62,73 @@ public class ZenModeAutomaticRulesPreferenceController extends
         super.displayPreference(screen);
         mPreferenceCategory = screen.findPreference(getPreferenceKey());
         mPreferenceCategory.setPersistent(false);
+
+        // if mPreferenceCategory was un-set, make sure to clear out mZenRulePreferences too, just
+        // in case
+        if (mPreferenceCategory.getPreferenceCount() == 0) {
+            mZenRulePreferences.clear();
+        }
     }
 
     @Override
     public void updateState(Preference preference) {
         super.updateState(preference);
         Map.Entry<String, AutomaticZenRule>[] sortedRules = getRules();
-        final int currNumPreferences = mPreferenceCategory.getPreferenceCount();
+
-        if (currNumPreferences == sortedRules.length) {
+        // refresh the whole preference category list if the total number of rules has changed, or
+        // if any individual rules have changed, so we can rebuild the list & keep things in sync
+        boolean refreshPrefs = false;
+        if (mPreferenceCategory.getPreferenceCount() != sortedRules.length) {
+            refreshPrefs = true;
+        } else {
+            // check whether any rules in sortedRules are not in mZenRulePreferences; that should
+            // be enough to see whether something has changed
             for (int i = 0; i < sortedRules.length; i++) {
-                ZenRulePreference pref = (ZenRulePreference) mPreferenceCategory.getPreference(i);
+                if (!mZenRulePreferences.containsKey(sortedRules[i].getKey())) {
-                // we are either:
+                    refreshPrefs = true;
-                // 1. updating everything about the rule
-                // 2. rule was added or deleted, so reload the entire list
-                if (Objects.equals(pref.mId, sortedRules[i].getKey())) {
-                    AutomaticZenRule rule = sortedRules[i].getValue();
-                    pref.updatePreference(rule);
-                } else {
-                    reloadAllRules(sortedRules);
                     break;
                 }
             }
-        } else {
-            reloadAllRules(sortedRules);
         }
-    }
 
-    @VisibleForTesting
+        // if we need to refresh the whole list, clear the preference category and also start a
-    void reloadAllRules(Map.Entry<String, AutomaticZenRule>[] rules) {
+        // new map of preferences according to the preference category contents
-        mPreferenceCategory.removeAll();
+        // we need to not update the existing one yet, as we'll need to know what preferences
-        for (Map.Entry<String, AutomaticZenRule> rule : rules) {
+        // previously existed in order to update and re-attach them to the preference category
-            ZenRulePreference pref = createZenRulePreference(rule);
+        Map<String, ZenRulePreference> newPrefs = new ArrayMap<>();
-            mPreferenceCategory.addPreference(pref);
+        if (refreshPrefs) {
+            mPreferenceCategory.removeAll();
+        }
+
+        // Loop through each rule, either updating the existing rule or creating the rule's
+        // preference if needed (and, in the case where we need to rebuild the preference category
+        // list, do so as well)
+        for (int i = 0; i < sortedRules.length; i++) {
+            String key = sortedRules[i].getKey();
+            if (mZenRulePreferences.containsKey(key)) {
+                // existing rule; update its info if it's changed since the last display
+                AutomaticZenRule rule = sortedRules[i].getValue();
+                ZenRulePreference pref = mZenRulePreferences.get(key);
+                pref.updatePreference(rule);
+
+                // only add to preference category if the overall set of rules has changed so this
+                // needs to be rearranged
+                if (refreshPrefs) {
+                    mPreferenceCategory.addPreference(pref);
+                    newPrefs.put(key, pref);
+                }
+            } else {
+                // new rule; create a new ZenRulePreference & add it to the preference category
+                // and the map so we'll know about it later
+                ZenRulePreference pref = createZenRulePreference(sortedRules[i]);
+                mPreferenceCategory.addPreference(pref);
+                newPrefs.put(key, pref);
+            }
+        }
+
+        // If anything was new, then make sure we overwrite mZenRulePreferences with our new data
+        if (refreshPrefs) {
+            mZenRulePreferences = newPrefs;
         }
     }
 

tests/robotests/src/com/android/settings/notification/zen/ZenModeAutomaticRulesPreferenceControllerTest.java

@@ -16,6 +16,8 @@
 
 package com.android.settings.notification.zen;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
@@ -33,10 +35,7 @@ import android.provider.Settings;
 import androidx.fragment.app.Fragment;
 import androidx.preference.PreferenceCategory;
 import androidx.preference.PreferenceScreen;
-
+import androidx.test.core.app.ApplicationProvider;
-import com.android.settings.notification.zen.ZenModeAutomaticRulesPreferenceController;
-import com.android.settings.notification.zen.ZenModeBackend;
-import com.android.settings.notification.zen.ZenRulePreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -45,7 +44,6 @@ import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 import org.mockito.internal.util.reflection.FieldSetter;
 import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 import org.robolectric.util.ReflectionHelpers;
 
 import java.util.HashMap;
@@ -68,7 +66,7 @@ public class ZenModeAutomaticRulesPreferenceControllerTest {
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = ApplicationProvider.getApplicationContext();
         mController = spy(new ZenModeAutomaticRulesPreferenceController(mContext, mock(Fragment.class),
             null));
         ReflectionHelpers.setField(mController, "mBackend", mBackend);
@@ -78,6 +76,16 @@ public class ZenModeAutomaticRulesPreferenceControllerTest {
         doReturn(mZenRulePreference).when(mController).createZenRulePreference(any());
     }
 
+    @Test
+    public void testDisplayPreference_resetsPreferencesWhenCategoryEmpty() {
+        // when the PreferenceCategory is empty (no preferences), make sure we clear out any
+        // stale state in the cached set of zen rule preferences
+        mController.mZenRulePreferences.put("test1_id", mZenRulePreference);
+        when(mockPref.getPreferenceCount()).thenReturn(0);
+        mController.displayPreference(mPreferenceScreen);
+        assertTrue(mController.mZenRulePreferences.isEmpty());
+    }
+
     @Test
     public void testUpdateState_clearsPreferencesWhenAddingNewPreferences() {
         final int NUM_RULES = 3;
@@ -103,6 +111,7 @@ public class ZenModeAutomaticRulesPreferenceControllerTest {
         mController.updateState(mockPref);
         verify(mockPref, times(1)).removeAll();
         verify(mockPref, times(NUM_RULES)).addPreference(any());
+        assertEquals(NUM_RULES, mController.mZenRulePreferences.size());
     }
 
     @Test
@@ -121,12 +130,49 @@ public class ZenModeAutomaticRulesPreferenceControllerTest {
         rMap.put(ruleId1, autoRule1);
         rMap.put(ruleId2, autoRule2);
 
+        // Add three preferences to the set of previously-known-about ZenRulePreferences; in this
+        // case, test3_id is "deleted"
+        mController.mZenRulePreferences.put("test1_id", mZenRulePreference);
+        mController.mZenRulePreferences.put("test2_id", mZenRulePreference);
+        mController.mZenRulePreferences.put("test3_id", mZenRulePreference);
+
         // update state should re-add all preferences since a preference was deleted
-        when(mockPref.getPreferenceCount()).thenReturn(NUM_RULES + 2);
+        when(mockPref.getPreferenceCount()).thenReturn(NUM_RULES + 1);
         mockGetAutomaticZenRules(NUM_RULES, rMap);
         mController.updateState(mockPref);
         verify(mockPref, times(1)).removeAll();
         verify(mockPref, times(NUM_RULES)).addPreference(any());
+        assertEquals(NUM_RULES, mController.mZenRulePreferences.size());
+    }
+
+    @Test
+    public void testUpdateState_clearsPreferencesWhenSameNumberButDifferentPrefs() {
+        final int NUM_RULES = 2;
+        Map<String, AutomaticZenRule> rMap = new HashMap<>();
+
+        String ruleId1 = "test1_id";
+        String ruleId2 = "test2_id";
+
+        AutomaticZenRule autoRule1 = new AutomaticZenRule("test_rule_1", null, null,
+                null, null, Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS, true, 10);
+        AutomaticZenRule autoRule2 = new AutomaticZenRule("test_rule_2", null, null,
+                null, null, Settings.Global.ZEN_MODE_IMPORTANT_INTERRUPTIONS, true, 20);
+
+        rMap.put(ruleId1, autoRule1);
+        rMap.put(ruleId2, autoRule2);
+
+        // Add two preferences to the set of previously-known-about ZenRulePreferences; in this
+        // case, test3_id is "deleted" but test2_id is "added"
+        mController.mZenRulePreferences.put("test1_id", mZenRulePreference);
+        mController.mZenRulePreferences.put("test3_id", mZenRulePreference);
+
+        // update state should re-add all preferences since a preference was deleted
+        when(mockPref.getPreferenceCount()).thenReturn(NUM_RULES);
+        mockGetAutomaticZenRules(NUM_RULES, rMap);
+        mController.updateState(mockPref);
+        verify(mockPref, times(1)).removeAll();
+        verify(mockPref, times(NUM_RULES)).addPreference(any());
+        assertEquals(NUM_RULES, mController.mZenRulePreferences.size());
     }
 
     @Test
@@ -140,6 +186,7 @@ public class ZenModeAutomaticRulesPreferenceControllerTest {
 
         when(mockPref.getPreferenceCount()).thenReturn(NUM_RULES);
         when(mockPref.getPreference(anyInt())).thenReturn(mZenRulePreference);
+        mController.mZenRulePreferences.put("test1_id", mZenRulePreference);
 
         // update state should NOT re-add all the preferences, should only update enable state
         rule.setEnabled(false);
@@ -148,7 +195,8 @@ public class ZenModeAutomaticRulesPreferenceControllerTest {
         FieldSetter.setField(mZenRulePreference, ZenRulePreference.class.getDeclaredField("mId"), testId);
         mController.updateState(mockPref);
         verify(mZenRulePreference, times(1)).updatePreference(any());
-        verify(mController, never()).reloadAllRules(any());
+        verify(mockPref, never()).removeAll();
+        assertEquals(NUM_RULES, mController.mZenRulePreferences.size());
     }
 
     private void mockGetAutomaticZenRules(int numRules, Map<String, AutomaticZenRule> rules) {