Merge "Avoid double quiet mode requests" into udc-dev
diff --git a/src/com/android/settings/accounts/WorkModePreferenceController.java b/src/com/android/settings/accounts/WorkModePreferenceController.java
index a261afc..46440f6 100644
--- a/src/com/android/settings/accounts/WorkModePreferenceController.java
+++ b/src/com/android/settings/accounts/WorkModePreferenceController.java
@@ -16,25 +16,29 @@
 package com.android.settings.accounts;
 
 import android.content.Context;
+import android.widget.Switch;
 
 import androidx.lifecycle.DefaultLifecycleObserver;
 import androidx.lifecycle.LifecycleOwner;
 import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
 
-import com.android.internal.annotations.VisibleForTesting;
 import com.android.settings.R;
+import com.android.settings.core.BasePreferenceController;
 import com.android.settings.slices.SliceData;
-import com.android.settings.widget.SettingsMainSwitchPreferenceController;
 import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
 
 import org.jetbrains.annotations.NotNull;
 
 
-public class WorkModePreferenceController extends SettingsMainSwitchPreferenceController
-        implements Preference.OnPreferenceChangeListener, DefaultLifecycleObserver,
+/** Controller for "Work apps" toggle that allows the user to enable/disable quiet mode. */
+public class WorkModePreferenceController extends BasePreferenceController
+        implements OnMainSwitchChangeListener, DefaultLifecycleObserver,
         ManagedProfileQuietModeEnabler.QuietModeChangeListener {
 
     private final ManagedProfileQuietModeEnabler mQuietModeEnabler;
+    private MainSwitchPreference mPreference;
 
     public WorkModePreferenceController(Context context, String key) {
         super(context, key);
@@ -42,6 +46,13 @@
     }
 
     @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        mPreference = screen.findPreference(getPreferenceKey());
+        mPreference.addOnSwitchChangeListener(this);
+    }
+
+    @Override
     public int getAvailabilityStatus() {
         return (mQuietModeEnabler.isAvailable()) ? AVAILABLE : DISABLED_FOR_USER;
     }
@@ -57,19 +68,18 @@
     }
 
     @Override
-    public boolean isChecked() {
-        return !mQuietModeEnabler.isQuietModeEnabled();
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
+        mQuietModeEnabler.setQuietModeEnabled(!isChecked);
     }
 
     @Override
-    public boolean setChecked(boolean isChecked) {
-        mQuietModeEnabler.setQuietModeEnabled(!isChecked);
-        return true;
+    public final void updateState(Preference preference) {
+        mPreference.updateStatus(!mQuietModeEnabler.isQuietModeEnabled());
     }
 
     @Override
     public void onQuietModeChanged() {
-        updateState(mSwitchPreference);
+        updateState(mPreference);
     }
 
     @Override
@@ -82,9 +92,4 @@
     public int getSliceHighlightMenuRes() {
         return R.string.menu_key_accounts;
     }
-
-    @VisibleForTesting
-    void setPreference(MainSwitchPreference preference) {
-        mSwitchPreference = preference;
-    }
 }
diff --git a/tests/robotests/src/com/android/settings/accounts/WorkModePreferenceControllerTest.java b/tests/robotests/src/com/android/settings/accounts/WorkModePreferenceControllerTest.java
index e862d10..d6be73d 100644
--- a/tests/robotests/src/com/android/settings/accounts/WorkModePreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/accounts/WorkModePreferenceControllerTest.java
@@ -19,6 +19,7 @@
 
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -27,7 +28,9 @@
 import android.content.pm.UserInfo;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.widget.Switch;
 
+import androidx.preference.PreferenceScreen;
 import androidx.test.core.app.ApplicationProvider;
 
 import com.android.settingslib.widget.MainSwitchPreference;
@@ -57,6 +60,10 @@
     private UserHandle mManagedUser;
     @Mock
     private UserInfo mUserInfo;
+    @Mock
+    private PreferenceScreen mScreen;
+    @Mock
+    Switch mSwitch;
 
     @Before
     public void setUp() {
@@ -69,7 +76,9 @@
         when(mUserManager.getProcessUserId()).thenReturn(0);
         when(mUserManager.getUserProfiles()).thenReturn(Collections.singletonList(mManagedUser));
         when(mManagedUser.getIdentifier()).thenReturn(MANAGED_USER_ID);
+        when(mScreen.findPreference(anyString())).thenReturn(mPreference);
         mController = new WorkModePreferenceController(mContext, PREF_KEY);
+        mController.displayPreference(mScreen);
     }
 
     @Test
@@ -106,13 +115,11 @@
 
     @Test
     public void onPreferenceChange_shouldRequestQuietModeEnabled() {
-        mController.setPreference(mPreference);
-
-        mController.onPreferenceChange(mPreference, true);
+        mController.onSwitchChanged(mSwitch, true);
 
         verify(mUserManager).requestQuietModeEnabled(false, mManagedUser);
 
-        mController.onPreferenceChange(mPreference, false);
+        mController.onSwitchChanged(mSwitch, false);
 
         verify(mUserManager).requestQuietModeEnabled(true, mManagedUser);
     }