Merge "Fix toggle when no_user_switch is applied without EnforcedAdmin" into main
diff --git a/src/com/android/settings/users/MultiUserSwitchBarController.java b/src/com/android/settings/users/MultiUserSwitchBarController.java
index 641ae51..07c03d7 100644
--- a/src/com/android/settings/users/MultiUserSwitchBarController.java
+++ b/src/com/android/settings/users/MultiUserSwitchBarController.java
@@ -17,6 +17,7 @@
 package com.android.settings.users;
 
 import android.content.Context;
+import android.multiuser.Flags;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.Settings;
@@ -25,6 +26,7 @@
 import androidx.annotation.VisibleForTesting;
 
 import com.android.settings.widget.SwitchWidgetController;
+import com.android.settingslib.RestrictedLockUtils;
 import com.android.settingslib.RestrictedLockUtilsInternal;
 import com.android.settingslib.core.lifecycle.LifecycleObserver;
 import com.android.settingslib.core.lifecycle.events.OnStart;
@@ -53,12 +55,24 @@
         mUserCapabilities = UserCapabilities.create(context);
         mSwitchBar.setChecked(mUserCapabilities.mUserSwitcherEnabled);
 
-        if (mUserCapabilities.mDisallowSwitchUser) {
-            mSwitchBar.setDisabledByAdmin(RestrictedLockUtilsInternal
+        if (Flags.fixDisablingOfMuToggleWhenRestrictionApplied()) {
+            RestrictedLockUtils.EnforcedAdmin enforcedAdmin = RestrictedLockUtilsInternal
                     .checkIfRestrictionEnforced(mContext, UserManager.DISALLOW_USER_SWITCH,
-                            UserHandle.myUserId()));
+                            UserHandle.myUserId());
+            if (enforcedAdmin != null) {
+                mSwitchBar.setDisabledByAdmin(enforcedAdmin);
+            } else {
+                mSwitchBar.setEnabled(mUserCapabilities.mIsMain
+                        && !mUserCapabilities.mDisallowSwitchUser);
+            }
         } else {
-            mSwitchBar.setEnabled(mUserCapabilities.mIsMain);
+            if (mUserCapabilities.mDisallowSwitchUser) {
+                mSwitchBar.setDisabledByAdmin(RestrictedLockUtilsInternal
+                        .checkIfRestrictionEnforced(mContext, UserManager.DISALLOW_USER_SWITCH,
+                                UserHandle.myUserId()));
+            } else {
+                mSwitchBar.setEnabled(mUserCapabilities.mIsMain);
+            }
         }
         mSwitchBar.setListener(this);
     }
diff --git a/tests/robotests/src/com/android/settings/users/MultiUserSwitchBarControllerTest.java b/tests/robotests/src/com/android/settings/users/MultiUserSwitchBarControllerTest.java
index bfab257..58288a5 100644
--- a/tests/robotests/src/com/android/settings/users/MultiUserSwitchBarControllerTest.java
+++ b/tests/robotests/src/com/android/settings/users/MultiUserSwitchBarControllerTest.java
@@ -23,28 +23,40 @@
 
 import android.content.Context;
 import android.content.pm.UserInfo;
+import android.multiuser.Flags;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.platform.test.annotations.RequiresFlagsDisabled;
+import android.platform.test.annotations.RequiresFlagsEnabled;
+import android.platform.test.flag.junit.CheckFlagsRule;
+import android.platform.test.flag.junit.DeviceFlagsValueProvider;
 
+import com.android.settings.testutils.shadow.ShadowDevicePolicyManager;
 import com.android.settings.testutils.shadow.ShadowUserManager;
 import com.android.settings.widget.SwitchWidgetController;
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
 import org.robolectric.RuntimeEnvironment;
 import org.robolectric.annotation.Config;
 
+import java.util.ArrayList;
+import java.util.List;
+
 @RunWith(RobolectricTestRunner.class)
-@Config(shadows = {ShadowUserManager.class})
+@Config(shadows = {ShadowUserManager.class, ShadowDevicePolicyManager.class})
 public class MultiUserSwitchBarControllerTest {
 
     private Context mContext;
     private ShadowUserManager mUserManager;
     private SwitchWidgetController mSwitchWidgetController;
 
+    @Rule
+    public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
 
     @Before
     public void setUp() {
@@ -60,6 +72,7 @@
     }
 
     @Test
+    @RequiresFlagsDisabled({Flags.FLAG_FIX_DISABLING_OF_MU_TOGGLE_WHEN_RESTRICTION_APPLIED})
     public void onStart_disallowUserSwitch_shouldSetDisabledByAdmin() {
         mUserManager.setUserRestriction(UserHandle.of(UserHandle.myUserId()),
                 UserManager.DISALLOW_USER_SWITCH, true);
@@ -71,6 +84,48 @@
     }
 
     @Test
+    @RequiresFlagsEnabled({Flags.FLAG_FIX_DISABLING_OF_MU_TOGGLE_WHEN_RESTRICTION_APPLIED})
+    public void onStart_disallowUserSwitchEnforcedByAdmin_shouldSetDisabledByAdminUnchecked() {
+        int userId = UserHandle.myUserId();
+        List<UserManager.EnforcingUser> enforcingUsers = new ArrayList<>();
+        enforcingUsers.add(new UserManager.EnforcingUser(userId,
+                UserManager.RESTRICTION_SOURCE_DEVICE_OWNER));
+        // Ensure that RestrictedLockUtils.checkIfRestrictionEnforced doesn't return null.
+        ShadowUserManager.getShadow().setUserRestrictionSources(
+                UserManager.DISALLOW_USER_SWITCH,
+                UserHandle.of(userId),
+                enforcingUsers);
+
+        new MultiUserSwitchBarController(mContext, mSwitchWidgetController, null);
+        verify(mSwitchWidgetController).setChecked(false);
+        verify(mSwitchWidgetController).setDisabledByAdmin(any());
+    }
+
+    @Test
+    @RequiresFlagsEnabled({Flags.FLAG_FIX_DISABLING_OF_MU_TOGGLE_WHEN_RESTRICTION_APPLIED})
+    public void onStart_disallowUserSwitch_userNotMain_shouldSetDisabledUnchecked() {
+        mUserManager.setUserRestriction(UserHandle.of(UserHandle.myUserId()),
+                UserManager.DISALLOW_USER_SWITCH, true);
+        new MultiUserSwitchBarController(mContext, mSwitchWidgetController, null);
+
+        verify(mSwitchWidgetController).setChecked(false);
+        verify(mSwitchWidgetController).setEnabled(false);
+        verify(mSwitchWidgetController, never()).setDisabledByAdmin(any());
+    }
+
+    @Test
+    @RequiresFlagsEnabled({Flags.FLAG_FIX_DISABLING_OF_MU_TOGGLE_WHEN_RESTRICTION_APPLIED})
+    public void onStart_allowUserSwitch_notMainUser_shouldSetDisabled() {
+        mUserManager.setUserRestriction(UserHandle.of(UserHandle.myUserId()),
+                UserManager.DISALLOW_USER_SWITCH, false);
+        mUserManager.addUser(10, "Test", UserInfo.FLAG_ADMIN);
+        mUserManager.switchUser(10);
+        new MultiUserSwitchBarController(mContext, mSwitchWidgetController, null);
+
+        verify(mSwitchWidgetController).setEnabled(false);
+    }
+
+    @Test
     public void onStart_allowUserSwitch_shouldNotSetDisabledByAdmin() {
         mUserManager.setUserRestriction(UserHandle.of(UserHandle.myUserId()),
                 UserManager.DISALLOW_USER_SWITCH, false);