diff --git a/src/com/android/settings/sound/AudioSwitchPreferenceController.java b/src/com/android/settings/sound/AudioSwitchPreferenceController.java
index 25a0518..28ad3f5 100644
--- a/src/com/android/settings/sound/AudioSwitchPreferenceController.java
+++ b/src/com/android/settings/sound/AudioSwitchPreferenceController.java
@@ -193,8 +193,12 @@
     public void onDeviceBondStateChanged(CachedBluetoothDevice cachedDevice, int bondState) {
     }
 
+    protected boolean isStreamFromOutputDevice(int streamType, int device) {
+        return mAudioManager.getDevicesForStream(streamType) == device;
+    }
+
     protected boolean isOngoingCallStatus() {
-        int audioMode = mAudioManager.getMode();
+        final int audioMode = mAudioManager.getMode();
         return audioMode == AudioManager.MODE_RINGTONE
                 || audioMode == AudioManager.MODE_IN_CALL
                 || audioMode == AudioManager.MODE_IN_COMMUNICATION;
diff --git a/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceController.java b/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceController.java
index b0b3dc5..2f21f1b 100644
--- a/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceController.java
+++ b/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceController.java
@@ -16,6 +16,9 @@
 
 package com.android.settings.sound;
 
+import static android.media.AudioManager.STREAM_VOICE_CALL;
+import static android.media.AudioSystem.DEVICE_OUT_USB_HEADSET;
+
 import android.bluetooth.BluetoothDevice;
 import android.content.Context;
 import android.support.v7.preference.Preference;
@@ -76,7 +79,7 @@
         // Setup devices entries, select active connected device
         setupPreferenceEntries(mediaOutputs, mediaValues, activeDevice);
 
-        if (mAudioManager.isWiredHeadsetOn() && !mAudioManager.isBluetoothScoOn()) {
+        if (isStreamFromOutputDevice(STREAM_VOICE_CALL, DEVICE_OUT_USB_HEADSET)) {
             // If wired headset is plugged in and active, select to default device.
             mSelectedIndex = getDefaultDeviceIndex();
         }
diff --git a/src/com/android/settings/sound/MediaOutputPreferenceController.java b/src/com/android/settings/sound/MediaOutputPreferenceController.java
index 2e52f77..74c77cd 100644
--- a/src/com/android/settings/sound/MediaOutputPreferenceController.java
+++ b/src/com/android/settings/sound/MediaOutputPreferenceController.java
@@ -16,6 +16,8 @@
 
 package com.android.settings.sound;
 
+import static android.media.AudioManager.STREAM_VOICE_CALL;
+import static android.media.AudioSystem.DEVICE_OUT_USB_HEADSET;
 import static android.media.MediaRouter.ROUTE_TYPE_REMOTE_DISPLAY;
 
 import android.bluetooth.BluetoothDevice;
@@ -91,7 +93,7 @@
         // Setup devices entries, select active connected device
         setupPreferenceEntries(mediaOutputs, mediaValues, activeDevice);
 
-        if (mAudioManager.isWiredHeadsetOn() && !mAudioManager.isBluetoothA2dpOn()) {
+        if (isStreamFromOutputDevice(STREAM_VOICE_CALL, DEVICE_OUT_USB_HEADSET)) {
             // If wired headset is plugged in and active, select to default device.
             mSelectedIndex = getDefaultDeviceIndex();
         }
diff --git a/tests/robotests/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceControllerTest.java
index c6c4b45..f3f1c83 100644
--- a/tests/robotests/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/sound/HandsFreeProfileOutputPreferenceControllerTest.java
@@ -17,6 +17,8 @@
 package com.android.settings.sound;
 
 
+import static android.media.AudioSystem.DEVICE_OUT_USB_HEADSET;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.mock;
@@ -205,10 +207,9 @@
      * Preference summary should be "This device"
      */
     @Test
-    public void hapBtDevicesAreAvailableButWiredHeadsetIsActivated_shouldSetDefaultSummary() {
+    public void updateState_withAvailableDevicesWiredHeadsetActivated_shouldSetDefaultSummary() {
         mShadowAudioManager.setMode(AudioManager.MODE_IN_COMMUNICATION);
-        mShadowAudioManager.setWiredHeadsetOn(true);
-        mShadowAudioManager.setBluetoothScoOn(false);
+        mShadowAudioManager.setStream(DEVICE_OUT_USB_HEADSET);
         when(mHeadsetProfile.getConnectedDevices()).thenReturn(mConnectedDevices);
         when(mHeadsetProfile.getActiveDevice()).thenReturn(
                 mBluetoothDevice); // BT device is still activated in this case
@@ -226,7 +227,7 @@
      * Preference summary should be "This device"
      */
     @Test
-    public void noAvailableHeadsetBtDevices_preferenceEnableIsFalse_shouldSetDefaultSummary() {
+    public void updateState_noAvailableHeadsetBtDevices_shouldSetDefaultSummary() {
         mShadowAudioManager.setMode(AudioManager.MODE_IN_COMMUNICATION);
         List<BluetoothDevice> emptyDeviceList = new ArrayList<>();
         when(mHeadsetProfile.getConnectedDevices()).thenReturn(emptyDeviceList);
diff --git a/tests/robotests/src/com/android/settings/sound/MediaOutputPreferenceControllerTest.java b/tests/robotests/src/com/android/settings/sound/MediaOutputPreferenceControllerTest.java
index 2b15b8e..4da5f27 100644
--- a/tests/robotests/src/com/android/settings/sound/MediaOutputPreferenceControllerTest.java
+++ b/tests/robotests/src/com/android/settings/sound/MediaOutputPreferenceControllerTest.java
@@ -17,6 +17,8 @@
 package com.android.settings.sound;
 
 
+import static android.media.AudioSystem.DEVICE_OUT_USB_HEADSET;
+
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.Mockito.mock;
@@ -256,8 +258,7 @@
     @Test
     public void updateState_a2dpDevicesAvailableWiredHeadsetIsActivated_shouldSetDefaultSummary() {
         mShadowAudioManager.setMode(AudioManager.MODE_NORMAL);
-        mShadowAudioManager.setWiredHeadsetOn(true);
-        mShadowAudioManager.setBluetoothA2dpOn(false);
+        mShadowAudioManager.setStream(DEVICE_OUT_USB_HEADSET);
         when(mA2dpProfile.getConnectedDevices()).thenReturn(mConnectedDevices);
         when(mA2dpProfile.getActiveDevice()).thenReturn(
                 mBluetoothDevice); // BT device is still activated in this case
diff --git a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowAudioManager.java b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowAudioManager.java
index 6817648..0de2156 100644
--- a/tests/robotests/src/com/android/settings/testutils/shadow/ShadowAudioManager.java
+++ b/tests/robotests/src/com/android/settings/testutils/shadow/ShadowAudioManager.java
@@ -16,6 +16,15 @@
 
 package com.android.settings.testutils.shadow;
 
+import static android.media.AudioManager.STREAM_ACCESSIBILITY;
+import static android.media.AudioManager.STREAM_ALARM;
+import static android.media.AudioManager.STREAM_MUSIC;
+import static android.media.AudioManager.STREAM_NOTIFICATION;
+import static android.media.AudioManager.STREAM_RING;
+import static android.media.AudioManager.STREAM_SYSTEM;
+import static android.media.AudioManager.STREAM_VOICE_CALL;
+import static android.media.AudioManager.STREAM_DTMF;
+
 import static org.robolectric.RuntimeEnvironment.application;
 
 import android.media.AudioDeviceCallback;
@@ -32,6 +41,7 @@
 @Implements(value = AudioManager.class, inheritImplementationMethods = true)
 public class ShadowAudioManager extends org.robolectric.shadows.ShadowAudioManager {
     private int mRingerMode;
+    private int mStream;
     private boolean mMusicActiveRemotely = false;
     private ArrayList<AudioDeviceCallback> mDeviceCallbacks = new ArrayList();
 
@@ -48,10 +58,12 @@
         mRingerMode = mode;
     }
 
+    @Implementation
     public void registerAudioDeviceCallback(AudioDeviceCallback callback, Handler handler) {
         mDeviceCallbacks.add(callback);
     }
 
+    @Implementation
     public void unregisterAudioDeviceCallback(AudioDeviceCallback callback) {
         if (mDeviceCallbacks.contains(callback)) {
             mDeviceCallbacks.remove(callback);
@@ -62,10 +74,32 @@
         mMusicActiveRemotely = flag;
     }
 
+    @Implementation
     public boolean isMusicActiveRemotely() {
         return mMusicActiveRemotely;
     }
 
+    public void setStream(int stream) {
+        mStream = stream;
+    }
+
+    @Implementation
+    public int getDevicesForStream(int streamType) {
+        switch (streamType) {
+            case STREAM_VOICE_CALL:
+            case STREAM_SYSTEM:
+            case STREAM_RING:
+            case STREAM_MUSIC:
+            case STREAM_ALARM:
+            case STREAM_NOTIFICATION:
+            case STREAM_DTMF:
+            case STREAM_ACCESSIBILITY:
+                return mStream;
+            default:
+                return 0;
+        }
+    }
+
     @Resetter
     public void reset() {
         mDeviceCallbacks.clear();