From a7b13c11cddce9c6ddd4fd89ed9c2e931f9086ec Mon Sep 17 00:00:00 2001
From: Weng Su <wengsu@google.com>
Date: Fri, 5 Jul 2024 06:51:01 +0800
Subject: [PATCH] Disable "Add network" if DISALLOW_ADD_WIFI_CONFIG

- When DISALLOW_ADD_WIFI_CONFIG is not restricted by Admin
  - Disable the preference directly

Bug: 335387874
Flag: EXEMPT bugfix
Test: Manual testing
atest -c AddWifiNetworkPreferenceTest

Change-Id: I0bb08656971107c73d7632a395a820d03723620b
---
 .../wifi/AddWifiNetworkPreference.java        | 15 ++++++-
 .../wifi/AddWifiNetworkPreferenceTest.java    | 39 +++++++++++++------
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/src/com/android/settings/wifi/AddWifiNetworkPreference.java b/src/com/android/settings/wifi/AddWifiNetworkPreference.java
index 999f4a4d0b5..ea229541886 100644
--- a/src/com/android/settings/wifi/AddWifiNetworkPreference.java
+++ b/src/com/android/settings/wifi/AddWifiNetworkPreference.java
@@ -26,11 +26,13 @@ import android.widget.ImageButton;
 
 import androidx.annotation.DrawableRes;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import androidx.preference.PreferenceViewHolder;
 
 import com.android.settings.R;
 import com.android.settings.wifi.dpp.WifiDppUtils;
 import com.android.settingslib.RestrictedPreference;
+import com.android.settingslib.wifi.WifiEnterpriseRestrictionUtils;
 
 /**
  * The Preference for users to add Wi-Fi networks in WifiSettings
@@ -53,7 +55,7 @@ public class AddWifiNetworkPreference extends RestrictedPreference {
         setTitle(R.string.wifi_add_network);
 
         mScanIconDrawable = getDrawable(R.drawable.ic_scan_24dp);
-        checkRestrictionAndSetDisabled(UserManager.DISALLOW_ADD_WIFI_CONFIG);
+        checkRestrictionAndSetDisabled();
     }
 
     @Override
@@ -80,4 +82,15 @@ public class AddWifiNetworkPreference extends RestrictedPreference {
         }
         return buttonIcon;
     }
+
+    @VisibleForTesting
+    void checkRestrictionAndSetDisabled() {
+        checkRestrictionAndSetDisabled(UserManager.DISALLOW_ADD_WIFI_CONFIG);
+        if (isDisabledByAdmin()) {
+            return;
+        }
+        if (!WifiEnterpriseRestrictionUtils.isAddWifiConfigAllowed(getContext())) {
+            setEnabled(false);
+        }
+    }
 }
diff --git a/tests/robotests/src/com/android/settings/wifi/AddWifiNetworkPreferenceTest.java b/tests/robotests/src/com/android/settings/wifi/AddWifiNetworkPreferenceTest.java
index 5f88eb2a083..c0adaf3ec36 100644
--- a/tests/robotests/src/com/android/settings/wifi/AddWifiNetworkPreferenceTest.java
+++ b/tests/robotests/src/com/android/settings/wifi/AddWifiNetworkPreferenceTest.java
@@ -15,19 +15,23 @@
  */
 package com.android.settings.wifi;
 
+import static android.os.UserManager.DISALLOW_ADD_WIFI_CONFIG;
+
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.when;
+
 import android.content.Context;
 import android.os.UserManager;
 
 import androidx.test.core.app.ApplicationProvider;
 
-import com.android.settingslib.RestrictedLockUtils;
-
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
 import org.mockito.Spy;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
@@ -42,31 +46,44 @@ public class AddWifiNetworkPreferenceTest {
     public final MockitoRule mMockitoRule = MockitoJUnit.rule();
     @Spy
     Context mContext = ApplicationProvider.getApplicationContext();
+    @Mock
+    private UserManager mUserManager;
 
     private AddWifiNetworkPreference mPreference;
 
     @Before
     public void setUp() {
-        mPreference = new AddWifiNetworkPreference(mContext);
+        when(mContext.getSystemService(UserManager.class)).thenReturn(mUserManager);
+
+        mPreference = spy(new AddWifiNetworkPreference(mContext));
     }
 
     @Test
-    public void updatePreferenceForRestriction_isAddWifiConfigAllowed_prefIsEnabled() {
-        // If the user is allowed to add Wi-Fi configuration then the EnforcedAdmin will be null.
-        RestrictedLockUtils.EnforcedAdmin enforcedAdmin = null;
+    public void checkRestrictionAndSetDisabled_disabledByAdmin_keepEnabledForClicks() {
+        when(mPreference.isDisabledByAdmin()).thenReturn(true);
 
-        mPreference.setDisabledByAdmin(enforcedAdmin);
+        mPreference.checkRestrictionAndSetDisabled();
 
         assertThat(mPreference.isEnabled()).isTrue();
     }
 
     @Test
-    public void updatePreferenceForRestriction_isAddWifiConfigNotAllowed_prefIsDisabled() {
-        RestrictedLockUtils.EnforcedAdmin enforcedAdmin = new RestrictedLockUtils.EnforcedAdmin(
-                null /* component */, UserManager.DISALLOW_ADD_WIFI_CONFIG, null /* user */);
+    public void checkRestrictionAndSetDisabled_notDisabledByAdmin_setDisabled() {
+        when(mPreference.isDisabledByAdmin()).thenReturn(false);
+        when(mUserManager.hasUserRestriction(DISALLOW_ADD_WIFI_CONFIG)).thenReturn(true);
 
-        mPreference.setDisabledByAdmin(enforcedAdmin);
+        mPreference.checkRestrictionAndSetDisabled();
 
         assertThat(mPreference.isEnabled()).isFalse();
     }
+
+    @Test
+    public void checkRestrictionAndSetDisabled_noRestriction_setEnabled() {
+        when(mPreference.isDisabledByAdmin()).thenReturn(false);
+        when(mUserManager.hasUserRestriction(DISALLOW_ADD_WIFI_CONFIG)).thenReturn(false);
+
+        mPreference.checkRestrictionAndSetDisabled();
+
+        assertThat(mPreference.isEnabled()).isTrue();
+    }
 }