Disallow modifying USB in Settings

Background
* The admin can disable signaling over USB
  on corporate owned device.

Changes
* If the admin has disabled signaling over
  USB, disable preferences in Settings to
  modify USB.
* Show policy transparency dialogs.

Manual Testing:
* Install TestDPC
* Set up as Device Owner
* Disable USB data signaling
* Verify Settings > Connected Devices
  > USB is disabled
* Verify Settings > Network & Internet
  > Hotspot & tethering > USB tethering
  is disabled
* Verify Developer options >
  USB debugging, default USB configurations
  and USB audio routing are disabled

Bug: 168301639
Test: Manual Testing
      make RunSettingsRoboTests -j ROBOTEST_FILTER=TetherSettingsTest
      make RunSettingsRoboTests -j
      ROBOTEST_FILTER=UsbAudioRoutingPreferenceControllerTest
      make RunSettingsRoboTests -j
      ROBOTEST_FILTER=UsbDefaultFragmentTest
      make RunSettingsRoboTests -j
      ROBOTEST_FILTER=DevelopmentSettingsDashboardFragmentTest
      make RunSettingsRoboTests -j
      ROBOTEST_FILTER=DefaultUsbConfigurationPreferenceControllerTest
      make RunSettingsLibRoboTests -j
      ROBOTEST_FILTER=EnableAdbPreferenceControllerTest

Change-Id: I3e8461ac37f406a407c874598335cde28fb89b2d
Merged-In: I3e8461ac37f406a407c874598335cde28fb89b2d
diff --git a/res/xml/development_settings.xml b/res/xml/development_settings.xml
index 67d857f..7827180 100644
--- a/res/xml/development_settings.xml
+++ b/res/xml/development_settings.xml
@@ -135,7 +135,7 @@
         android:title="@string/debug_debugging_category"
         android:order="200">
 
-        <SwitchPreference
+        <com.android.settingslib.RestrictedSwitchPreference
             android:key="enable_adb"
             android:title="@string/enable_adb"
             android:summary="@string/enable_adb_summary" />
@@ -294,7 +294,7 @@
             android:title="@string/tethering_hardware_offload"
             android:summary="@string/tethering_hardware_offload_summary" />
 
-        <Preference
+        <com.android.settingslib.RestrictedSwitchPreference
             android:key="default_usb_configuration"
             android:fragment="com.android.settings.connecteddevice.usb.UsbDefaultFragment"
             android:icon="@drawable/ic_usb"
@@ -519,7 +519,7 @@
         android:title="@string/media_category"
         android:order="800">
 
-        <SwitchPreference
+        <com.android.settingslib.RestrictedSwitchPreference
             android:key="usb_audio"
             android:title="@string/usb_audio_disable_routing"
             android:summary="@string/usb_audio_disable_routing_summary" />
diff --git a/res/xml/tether_prefs.xml b/res/xml/tether_prefs.xml
index cca4e5f..37194d9 100644
--- a/res/xml/tether_prefs.xml
+++ b/res/xml/tether_prefs.xml
@@ -31,7 +31,7 @@
         settings:allowDividerAbove="true"
         settings:summaryLineCount="2"/>
 
-    <SwitchPreference
+    <com.android.settingslib.RestrictedSwitchPreference
         android:key="usb_tether_settings"
         android:title="@string/usb_tethering_button_text"
         android:summary="@string/usb_tethering_subtext"
diff --git a/src/com/android/settings/TetherSettings.java b/src/com/android/settings/TetherSettings.java
index c33f7db..48dd218 100644
--- a/src/com/android/settings/TetherSettings.java
+++ b/src/com/android/settings/TetherSettings.java
@@ -20,6 +20,8 @@
 import static android.net.ConnectivityManager.TETHERING_USB;
 import static android.net.TetheringManager.TETHERING_ETHERNET;
 
+import static com.android.settingslib.RestrictedLockUtilsInternal.checkIfUsbDataSignalingIsDisabled;
+
 import android.app.Activity;
 import android.app.settings.SettingsEnums;
 import android.bluetooth.BluetoothAdapter;
@@ -38,6 +40,7 @@
 import android.os.Environment;
 import android.os.Handler;
 import android.os.HandlerExecutor;
+import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.SearchIndexableResource;
 import android.text.TextUtils;
@@ -51,6 +54,7 @@
 import com.android.settings.datausage.DataSaverBackend;
 import com.android.settings.search.BaseSearchIndexProvider;
 import com.android.settings.wifi.tether.WifiTetherPreferenceController;
+import com.android.settingslib.RestrictedSwitchPreference;
 import com.android.settingslib.TetherUtil;
 import com.android.settingslib.search.SearchIndexable;
 
@@ -82,7 +86,7 @@
 
     private static final String TAG = "TetheringSettings";
 
-    private SwitchPreference mUsbTether;
+    private RestrictedSwitchPreference mUsbTether;
 
     private SwitchPreference mBluetoothTether;
 
@@ -208,7 +212,7 @@
 
     @VisibleForTesting
     void setupTetherPreference() {
-        mUsbTether = (SwitchPreference) findPreference(KEY_USB_TETHER_SETTINGS);
+        mUsbTether = (RestrictedSwitchPreference) findPreference(KEY_USB_TETHER_SETTINGS);
         mBluetoothTether = (SwitchPreference) findPreference(KEY_ENABLE_BLUETOOTH_TETHERING);
         mEthernetTether = (SwitchPreference) findPreference(KEY_ENABLE_ETHERNET_TETHERING);
     }
@@ -415,6 +419,8 @@
             mUsbTether.setEnabled(false);
             mUsbTether.setChecked(false);
         }
+        mUsbTether.setDisabledByAdmin(
+                checkIfUsbDataSignalingIsDisabled(getContext(), UserHandle.myUserId()));
     }
 
     @VisibleForTesting
diff --git a/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdater.java b/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdater.java
index 6a13a96..dc73ee0 100644
--- a/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdater.java
+++ b/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdater.java
@@ -19,8 +19,11 @@
 import static android.hardware.usb.UsbPortStatus.POWER_ROLE_SINK;
 import static android.hardware.usb.UsbPortStatus.POWER_ROLE_SOURCE;
 
+import static com.android.settingslib.RestrictedLockUtilsInternal.checkIfUsbDataSignalingIsDisabled;
+
 import android.content.Context;
 import android.hardware.usb.UsbManager;
+import android.os.UserHandle;
 
 import androidx.annotation.VisibleForTesting;
 import androidx.preference.Preference;
@@ -30,6 +33,7 @@
 import com.android.settings.core.SubSettingLauncher;
 import com.android.settings.dashboard.DashboardFragment;
 import com.android.settings.overlay.FeatureFactory;
+import com.android.settingslib.RestrictedPreference;
 import com.android.settingslib.core.instrumentation.MetricsFeatureProvider;
 
 /**
@@ -45,7 +49,7 @@
     private UsbBackend mUsbBackend;
     private DevicePreferenceCallback mDevicePreferenceCallback;
     @VisibleForTesting
-    Preference mUsbPreference;
+    RestrictedPreference mUsbPreference;
     @VisibleForTesting
     UsbConnectionBroadcastReceiver mUsbReceiver;
 
@@ -88,10 +92,12 @@
     }
 
     public void initUsbPreference(Context context) {
-        mUsbPreference = new Preference(context, null /* AttributeSet */);
+        mUsbPreference = new RestrictedPreference(context, null /* AttributeSet */);
         mUsbPreference.setTitle(R.string.usb_pref);
         mUsbPreference.setIcon(R.drawable.ic_usb);
         mUsbPreference.setKey(PREF_KEY);
+        mUsbPreference.setDisabledByAdmin(
+                checkIfUsbDataSignalingIsDisabled(context, UserHandle.myUserId()));
         mUsbPreference.setOnPreferenceClickListener((Preference p) -> {
             mMetricsFeatureProvider.logClickedPreference(p, mFragment.getMetricsCategory());
             // New version - uses a separate screen.
diff --git a/src/com/android/settings/connecteddevice/usb/UsbDetailsFragment.java b/src/com/android/settings/connecteddevice/usb/UsbDetailsFragment.java
index bc76568..8d14cf9 100644
--- a/src/com/android/settings/connecteddevice/usb/UsbDetailsFragment.java
+++ b/src/com/android/settings/connecteddevice/usb/UsbDetailsFragment.java
@@ -16,9 +16,12 @@
 
 package com.android.settings.connecteddevice.usb;
 
+import static com.android.settingslib.RestrictedLockUtilsInternal.checkIfUsbDataSignalingIsDisabled;
+
 import android.app.settings.SettingsEnums;
 import android.content.Context;
 import android.os.Bundle;
+import android.os.UserHandle;
 import android.view.View;
 
 import androidx.annotation.VisibleForTesting;
@@ -101,6 +104,11 @@
      */
     public static final BaseSearchIndexProvider SEARCH_INDEX_DATA_PROVIDER =
             new BaseSearchIndexProvider(R.xml.usb_details_fragment) {
+                @Override
+                protected boolean isPageSearchEnabled(Context context) {
+                    return checkIfUsbDataSignalingIsDisabled(
+                            context, UserHandle.myUserId()) != null;
+                }
 
                 @Override
                 public List<AbstractPreferenceController> createPreferenceControllers(
diff --git a/src/com/android/settings/development/DefaultUsbConfigurationPreferenceController.java b/src/com/android/settings/development/DefaultUsbConfigurationPreferenceController.java
new file mode 100644
index 0000000..905c552
--- /dev/null
+++ b/src/com/android/settings/development/DefaultUsbConfigurationPreferenceController.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.development;
+
+import static com.android.settingslib.RestrictedLockUtilsInternal.checkIfUsbDataSignalingIsDisabled;
+
+import android.content.Context;
+import android.os.UserHandle;
+
+import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
+
+import com.android.settingslib.RestrictedSwitchPreference;
+import com.android.settingslib.development.DeveloperOptionsPreferenceController;
+
+public class DefaultUsbConfigurationPreferenceController extends
+        DeveloperOptionsPreferenceController {
+
+    private static final String PREFERENCE_KEY = "default_usb_configuration";
+
+    private RestrictedSwitchPreference mPreference;
+
+    public DefaultUsbConfigurationPreferenceController(Context context) {
+        super(context);
+    }
+
+    @Override
+    public String getPreferenceKey() {
+        return PREFERENCE_KEY;
+    }
+
+    @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        mPreference = screen.findPreference(getPreferenceKey());
+    }
+
+    @Override
+    public void updateState(Preference preference) {
+        mPreference.setDisabledByAdmin(
+                checkIfUsbDataSignalingIsDisabled(mContext, UserHandle.myUserId()));
+    }
+}
diff --git a/src/com/android/settings/development/DevelopmentSettingsDashboardFragment.java b/src/com/android/settings/development/DevelopmentSettingsDashboardFragment.java
index e9dd9fd..367e7c5 100644
--- a/src/com/android/settings/development/DevelopmentSettingsDashboardFragment.java
+++ b/src/com/android/settings/development/DevelopmentSettingsDashboardFragment.java
@@ -560,8 +560,7 @@
         controllers.add(new DefaultLaunchPreferenceController(context, "demo_mode"));
         controllers.add(new DefaultLaunchPreferenceController(context, "quick_settings_tiles"));
         controllers.add(new DefaultLaunchPreferenceController(context, "feature_flags_dashboard"));
-        controllers.add(
-                new DefaultLaunchPreferenceController(context, "default_usb_configuration"));
+        controllers.add(new DefaultUsbConfigurationPreferenceController(context));
         controllers.add(new DefaultLaunchPreferenceController(context, "density"));
         controllers.add(new DefaultLaunchPreferenceController(context, "background_check"));
         controllers.add(new DefaultLaunchPreferenceController(context, "inactive_apps"));
diff --git a/src/com/android/settings/development/UsbAudioRoutingPreferenceController.java b/src/com/android/settings/development/UsbAudioRoutingPreferenceController.java
index 335a48d..8aa4f3c 100644
--- a/src/com/android/settings/development/UsbAudioRoutingPreferenceController.java
+++ b/src/com/android/settings/development/UsbAudioRoutingPreferenceController.java
@@ -16,14 +16,19 @@
 
 package com.android.settings.development;
 
+import static com.android.settingslib.RestrictedLockUtilsInternal.checkIfUsbDataSignalingIsDisabled;
+
 import android.content.Context;
+import android.os.UserHandle;
 import android.provider.Settings;
 
 import androidx.annotation.VisibleForTesting;
 import androidx.preference.Preference;
+import androidx.preference.PreferenceScreen;
 import androidx.preference.SwitchPreference;
 
 import com.android.settings.core.PreferenceControllerMixin;
+import com.android.settingslib.RestrictedSwitchPreference;
 import com.android.settingslib.development.DeveloperOptionsPreferenceController;
 
 public class UsbAudioRoutingPreferenceController extends DeveloperOptionsPreferenceController
@@ -36,6 +41,8 @@
     @VisibleForTesting
     static final int SETTING_VALUE_OFF = 0;
 
+    private RestrictedSwitchPreference mPreference;
+
     public UsbAudioRoutingPreferenceController(Context context) {
         super(context);
     }
@@ -46,6 +53,12 @@
     }
 
     @Override
+    public void displayPreference(PreferenceScreen screen) {
+        super.displayPreference(screen);
+        mPreference = screen.findPreference(getPreferenceKey());
+    }
+
+    @Override
     public boolean onPreferenceChange(Preference preference, Object newValue) {
         final boolean isEnabled = (Boolean) newValue;
         Settings.Secure.putInt(mContext.getContentResolver(),
@@ -58,7 +71,9 @@
     public void updateState(Preference preference) {
         final int usbAudioRoutingMode = Settings.Secure.getInt(mContext.getContentResolver(),
                 Settings.Secure.USB_AUDIO_AUTOMATIC_ROUTING_DISABLED, SETTING_VALUE_OFF);
-        ((SwitchPreference) mPreference).setChecked(usbAudioRoutingMode != SETTING_VALUE_OFF);
+        mPreference.setChecked(usbAudioRoutingMode != SETTING_VALUE_OFF);
+        mPreference.setDisabledByAdmin(
+                checkIfUsbDataSignalingIsDisabled(mContext, UserHandle.myUserId()));
     }
 
     @Override
diff --git a/tests/robotests/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdaterTest.java b/tests/robotests/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdaterTest.java
index 7c6a305..5c66db7 100644
--- a/tests/robotests/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdaterTest.java
+++ b/tests/robotests/src/com/android/settings/connecteddevice/usb/ConnectedUsbDeviceUpdaterTest.java
@@ -22,11 +22,17 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.admin.DevicePolicyManager;
 import android.content.Context;
 import android.hardware.usb.UsbManager;
+import android.os.UserHandle;
 
 import com.android.settings.R;
 import com.android.settings.connecteddevice.DevicePreferenceCallback;
@@ -55,21 +61,30 @@
     private DevicePreferenceCallback mDevicePreferenceCallback;
     @Mock
     private UsbBackend mUsbBackend;
+    @Mock
+    private DevicePolicyManager mDevicePolicyManager;
+
 
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
-        mContext = RuntimeEnvironment.application;
+        mContext = spy(RuntimeEnvironment.application);
         when(mFragment.getContext()).thenReturn(mContext);
         mDeviceUpdater =
                 new ConnectedUsbDeviceUpdater(mContext, mFragment, mDevicePreferenceCallback,
                         mUsbBackend);
         mDeviceUpdater.mUsbReceiver = mUsbReceiver;
+        when(mContext.getSystemService(DevicePolicyManager.class)).thenReturn(mDevicePolicyManager);
+        doReturn(mContext).when(mContext).createPackageContextAsUser(
+                any(String.class), anyInt(), any(UserHandle.class));
     }
 
     @Test
     public void initUsbPreference_preferenceInit() {
+        when(mDevicePolicyManager.isUsbDataSignalingEnabledForUser(
+                UserHandle.myUserId())).thenReturn(true);
+
         mDeviceUpdater.initUsbPreference(mContext);
 
         assertThat(mDeviceUpdater.mUsbPreference.getTitle()).isEqualTo("USB");
diff --git a/tests/robotests/src/com/android/settings/development/AdbPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/AdbPreferenceControllerTest.java
index 2e63f5a..ba9a7aa 100644
--- a/tests/robotests/src/com/android/settings/development/AdbPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/AdbPreferenceControllerTest.java
@@ -18,16 +18,21 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.admin.DevicePolicyManager;
 import android.content.Context;
+import android.os.UserHandle;
 import android.provider.Settings;
 
 import androidx.preference.PreferenceScreen;
-import androidx.preference.SwitchPreference;
+
+import com.android.settingslib.RestrictedSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -41,24 +46,29 @@
 public class AdbPreferenceControllerTest {
 
     @Mock
-    private SwitchPreference mPreference;
+    private RestrictedSwitchPreference mPreference;
     @Mock
     private PreferenceScreen mPreferenceScreen;
     @Mock
     private DevelopmentSettingsDashboardFragment mFragment;
+    @Mock
+    private DevicePolicyManager mDevicePolicyManager;
 
     private Context mContext;
     private AdbPreferenceController mController;
 
     @Before
-    public void setup() {
+    public void setup() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = spy(RuntimeEnvironment.application);
         mController = spy(new AdbPreferenceController(mContext, mFragment));
         doReturn(true).when(mController).isAvailable();
         when(mPreferenceScreen.findPreference(mController.getPreferenceKey()))
             .thenReturn(mPreference);
         mController.displayPreference(mPreferenceScreen);
+        when(mContext.getSystemService(DevicePolicyManager.class)).thenReturn(mDevicePolicyManager);
+        doReturn(mContext).when(mContext).createPackageContextAsUser(
+                any(String.class), anyInt(), any(UserHandle.class));
     }
 
     @Test
diff --git a/tests/robotests/src/com/android/settings/development/DefaultUsbConfigurationPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/DefaultUsbConfigurationPreferenceControllerTest.java
new file mode 100644
index 0000000..15f4eb9
--- /dev/null
+++ b/tests/robotests/src/com/android/settings/development/DefaultUsbConfigurationPreferenceControllerTest.java
@@ -0,0 +1,94 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.settings.development;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.admin.DevicePolicyManager;
+import android.content.ComponentName;
+import android.content.Context;
+import android.os.UserHandle;
+
+import androidx.preference.PreferenceScreen;
+
+import com.android.settingslib.RestrictedLockUtils;
+import com.android.settingslib.RestrictedSwitchPreference;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RobolectricTestRunner;
+import org.robolectric.RuntimeEnvironment;
+
+@RunWith(RobolectricTestRunner.class)
+public class DefaultUsbConfigurationPreferenceControllerTest {
+
+    private static final ComponentName TEST_COMPONENT_NAME = new ComponentName("test", "test");
+
+    @Mock
+    private RestrictedSwitchPreference mPreference;
+    @Mock
+    private PreferenceScreen mPreferenceScreen;
+    @Mock
+    private DevicePolicyManager mDevicePolicyManager;
+
+    private DefaultUsbConfigurationPreferenceController mController;
+
+    @Before
+    public void setup() throws Exception {
+        MockitoAnnotations.initMocks(this);
+        Context context = spy(RuntimeEnvironment.application);
+        mController = new DefaultUsbConfigurationPreferenceController(context);
+        when(mPreferenceScreen.findPreference(mController.getPreferenceKey()))
+                .thenReturn(mPreference);
+        mController.displayPreference(mPreferenceScreen);
+        when(context.getSystemService(DevicePolicyManager.class)).thenReturn(mDevicePolicyManager);
+        doReturn(context).when(context).createPackageContextAsUser(
+                any(String.class), anyInt(), any(UserHandle.class));
+    }
+
+    @Test
+    public void updateState_usbDataSignalingEnabled_shouldNotDisablePreference() {
+        when(mDevicePolicyManager.isUsbDataSignalingEnabledForUser(
+                UserHandle.myUserId())).thenReturn(true);
+        when(mDevicePolicyManager.getProfileOwner()).thenReturn(TEST_COMPONENT_NAME);
+
+        mController.updateState(mPreference);
+
+        verify(mPreference).setDisabledByAdmin(null);
+    }
+
+    @Test
+    public void updateState_usbDataSignalingDisabled_shouldDisablePreference() {
+        when(mDevicePolicyManager.isUsbDataSignalingEnabledForUser(
+                UserHandle.myUserId())).thenReturn(false);
+        when(mDevicePolicyManager.getProfileOwner()).thenReturn(TEST_COMPONENT_NAME);
+
+        mController.updateState(mPreference);
+
+        verify(mPreference).setDisabledByAdmin(eq(new RestrictedLockUtils.EnforcedAdmin(
+                TEST_COMPONENT_NAME, null, UserHandle.SYSTEM)));
+    }
+}
diff --git a/tests/robotests/src/com/android/settings/development/UsbAudioRoutingPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/development/UsbAudioRoutingPreferenceControllerTest.java
index affae70..1d45c1b 100644
--- a/tests/robotests/src/com/android/settings/development/UsbAudioRoutingPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/development/UsbAudioRoutingPreferenceControllerTest.java
@@ -18,14 +18,24 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.admin.DevicePolicyManager;
+import android.content.ComponentName;
 import android.content.Context;
+import android.os.UserHandle;
 import android.provider.Settings;
 
 import androidx.preference.PreferenceScreen;
-import androidx.preference.SwitchPreference;
+
+import com.android.settingslib.RestrictedLockUtils;
+import com.android.settingslib.RestrictedSwitchPreference;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -38,26 +48,36 @@
 @RunWith(RobolectricTestRunner.class)
 public class UsbAudioRoutingPreferenceControllerTest {
 
+    private static final ComponentName TEST_COMPONENT_NAME = new ComponentName("test", "test");
+
     @Mock
     private PreferenceScreen mScreen;
     @Mock
-    private SwitchPreference mPreference;
+    private RestrictedSwitchPreference mPreference;
+    @Mock
+    private DevicePolicyManager mDevicePolicyManager;
 
     private Context mContext;
 
     private UsbAudioRoutingPreferenceController mController;
 
     @Before
-    public void setUp() {
+    public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
-        mContext = RuntimeEnvironment.application;
+        mContext = spy(RuntimeEnvironment.application);
         mController = new UsbAudioRoutingPreferenceController(mContext);
         when(mScreen.findPreference(mController.getPreferenceKey())).thenReturn(mPreference);
         mController.displayPreference(mScreen);
+        when(mContext.getSystemService(DevicePolicyManager.class)).thenReturn(mDevicePolicyManager);
+        doReturn(mContext).when(mContext).createPackageContextAsUser(
+                any(String.class), anyInt(), any(UserHandle.class));
     }
 
     @Test
     public void updateState_usbAudioRoutingEnabled_shouldCheckedPreference() {
+        when(mDevicePolicyManager.isUsbDataSignalingEnabledForUser(
+                UserHandle.myUserId())).thenReturn(true);
+        when(mDevicePolicyManager.getProfileOwner()).thenReturn(TEST_COMPONENT_NAME);
         Settings.Secure.putInt(mContext.getContentResolver(),
                 Settings.Secure.USB_AUDIO_AUTOMATIC_ROUTING_DISABLED,
                 UsbAudioRoutingPreferenceController.SETTING_VALUE_ON);
@@ -69,6 +89,9 @@
 
     @Test
     public void updateState_usbAudioRoutingDisabled_shouldUncheckedPreference() {
+        when(mDevicePolicyManager.isUsbDataSignalingEnabledForUser(
+                UserHandle.myUserId())).thenReturn(true);
+        when(mDevicePolicyManager.getProfileOwner()).thenReturn(TEST_COMPONENT_NAME);
         Settings.Secure.putInt(mContext.getContentResolver(),
                 Settings.Secure.USB_AUDIO_AUTOMATIC_ROUTING_DISABLED,
                 UsbAudioRoutingPreferenceController.SETTING_VALUE_OFF);
@@ -79,6 +102,18 @@
     }
 
     @Test
+    public void updateState_usbDataSignalingDisabled_shouldDisablePreference() {
+        when(mDevicePolicyManager.isUsbDataSignalingEnabledForUser(
+                UserHandle.myUserId())).thenReturn(false);
+        when(mDevicePolicyManager.getProfileOwner()).thenReturn(TEST_COMPONENT_NAME);
+
+        mController.updateState(mPreference);
+
+        verify(mPreference).setDisabledByAdmin(eq(new RestrictedLockUtils.EnforcedAdmin(
+                TEST_COMPONENT_NAME, null, UserHandle.SYSTEM)));
+    }
+
+    @Test
     public void onPreferenceChange_preferenceChecked_shouldEnableUsbAudioRouting() {
         mController.onPreferenceChange(mPreference, true /* new value */);