Add flag for communication device refactor.

Add missing flag dependency for the call audio communication device
refactor work that was introduced as part of U. This CL also ensures
that the device set for communication is cleared when the audio mode is
set back to MODE_NORMAL (after call ends) and that the audio lost signal
is sent after the device has been cleared for communication. The latter
addresses issues with switching from a bluetooth LE audio device to
speaker to ensure that the device isn't considered an active device
still when switching audio routes.

Bug: 308968392
Test: atest TelecomUnitTests
Test: Manual to ensure that device set for communication is cleared
after a call ends.

Change-Id: I62a0d65c53b67a307b23d1015b994796e955a6aa

15 files changed, 273 insertions, 112 deletions

diff --git a/flags/telecom_callaudioroutestatemachine_flags.aconfig b/flags/telecom_callaudioroutestatemachine_flags.aconfig
index 7ff144063..a3294c0b2 100644
--- a/flags/telecom_callaudioroutestatemachine_flags.aconfig
+++ b/flags/telecom_callaudioroutestatemachine_flags.aconfig
@@ -26,4 +26,11 @@ flag {
   namespace: "telecom"
   description: "Fix audio route transition issue on call disconnection when bt audio connected."
   bug: "306113816"
+}
+
+flag {
+  name: "call_audio_communication_device_refactor"
+  namespace: "telecom"
+  description: "Refactor call audio set/clear communication device and include unsupported routes."
+  bug: "308968392"
 }
\ No newline at end of file
diff --git a/src/com/android/server/telecom/CallAudioCommunicationDeviceTracker.java b/src/com/android/server/telecom/CallAudioCommunicationDeviceTracker.java
index 43624a3c3..38a96eb3e 100644
--- a/src/com/android/server/telecom/CallAudioCommunicationDeviceTracker.java
+++ b/src/com/android/server/telecom/CallAudioCommunicationDeviceTracker.java
@@ -63,6 +63,10 @@ public class CallAudioCommunicationDeviceTracker {
         return mAudioDeviceType == audioDeviceType;
     }
 
+    public int getCurrentLocallyRequestedCommunicationDevice() {
+       return mAudioDeviceType;
+    }
+
     @VisibleForTesting
     public void setTestCommunicationDevice(int audioDeviceType) {
         mAudioDeviceType = audioDeviceType;
@@ -177,12 +181,6 @@ public class CallAudioCommunicationDeviceTracker {
             return;
         }
 
-        if (isBtDevice && mBtAudioDevice != null) {
-            // Signal that BT audio was lost for device.
-            mBluetoothRouteManager.onAudioLost(mBtAudioDevice);
-            mBtAudioDevice = null;
-        }
-
         if (mAudioManager == null) {
             Log.i(this, "clearCommunicationDevice: mAudioManager is null");
             return;
@@ -191,6 +189,12 @@ public class CallAudioCommunicationDeviceTracker {
         // Clear device and reset locally saved device type.
         mAudioManager.clearCommunicationDevice();
         mAudioDeviceType = sAUDIO_DEVICE_TYPE_INVALID;
+
+        if (isBtDevice && mBtAudioDevice != null) {
+            // Signal that BT audio was lost for device.
+            mBluetoothRouteManager.onAudioLost(mBtAudioDevice);
+            mBtAudioDevice = null;
+        }
     }
 
     private boolean isUsbHeadsetType(int audioDeviceType, int sourceType) {
diff --git a/src/com/android/server/telecom/CallAudioModeStateMachine.java b/src/com/android/server/telecom/CallAudioModeStateMachine.java
index 90da8ed54..f8c199790 100644
--- a/src/com/android/server/telecom/CallAudioModeStateMachine.java
+++ b/src/com/android/server/telecom/CallAudioModeStateMachine.java
@@ -41,8 +41,10 @@ public class CallAudioModeStateMachine extends StateMachine {
     private LocalLog mLocalLog = new LocalLog(20);
     public static class Factory {
         public CallAudioModeStateMachine create(SystemStateHelper systemStateHelper,
-                                                AudioManager am, FeatureFlags featureFlags) {
-            return new CallAudioModeStateMachine(systemStateHelper, am, featureFlags);
+                AudioManager am, FeatureFlags featureFlags,
+                CallAudioCommunicationDeviceTracker callAudioCommunicationDeviceTracker) {
+            return new CallAudioModeStateMachine(systemStateHelper, am,
+                    featureFlags, callAudioCommunicationDeviceTracker);
         }
     }
 
@@ -277,6 +279,11 @@ public class CallAudioModeStateMachine extends StateMachine {
             Log.i(LOG_TAG, "Audio focus entering UNFOCUSED state");
             mLocalLog.log("Enter UNFOCUSED");
             if (mIsInitialized) {
+                // Clear any communication device that was requested previously.
+                if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                    mCommunicationDeviceTracker.clearCommunicationDevice(mCommunicationDeviceTracker
+                            .getCurrentLocallyRequestedCommunicationDevice());
+                }
                 if (mFeatureFlags.setAudioModeBeforeAbandonFocus()) {
                     mAudioManager.setMode(AudioManager.MODE_NORMAL);
                     mCallAudioManager.setCallAudioRouteFocusState(
@@ -878,17 +885,20 @@ public class CallAudioModeStateMachine extends StateMachine {
     private final SystemStateHelper mSystemStateHelper;
     private CallAudioManager mCallAudioManager;
     private FeatureFlags mFeatureFlags;
+    private CallAudioCommunicationDeviceTracker mCommunicationDeviceTracker;
 
     private int mMostRecentMode;
     private boolean mIsInitialized = false;
 
     public CallAudioModeStateMachine(SystemStateHelper systemStateHelper,
-            AudioManager audioManager, FeatureFlags featureFlags) {
+            AudioManager audioManager, FeatureFlags featureFlags,
+            CallAudioCommunicationDeviceTracker callAudioCommunicationDeviceTracker) {
         super(CallAudioModeStateMachine.class.getSimpleName());
         mAudioManager = audioManager;
         mSystemStateHelper = systemStateHelper;
         mMostRecentMode = AudioManager.MODE_NORMAL;
         mFeatureFlags = featureFlags;
+        mCommunicationDeviceTracker = callAudioCommunicationDeviceTracker;
 
         createStates();
     }
@@ -897,12 +907,14 @@ public class CallAudioModeStateMachine extends StateMachine {
      * Used for testing
      */
     public CallAudioModeStateMachine(SystemStateHelper systemStateHelper,
-            AudioManager audioManager, Looper looper, FeatureFlags featureFlags) {
+            AudioManager audioManager, Looper looper, FeatureFlags featureFlags,
+            CallAudioCommunicationDeviceTracker communicationDeviceTracker) {
         super(CallAudioModeStateMachine.class.getSimpleName(), looper);
         mAudioManager = audioManager;
         mSystemStateHelper = systemStateHelper;
         mMostRecentMode = AudioManager.MODE_NORMAL;
         mFeatureFlags = featureFlags;
+        mCommunicationDeviceTracker = communicationDeviceTracker;
 
         createStates();
     }
diff --git a/src/com/android/server/telecom/CallAudioRouteStateMachine.java b/src/com/android/server/telecom/CallAudioRouteStateMachine.java
index f63463f6a..b89dbe297 100644
--- a/src/com/android/server/telecom/CallAudioRouteStateMachine.java
+++ b/src/com/android/server/telecom/CallAudioRouteStateMachine.java
@@ -377,8 +377,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
         public void enter() {
             super.enter();
             setSpeakerphoneOn(false);
-            mCommunicationDeviceTracker.setCommunicationDevice(
-                    AudioDeviceInfo.TYPE_BUILTIN_EARPIECE, null);
+            if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                mCommunicationDeviceTracker.setCommunicationDevice(
+                        AudioDeviceInfo.TYPE_BUILTIN_EARPIECE, null);
+            }
             CallAudioState newState = new CallAudioState(mIsMuted, ROUTE_EARPIECE,
                     mAvailableRoutes, null,
                     mBluetoothRouteManager.getConnectedDevices());
@@ -409,8 +411,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
                 case SWITCH_BLUETOOTH:
                 case USER_SWITCH_BLUETOOTH:
                     if ((mAvailableRoutes & ROUTE_BLUETOOTH) != 0) {
-                        mCommunicationDeviceTracker.clearCommunicationDevice(
-                                AudioDeviceInfo.TYPE_BUILTIN_EARPIECE);
+                        if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                            mCommunicationDeviceTracker.clearCommunicationDevice(
+                                    AudioDeviceInfo.TYPE_BUILTIN_EARPIECE);
+                        }
                         if (mAudioFocusType == ACTIVE_FOCUS
                                 || mBluetoothRouteManager.isInbandRingingEnabled()) {
                             String address = (msg.obj instanceof SomeArgs) ?
@@ -427,8 +431,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
                 case SWITCH_HEADSET:
                 case USER_SWITCH_HEADSET:
                     if ((mAvailableRoutes & ROUTE_WIRED_HEADSET) != 0) {
-                        mCommunicationDeviceTracker.clearCommunicationDevice(
-                                AudioDeviceInfo.TYPE_BUILTIN_EARPIECE);
+                        if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                            mCommunicationDeviceTracker.clearCommunicationDevice(
+                                    AudioDeviceInfo.TYPE_BUILTIN_EARPIECE);
+                        }
                         transitionTo(mActiveHeadsetRoute);
                     } else {
                         Log.w(this, "Ignoring switch to headset command. Not available.");
@@ -438,8 +444,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
                     // fall through; we want to switch to speaker mode when docked and in a call.
                 case SWITCH_SPEAKER:
                 case USER_SWITCH_SPEAKER:
-                    mCommunicationDeviceTracker.clearCommunicationDevice(
-                            AudioDeviceInfo.TYPE_BUILTIN_EARPIECE);
+                    if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                        mCommunicationDeviceTracker.clearCommunicationDevice(
+                                AudioDeviceInfo.TYPE_BUILTIN_EARPIECE);
+                    }
                     setSpeakerphoneOn(true);
                     // fall through
                 case SPEAKER_ON:
@@ -593,8 +601,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
         public void enter() {
             super.enter();
             setSpeakerphoneOn(false);
-            mCommunicationDeviceTracker.setCommunicationDevice(
-                    AudioDeviceInfo.TYPE_WIRED_HEADSET, null);
+            if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                mCommunicationDeviceTracker.setCommunicationDevice(
+                        AudioDeviceInfo.TYPE_WIRED_HEADSET, null);
+            }
             CallAudioState newState = new CallAudioState(mIsMuted, ROUTE_WIRED_HEADSET,
                     mAvailableRoutes, null, mBluetoothRouteManager.getConnectedDevices());
             setSystemAudioState(newState, true);
@@ -616,8 +626,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
                 case SWITCH_EARPIECE:
                 case USER_SWITCH_EARPIECE:
                     if ((mAvailableRoutes & ROUTE_EARPIECE) != 0) {
-                        mCommunicationDeviceTracker.clearCommunicationDevice(
-                                AudioDeviceInfo.TYPE_WIRED_HEADSET);
+                        if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                            mCommunicationDeviceTracker.clearCommunicationDevice(
+                                    AudioDeviceInfo.TYPE_WIRED_HEADSET);
+                        }
                         transitionTo(mActiveEarpieceRoute);
                     } else {
                         Log.w(this, "Ignoring switch to earpiece command. Not available.");
@@ -633,8 +645,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
                                 || mBluetoothRouteManager.isInbandRingingEnabled()) {
                             String address = (msg.obj instanceof SomeArgs) ?
                                     (String) ((SomeArgs) msg.obj).arg2 : null;
-                            mCommunicationDeviceTracker.clearCommunicationDevice(
-                                    AudioDeviceInfo.TYPE_WIRED_HEADSET);
+                            if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                                mCommunicationDeviceTracker.clearCommunicationDevice(
+                                        AudioDeviceInfo.TYPE_WIRED_HEADSET);
+                            }
                             // Omit transition to ActiveBluetoothRoute until actual connection.
                             setBluetoothOn(address);
                         } else {
@@ -651,8 +665,10 @@ public class CallAudioRouteStateMachine extends StateMachine {
                     return HANDLED;
                 case SWITCH_SPEAKER:
                 case USER_SWITCH_SPEAKER:
-                    mCommunicationDeviceTracker.clearCommunicationDevice(
-                            AudioDeviceInfo.TYPE_WIRED_HEADSET);
+                    if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                        mCommunicationDeviceTracker.clearCommunicationDevice(
+                                AudioDeviceInfo.TYPE_WIRED_HEADSET);
+                    }
                     setSpeakerphoneOn(true);
                     // fall through
                 case SPEAKER_ON:
@@ -820,7 +836,9 @@ public class CallAudioRouteStateMachine extends StateMachine {
             // the BT connection fails to be set. Previously, the logic was to setBluetoothOn in
             // ACTIVE_FOCUS but the route would still remain in a quiescent route, so instead we
             // should be transitioning directly into the active route.
-            setBluetoothOn(null);
+            if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                setBluetoothOn(null);
+            }
             CallAudioState newState = new CallAudioState(mIsMuted, ROUTE_BLUETOOTH,
                     mAvailableRoutes, mBluetoothRouteManager.getBluetoothAudioConnectedDevice(),
                     mBluetoothRouteManager.getConnectedDevices());
@@ -1100,7 +1118,11 @@ public class CallAudioRouteStateMachine extends StateMachine {
                     if (msg.arg1 == ACTIVE_FOCUS) {
                         // It is possible that the connection to BT will fail while in-call, in
                         // which case, we want to transition into the active route.
-                        transitionTo(mActiveBluetoothRoute);
+                        if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                            transitionTo(mActiveBluetoothRoute);
+                        } else {
+                            setBluetoothOn(null);
+                        }
                     } else if (msg.arg1 == RINGING_FOCUS) {
                         if (mBluetoothRouteManager.isInbandRingingEnabled()) {
                             setBluetoothOn(null);
@@ -1802,12 +1824,16 @@ public class CallAudioRouteStateMachine extends StateMachine {
         // These APIs are all via two-way binder calls so can potentially block Telecom.  Since none
         // of this has to happen in the Telecom lock we'll offload it to the async executor.
         boolean speakerOn = false;
-        if (on) {
-            speakerOn = mCommunicationDeviceTracker.setCommunicationDevice(
-                    AudioDeviceInfo.TYPE_BUILTIN_SPEAKER, null);
+        if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+            if (on) {
+                speakerOn = mCommunicationDeviceTracker.setCommunicationDevice(
+                        AudioDeviceInfo.TYPE_BUILTIN_SPEAKER, null);
+            } else {
+                mCommunicationDeviceTracker.clearCommunicationDevice(
+                        AudioDeviceInfo.TYPE_BUILTIN_SPEAKER);
+            }
         } else {
-            mCommunicationDeviceTracker.clearCommunicationDevice(
-                    AudioDeviceInfo.TYPE_BUILTIN_SPEAKER);
+            processLegacySpeakerCommunicationDevice(on);
         }
         mStatusBarNotifier.notifySpeakerphone(hasAnyCalls && speakerOn);
     }
@@ -2057,6 +2083,30 @@ public class CallAudioRouteStateMachine extends StateMachine {
         return containsWatchDevice && !containsNonWatchDevice && !isActiveDeviceWatch;
     }
 
+    private boolean processLegacySpeakerCommunicationDevice(boolean on) {
+        AudioDeviceInfo speakerDevice = null;
+        for (AudioDeviceInfo info : mAudioManager.getAvailableCommunicationDevices()) {
+            if (info.getType() == AudioDeviceInfo.TYPE_BUILTIN_SPEAKER) {
+                speakerDevice = info;
+                break;
+            }
+        }
+        boolean speakerOn = false;
+        if (speakerDevice != null && on) {
+            boolean result = mAudioManager.setCommunicationDevice(speakerDevice);
+            if (result) {
+                speakerOn = true;
+            }
+        } else {
+            AudioDeviceInfo curDevice = mAudioManager.getCommunicationDevice();
+            if (curDevice != null
+                    && curDevice.getType() == AudioDeviceInfo.TYPE_BUILTIN_SPEAKER) {
+                mAudioManager.clearCommunicationDevice();
+            }
+        }
+        return speakerOn;
+    }
+
     private int calculateBaselineRouteMessage(boolean isExplicitUserRequest,
             boolean includeBluetooth) {
         boolean isSkipEarpiece = false;
diff --git a/src/com/android/server/telecom/CallsManager.java b/src/com/android/server/telecom/CallsManager.java
index 1c2b015e7..6afa01d7e 100755
--- a/src/com/android/server/telecom/CallsManager.java
+++ b/src/com/android/server/telecom/CallsManager.java
@@ -656,7 +656,8 @@ public class CallsManager extends Call.ListenerBase
                 mTimeoutsAdapter, mLock);
         mCallAudioManager = new CallAudioManager(callAudioRouteStateMachine,
                 this, callAudioModeStateMachineFactory.create(systemStateHelper,
-                (AudioManager) mContext.getSystemService(Context.AUDIO_SERVICE), featureFlags),
+                (AudioManager) mContext.getSystemService(Context.AUDIO_SERVICE),
+                featureFlags, communicationDeviceTracker),
                 playerFactory, mRinger, new RingbackPlayer(playerFactory),
                 bluetoothStateReceiver, mDtmfLocalTonePlayer, featureFlags);
 
diff --git a/src/com/android/server/telecom/TelecomSystem.java b/src/com/android/server/telecom/TelecomSystem.java
index 379487d37..101cd2d09 100644
--- a/src/com/android/server/telecom/TelecomSystem.java
+++ b/src/com/android/server/telecom/TelecomSystem.java
@@ -259,12 +259,13 @@ public class TelecomSystem {
                     CallAudioCommunicationDeviceTracker(mContext);
             BluetoothDeviceManager bluetoothDeviceManager = new BluetoothDeviceManager(mContext,
                     mContext.getSystemService(BluetoothManager.class).getAdapter(),
-                    communicationDeviceTracker);
+                    communicationDeviceTracker, featureFlags);
             BluetoothRouteManager bluetoothRouteManager = new BluetoothRouteManager(mContext, mLock,
                     bluetoothDeviceManager, new Timeouts.Adapter(),
                     communicationDeviceTracker, featureFlags);
             BluetoothStateReceiver bluetoothStateReceiver = new BluetoothStateReceiver(
-                    bluetoothDeviceManager, bluetoothRouteManager, communicationDeviceTracker);
+                    bluetoothDeviceManager, bluetoothRouteManager,
+                    communicationDeviceTracker, featureFlags);
             mContext.registerReceiver(bluetoothStateReceiver, BluetoothStateReceiver.INTENT_FILTER);
             communicationDeviceTracker.setBluetoothRouteManager(bluetoothRouteManager);
 
diff --git a/src/com/android/server/telecom/bluetooth/BluetoothDeviceManager.java b/src/com/android/server/telecom/bluetooth/BluetoothDeviceManager.java
index e32f72c2d..c5e56b377 100644
--- a/src/com/android/server/telecom/bluetooth/BluetoothDeviceManager.java
+++ b/src/com/android/server/telecom/bluetooth/BluetoothDeviceManager.java
@@ -35,6 +35,7 @@ import android.util.LocalLog;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.util.IndentingPrintWriter;
 import com.android.server.telecom.CallAudioCommunicationDeviceTracker;
+import com.android.server.telecom.flags.FeatureFlags;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -211,9 +212,11 @@ public class BluetoothDeviceManager {
     private AudioManager mAudioManager;
     private Executor mExecutor;
     private CallAudioCommunicationDeviceTracker mCommunicationDeviceTracker;
+    private FeatureFlags mFeatureFlags;
 
     public BluetoothDeviceManager(Context context, BluetoothAdapter bluetoothAdapter,
-            CallAudioCommunicationDeviceTracker communicationDeviceTracker) {
+            CallAudioCommunicationDeviceTracker communicationDeviceTracker,
+            FeatureFlags featureFlags) {
         if (bluetoothAdapter != null) {
             mBluetoothAdapter = bluetoothAdapter;
             bluetoothAdapter.getProfileProxy(context, mBluetoothProfileServiceListener,
@@ -225,6 +228,7 @@ public class BluetoothDeviceManager {
             mAudioManager = context.getSystemService(AudioManager.class);
             mExecutor = context.getMainExecutor();
             mCommunicationDeviceTracker = communicationDeviceTracker;
+            mFeatureFlags = featureFlags;
         }
     }
 
@@ -449,7 +453,17 @@ public class BluetoothDeviceManager {
     }
 
     public void disconnectAudio() {
-        mCommunicationDeviceTracker.clearBtCommunicationDevice();
+        if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+            mCommunicationDeviceTracker.clearBtCommunicationDevice();
+            disconnectSco();
+        } else {
+            disconnectSco();
+            clearLeAudioOrSpeakerCommunicationDevice();
+            clearHearingAidOrSpeakerCommunicationDevice();
+        }
+    }
+
+    public void disconnectSco() {
         if (mBluetoothHeadset == null) {
             Log.w(this, "Trying to disconnect audio but no headset service exists.");
         } else {
@@ -481,9 +495,11 @@ public class BluetoothDeviceManager {
         AudioDeviceInfo audioDeviceInfo = mAudioManager.getCommunicationDevice();
         if (audioDeviceInfo != null) {
             if (audioDeviceInfo.getType() == AudioDeviceInfo.TYPE_BLE_HEADSET) {
+                Log.i(this, "clearLeAudioCommunicationDevice: clearing le audio");
                 mBluetoothRouteManager.onAudioLost(audioDeviceInfo.getAddress());
                 mAudioManager.clearCommunicationDevice();
             } else if (audioDeviceInfo.getType() == AudioDeviceInfo.TYPE_BUILTIN_SPEAKER) {
+                Log.i(this, "clearLeAudioCommunicationDevice: clearing speaker");
                 mAudioManager.clearCommunicationDevice();
             }
         }
@@ -505,10 +521,12 @@ public class BluetoothDeviceManager {
 
         AudioDeviceInfo audioDeviceInfo = mAudioManager.getCommunicationDevice();
         if (audioDeviceInfo != null) {
-            if (audioDeviceInfo.getType() == AudioDeviceInfo.TYPE_BLE_HEADSET) {
+            if (audioDeviceInfo.getType() == AudioDeviceInfo.TYPE_HEARING_AID) {
+                Log.i(this, "clearHearingAidCommunicationDevice: clearing hearing aid");
                 mBluetoothRouteManager.onAudioLost(audioDeviceInfo.getAddress());
                 mAudioManager.clearCommunicationDevice();
             } else if (audioDeviceInfo.getType() == AudioDeviceInfo.TYPE_BUILTIN_SPEAKER) {
+                Log.i(this, "clearHearingAidCommunicationDevice: clearing speaker");
                 mAudioManager.clearCommunicationDevice();
             }
         }
@@ -664,8 +682,10 @@ public class BluetoothDeviceManager {
                  * Only after receiving ACTION_ACTIVE_DEVICE_CHANGED it is known that device that
                  * will be audio switched to is available to be choose as communication device */
                 if (!switchingBtDevices) {
-                    return mCommunicationDeviceTracker.setCommunicationDevice(
-                            AudioDeviceInfo.TYPE_BLE_HEADSET, device);
+                    return mFeatureFlags.callAudioCommunicationDeviceRefactor() ?
+                            mCommunicationDeviceTracker.setCommunicationDevice(
+                                    AudioDeviceInfo.TYPE_BLE_HEADSET, device)
+                            : setLeAudioCommunicationDevice();
                 }
                 return true;
             }
@@ -676,8 +696,10 @@ public class BluetoothDeviceManager {
                  * Only after receiving ACTION_ACTIVE_DEVICE_CHANGED it is known that device that
                  * will be audio switched to is available to be choose as communication device */
                 if (!switchingBtDevices) {
-                    return mCommunicationDeviceTracker.setCommunicationDevice(
-                            AudioDeviceInfo.TYPE_HEARING_AID, null);
+                    return mFeatureFlags.callAudioCommunicationDeviceRefactor() ?
+                            mCommunicationDeviceTracker.setCommunicationDevice(
+                                    AudioDeviceInfo.TYPE_HEARING_AID, null)
+                            : setHearingAidCommunicationDevice();
                 }
                 return true;
             }
diff --git a/src/com/android/server/telecom/bluetooth/BluetoothRouteManager.java b/src/com/android/server/telecom/bluetooth/BluetoothRouteManager.java
index a88af5b0a..6cf6ae432 100644
--- a/src/com/android/server/telecom/bluetooth/BluetoothRouteManager.java
+++ b/src/com/android/server/telecom/bluetooth/BluetoothRouteManager.java
@@ -629,14 +629,22 @@ public class BluetoothRouteManager extends StateMachine {
         if (deviceType == BluetoothDeviceManager.DEVICE_TYPE_LE_AUDIO) {
             mLeAudioActiveDeviceCache = device;
             if (device == null) {
-                mCommunicationDeviceTracker.clearCommunicationDevice(
-                        AudioDeviceInfo.TYPE_BLE_HEADSET);
+                if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                    mCommunicationDeviceTracker.clearCommunicationDevice(
+                            AudioDeviceInfo.TYPE_BLE_HEADSET);
+                } else {
+                    mDeviceManager.clearLeAudioOrSpeakerCommunicationDevice();
+                }
             }
         } else if (deviceType == BluetoothDeviceManager.DEVICE_TYPE_HEARING_AID) {
             mHearingAidActiveDeviceCache = device;
             if (device == null) {
-                mCommunicationDeviceTracker.clearCommunicationDevice(
-                        AudioDeviceInfo.TYPE_HEARING_AID);
+                if (mFeatureFlags.callAudioCommunicationDeviceRefactor()) {
+                    mCommunicationDeviceTracker.clearCommunicationDevice(
+                            AudioDeviceInfo.TYPE_HEARING_AID);
+                } else {
+                    mDeviceManager.clearHearingAidOrSpeakerCommunicationDevice();
+                }
             }
         } else if (deviceType == BluetoothDeviceManager.DEVICE_TYPE_HEADSET) {
             mHfpActiveDeviceCache = device;
@@ -848,9 +856,13 @@ public class BluetoothRouteManager extends StateMachine {
             }
         }
 
+        boolean isHearingAidSetForCommunication =
+                mFeatureFlags.callAudioCommunicationDeviceRefactor()
+                ? mCommunicationDeviceTracker.isAudioDeviceSetForType(
+                        AudioDeviceInfo.TYPE_HEARING_AID)
+                : mDeviceManager.isHearingAidSetAsCommunicationDevice();
         if (bluetoothHearingAid != null) {
-            if (mCommunicationDeviceTracker.isAudioDeviceSetForType(
-                    AudioDeviceInfo.TYPE_HEARING_AID)) {
+            if (isHearingAidSetForCommunication) {
                 for (BluetoothDevice device : bluetoothAdapter.getActiveDevices(
                         BluetoothProfile.HEARING_AID)) {
                     if (device != null) {
@@ -862,9 +874,13 @@ public class BluetoothRouteManager extends StateMachine {
             }
         }
 
+        boolean isLeAudioSetForCommunication =
+                mFeatureFlags.callAudioCommunicationDeviceRefactor()
+                        ? mCommunicationDeviceTracker.isAudioDeviceSetForType(
+                        AudioDeviceInfo.TYPE_BLE_HEADSET)
+                        : mDeviceManager.isLeAudioCommunicationDevice();
         if (bluetoothLeAudio != null) {
-            if (mCommunicationDeviceTracker.isAudioDeviceSetForType(
-                    AudioDeviceInfo.TYPE_BLE_HEADSET)) {
+            if (isLeAudioSetForCommunication) {
                 for (BluetoothDevice device : bluetoothAdapter.getActiveDevices(
                         BluetoothProfile.LE_AUDIO)) {
                     if (device != null) {
diff --git a/src/com/android/server/telecom/bluetooth/BluetoothStateReceiver.java b/src/com/android/server/telecom/bluetooth/BluetoothStateReceiver.java
index ec4f2636e..d2521ac72 100644
--- a/src/com/android/server/telecom/bluetooth/BluetoothStateReceiver.java
+++ b/src/com/android/server/telecom/bluetooth/BluetoothStateReceiver.java
@@ -33,6 +33,7 @@ import android.telecom.Logging.Session;
 
 import com.android.internal.os.SomeArgs;
 import com.android.server.telecom.CallAudioCommunicationDeviceTracker;
+import com.android.server.telecom.flags.FeatureFlags;
 
 import static com.android.server.telecom.bluetooth.BluetoothRouteManager.BT_AUDIO_IS_ON;
 import static com.android.server.telecom.bluetooth.BluetoothRouteManager.BT_AUDIO_LOST;
@@ -59,6 +60,7 @@ public class BluetoothStateReceiver extends BroadcastReceiver {
     private final BluetoothRouteManager mBluetoothRouteManager;
     private final BluetoothDeviceManager mBluetoothDeviceManager;
     private CallAudioCommunicationDeviceTracker mCommunicationDeviceTracker;
+    private FeatureFlags mFeatureFlags;
 
     public void onReceive(Context context, Intent intent) {
         Log.startSession("BSR.oR");
@@ -208,19 +210,27 @@ public class BluetoothStateReceiver extends BroadcastReceiver {
                     /* In Le Audio case, once device got Active, the Telecom needs to make sure it
                      * is set as communication device before we can say that BT_AUDIO_IS_ON
                      */
+                    boolean isLeAudioSetForCommunication =
+                            mFeatureFlags.callAudioCommunicationDeviceRefactor()
+                                    ? mCommunicationDeviceTracker.setCommunicationDevice(
+                                    AudioDeviceInfo.TYPE_BLE_HEADSET, device)
+                                    : mBluetoothDeviceManager.setLeAudioCommunicationDevice();
                     if ((!usePreferredAudioProfile
                             || preferredDuplexProfile == BluetoothProfile.LE_AUDIO)
-                            && !mCommunicationDeviceTracker.setCommunicationDevice(
-                            AudioDeviceInfo.TYPE_BLE_HEADSET, device)) {
+                            && !isLeAudioSetForCommunication) {
                         Log.w(LOG_TAG,
                                 "Device %s cannot be use as LE audio communication device.",
                                 device);
                         return;
                     }
                 } else {
+                    boolean isHearingAidSetForCommunication =
+                            mFeatureFlags.callAudioCommunicationDeviceRefactor()
+                            ? mCommunicationDeviceTracker.setCommunicationDevice(
+                                    AudioDeviceInfo.TYPE_HEARING_AID, null)
+                            : mBluetoothDeviceManager.setHearingAidCommunicationDevice();
                     /* deviceType == BluetoothDeviceManager.DEVICE_TYPE_HEARING_AID */
-                    if (!mCommunicationDeviceTracker.setCommunicationDevice(
-                            AudioDeviceInfo.TYPE_HEARING_AID, null)) {
+                    if (!isHearingAidSetForCommunication) {
                         Log.w(LOG_TAG,
                                 "Device %s cannot be use as hearing aid communication device.",
                                 device);
@@ -238,10 +248,12 @@ public class BluetoothStateReceiver extends BroadcastReceiver {
 
     public BluetoothStateReceiver(BluetoothDeviceManager deviceManager,
             BluetoothRouteManager routeManager,
-            CallAudioCommunicationDeviceTracker communicationDeviceTracker) {
+            CallAudioCommunicationDeviceTracker communicationDeviceTracker,
+            FeatureFlags featureFlags) {
         mBluetoothDeviceManager = deviceManager;
         mBluetoothRouteManager = routeManager;
         mCommunicationDeviceTracker = communicationDeviceTracker;
+        mFeatureFlags = featureFlags;
     }
 
     public void setIsInCall(boolean isInCall) {
diff --git a/tests/src/com/android/server/telecom/tests/BluetoothDeviceManagerTest.java b/tests/src/com/android/server/telecom/tests/BluetoothDeviceManagerTest.java
index da3f40cf8..86d2beefb 100644
--- a/tests/src/com/android/server/telecom/tests/BluetoothDeviceManagerTest.java
+++ b/tests/src/com/android/server/telecom/tests/BluetoothDeviceManagerTest.java
@@ -52,6 +52,7 @@ import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -107,7 +108,7 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
         mContext = mComponentContextFixture.getTestDouble().getApplicationContext();
         mCommunicationDeviceTracker = new CallAudioCommunicationDeviceTracker(mContext);
         mBluetoothDeviceManager = new BluetoothDeviceManager(mContext, mAdapter,
-                mCommunicationDeviceTracker);
+                mCommunicationDeviceTracker, mFeatureFlags);
         mBluetoothDeviceManager.setBluetoothRouteManager(mRouteManager);
         mCommunicationDeviceTracker.setBluetoothRouteManager(mRouteManager);
 
@@ -120,7 +121,7 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
         serviceListenerUnderTest = serviceCaptor.getValue();
 
         receiverUnderTest = new BluetoothStateReceiver(mBluetoothDeviceManager,
-                mRouteManager, mCommunicationDeviceTracker);
+                mRouteManager, mCommunicationDeviceTracker, mFeatureFlags);
 
         mBluetoothDeviceManager.setHeadsetServiceForTesting(mBluetoothHeadset);
         mBluetoothDeviceManager.setHearingAidServiceForTesting(mBluetoothHearingAid);
@@ -131,6 +132,7 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
         verify(mBluetoothLeAudio).registerCallback(any(), leAudioCallbacksTest.capture());
 
         when(mSpeakerInfo.getType()).thenReturn(TYPE_BUILTIN_SPEAKER);
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(false);
     }
 
     @Override
@@ -440,7 +442,7 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
 
         when(mockAudioManager.getCommunicationDevice()).thenReturn(mockAudioDeviceInfo);
         mBluetoothDeviceManager.disconnectAudio();
-        verify(mockAudioManager).clearCommunicationDevice();
+        verify(mockAudioManager, atLeastOnce()).clearCommunicationDevice();
     }
 
     @SmallTest
@@ -477,7 +479,7 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
                 BluetoothDeviceManager.DEVICE_TYPE_LE_AUDIO));
 
         mBluetoothDeviceManager.disconnectAudio();
-        verify(mockAudioManager).clearCommunicationDevice();
+        verify(mockAudioManager, atLeastOnce()).clearCommunicationDevice();
     }
 
     @SmallTest
@@ -516,6 +518,7 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
     @SmallTest
     @Test
     public void testConnectMultipleLeAudioDevices() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         receiverUnderTest.setIsInCall(true);
         receiverUnderTest.onReceive(mContext,
                 buildConnectionActionIntent(BluetoothHeadset.STATE_CONNECTED, device1,
@@ -563,13 +566,8 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
     @SmallTest
     @Test
     public void testClearCommunicationDeviceOnActiveDeviceChange() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         receiverUnderTest.setIsInCall(true);
-//        receiverUnderTest.onReceive(mContext,
-//                buildConnectionActionIntent(BluetoothHeadset.STATE_CONNECTED, device1,
-//                        BluetoothDeviceManager.DEVICE_TYPE_LE_AUDIO));
-//        leAudioCallbacksTest.getValue().onGroupNodeAdded(device1, 1);
-//        when(mAdapter.setActiveDevice(nullable(BluetoothDevice.class),
-//                eq(BluetoothAdapter.ACTIVE_DEVICE_ALL))).thenReturn(true);
 
         List<AudioDeviceInfo> devices = new ArrayList<>();
         AudioDeviceInfo leAudioDevice1 = createMockAudioDeviceInfo(device1.getAddress(),
@@ -598,6 +596,7 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
     @SmallTest
     @Test
     public void testConnectDualModeEarbud() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         receiverUnderTest.setIsInCall(true);
 
         // LE Audio earbuds connected
@@ -691,46 +690,28 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
 
     @SmallTest
     @Test
-    public void testClearHearingAidCommunicationDevice() {
-        AudioDeviceInfo mockAudioDeviceInfo = mock(AudioDeviceInfo.class);
-        when(mockAudioDeviceInfo.getAddress()).thenReturn(DEVICE_ADDRESS_1);
-        when(mockAudioDeviceInfo.getType()).thenReturn(AudioDeviceInfo.TYPE_HEARING_AID);
-        List<AudioDeviceInfo> devices = new ArrayList<>();
-        devices.add(mockAudioDeviceInfo);
-
-        when(mockAudioManager.getAvailableCommunicationDevices())
-                .thenReturn(devices);
-        when(mockAudioManager.setCommunicationDevice(eq(mockAudioDeviceInfo)))
-                .thenReturn(true);
-
-        mCommunicationDeviceTracker.setCommunicationDevice(AudioDeviceInfo.TYPE_HEARING_AID, null);
-        when(mockAudioManager.getCommunicationDevice()).thenReturn(mSpeakerInfo);
-        mCommunicationDeviceTracker.clearCommunicationDevice(AudioDeviceInfo.TYPE_HEARING_AID);
-        verify(mRouteManager).onAudioLost(eq(DEVICE_ADDRESS_1));
-        assertFalse(mCommunicationDeviceTracker.isAudioDeviceSetForType(
-                AudioDeviceInfo.TYPE_HEARING_AID));
+    public void testClearHearingAidCommunicationDeviceLegacy() {
+        assertClearHearingAidOrLeCommunicationDevice(false, AudioDeviceInfo.TYPE_HEARING_AID);
     }
 
     @SmallTest
     @Test
-    public void testClearLeAudioCommunicationDevice() {
-        AudioDeviceInfo mockAudioDeviceInfo = createMockAudioDeviceInfo(DEVICE_ADDRESS_1,
-                AudioDeviceInfo.TYPE_BLE_HEADSET);
-        List<AudioDeviceInfo> devices = new ArrayList<>();
-        devices.add(mockAudioDeviceInfo);
+    public void testClearHearingAidCommunicationDeviceWithFlag() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
+        assertClearHearingAidOrLeCommunicationDevice(true, AudioDeviceInfo.TYPE_HEARING_AID);
+    }
 
-        when(mockAudioManager.getAvailableCommunicationDevices())
-                .thenReturn(devices);
-        when(mockAudioManager.setCommunicationDevice(eq(mockAudioDeviceInfo)))
-                .thenReturn(true);
+    @SmallTest
+    @Test
+    public void testClearLeAudioCommunicationDeviceLegacy() {
+        assertClearHearingAidOrLeCommunicationDevice(false, AudioDeviceInfo.TYPE_BLE_HEADSET);
+    }
 
-        mCommunicationDeviceTracker.setCommunicationDevice(
-                AudioDeviceInfo.TYPE_BLE_HEADSET, device1);
-        when(mockAudioManager.getCommunicationDevice()).thenReturn(mSpeakerInfo);
-        mCommunicationDeviceTracker.clearCommunicationDevice(AudioDeviceInfo.TYPE_BLE_HEADSET);
-        verify(mRouteManager).onAudioLost(eq(DEVICE_ADDRESS_1));
-        assertFalse(mCommunicationDeviceTracker.isAudioDeviceSetForType(
-                AudioDeviceInfo.TYPE_BLE_HEADSET));
+    @SmallTest
+    @Test
+    public void testClearLeAudioCommunicationDeviceWithFlag() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
+        assertClearHearingAidOrLeCommunicationDevice(true, AudioDeviceInfo.TYPE_BLE_HEADSET);
     }
 
     @SmallTest
@@ -769,6 +750,47 @@ public class BluetoothDeviceManagerTest extends TelecomTestCase {
         assertTrue(mBluetoothDeviceManager.isInbandRingingEnabled());
     }
 
+    private void assertClearHearingAidOrLeCommunicationDevice(
+            boolean flagEnabled, int device_type
+    ) {
+        AudioDeviceInfo mockAudioDeviceInfo = mock(AudioDeviceInfo.class);
+        when(mockAudioDeviceInfo.getAddress()).thenReturn(DEVICE_ADDRESS_1);
+        when(mockAudioDeviceInfo.getType()).thenReturn(device_type);
+        List<AudioDeviceInfo> devices = new ArrayList<>();
+        devices.add(mockAudioDeviceInfo);
+
+        when(mockAudioManager.getAvailableCommunicationDevices())
+                .thenReturn(devices);
+        when(mockAudioManager.setCommunicationDevice(eq(mockAudioDeviceInfo)))
+                .thenReturn(true);
+
+        if (flagEnabled) {
+            BluetoothDevice btDevice = device_type == AudioDeviceInfo.TYPE_BLE_HEADSET
+                    ? device1 : null;
+            mCommunicationDeviceTracker.setCommunicationDevice(device_type, btDevice);
+        } else {
+            if (device_type == AudioDeviceInfo.TYPE_BLE_HEADSET) {
+                mBluetoothDeviceManager.setLeAudioCommunicationDevice();
+            } else {
+                mBluetoothDeviceManager.setHearingAidCommunicationDevice();
+            }
+        }
+        when(mockAudioManager.getCommunicationDevice()).thenReturn(mSpeakerInfo);
+        if (flagEnabled) {
+            mCommunicationDeviceTracker.clearCommunicationDevice(device_type);
+            assertFalse(mCommunicationDeviceTracker.isAudioDeviceSetForType(device_type));
+        } else {
+            if (device_type == AudioDeviceInfo.TYPE_BLE_HEADSET) {
+                mBluetoothDeviceManager.clearLeAudioOrSpeakerCommunicationDevice();
+                assertFalse(mBluetoothDeviceManager.isLeAudioCommunicationDevice());
+            } else {
+                mBluetoothDeviceManager.clearHearingAidOrSpeakerCommunicationDevice();
+                assertFalse(mBluetoothDeviceManager.isHearingAidSetAsCommunicationDevice());
+            }
+        }
+        verify(mRouteManager).onAudioLost(eq(DEVICE_ADDRESS_1));
+    }
+
     private AudioDeviceInfo createMockAudioDeviceInfo(String address, int audioType) {
         AudioDeviceInfo mockAudioDeviceInfo = mock(AudioDeviceInfo.class);
         when(mockAudioDeviceInfo.getType()).thenReturn(audioType);
diff --git a/tests/src/com/android/server/telecom/tests/CallAudioModeStateMachineTest.java b/tests/src/com/android/server/telecom/tests/CallAudioModeStateMachineTest.java
index 8c6d84c7a..652e5ce60 100644
--- a/tests/src/com/android/server/telecom/tests/CallAudioModeStateMachineTest.java
+++ b/tests/src/com/android/server/telecom/tests/CallAudioModeStateMachineTest.java
@@ -35,6 +35,7 @@ import android.media.AudioManager;
 import android.os.HandlerThread;
 import android.test.suitebuilder.annotation.SmallTest;
 
+import com.android.server.telecom.CallAudioCommunicationDeviceTracker;
 import com.android.server.telecom.CallAudioManager;
 import com.android.server.telecom.CallAudioModeStateMachine;
 import com.android.server.telecom.CallAudioModeStateMachine.MessageArgs.Builder;
@@ -57,6 +58,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Mock private AudioManager mAudioManager;
     @Mock private CallAudioManager mCallAudioManager;
     @Mock private CallAudioRouteStateMachine mCallAudioRouteStateMachine;
+    @Mock private CallAudioCommunicationDeviceTracker mCommunicationDeviceTracker;
 
     private HandlerThread mTestThread;
 
@@ -83,7 +85,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testNoFocusWhenRingerSilenced() throws Throwable {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(CallAudioModeStateMachine.ABANDON_FOCUS_FOR_TESTING);
         waitForHandlerAction(sm.getHandler(), TEST_TIMEOUT);
@@ -115,7 +117,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testSwitchToStreamingMode() {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(CallAudioModeStateMachine.ABANDON_FOCUS_FOR_TESTING);
         waitForHandlerAction(sm.getHandler(), TEST_TIMEOUT);
@@ -145,7 +147,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testExitStreamingMode() {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(CallAudioModeStateMachine.ENTER_STREAMING_FOCUS_FOR_TESTING);
         waitForHandlerAction(sm.getHandler(), TEST_TIMEOUT);
@@ -173,7 +175,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testNoRingWhenDeviceIsAtEar() {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(CallAudioModeStateMachine.ABANDON_FOCUS_FOR_TESTING);
         sm.sendMessage(CallAudioModeStateMachine.NEW_HOLDING_CALL, new Builder()
@@ -209,7 +211,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testRegainFocusWhenHfpIsConnectedSilenced() throws Throwable {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(CallAudioModeStateMachine.ABANDON_FOCUS_FOR_TESTING);
         waitForHandlerAction(sm.getHandler(), TEST_TIMEOUT);
@@ -253,7 +255,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testDoNotRingTwiceWhenHfpConnected() {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(CallAudioModeStateMachine.ABANDON_FOCUS_FOR_TESTING);
         waitForHandlerAction(sm.getHandler(), TEST_TIMEOUT);
@@ -291,7 +293,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testStartRingingAfterHfpConnectedIfNotAlreadyPlaying() {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(CallAudioModeStateMachine.ABANDON_FOCUS_FOR_TESTING);
         waitForHandlerAction(sm.getHandler(), TEST_TIMEOUT);
@@ -329,7 +331,7 @@ public class CallAudioModeStateMachineTest extends TelecomTestCase {
     @Test
     public void testAudioFocusRequestWithResolveHiddenDependencies() {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         when(mFeatureFlags.telecomResolveHiddenDependencies()).thenReturn(true);
         ArgumentCaptor<AudioFocusRequest> captor = ArgumentCaptor.forClass(AudioFocusRequest.class);
         sm.setCallAudioManager(mCallAudioManager);
diff --git a/tests/src/com/android/server/telecom/tests/CallAudioModeTransitionTests.java b/tests/src/com/android/server/telecom/tests/CallAudioModeTransitionTests.java
index 21b5312d0..3690d5f11 100644
--- a/tests/src/com/android/server/telecom/tests/CallAudioModeTransitionTests.java
+++ b/tests/src/com/android/server/telecom/tests/CallAudioModeTransitionTests.java
@@ -24,6 +24,7 @@ import android.media.AudioManager;
 import android.os.HandlerThread;
 import android.test.suitebuilder.annotation.SmallTest;
 
+import com.android.server.telecom.CallAudioCommunicationDeviceTracker;
 import com.android.server.telecom.CallAudioManager;
 import com.android.server.telecom.CallAudioModeStateMachine;
 import com.android.server.telecom.CallAudioModeStateMachine.MessageArgs;
@@ -108,6 +109,7 @@ public class CallAudioModeTransitionTests extends TelecomTestCase {
     @Mock private SystemStateHelper mSystemStateHelper;
     @Mock private AudioManager mAudioManager;
     @Mock private CallAudioManager mCallAudioManager;
+    @Mock private CallAudioCommunicationDeviceTracker mCommunicationDeviceTracker;
     private final ModeTestParameters mParams;
     private HandlerThread mTestThread;
 
@@ -135,7 +137,7 @@ public class CallAudioModeTransitionTests extends TelecomTestCase {
     @SmallTest
     public void modeTransitionTest() {
         CallAudioModeStateMachine sm = new CallAudioModeStateMachine(mSystemStateHelper,
-                mAudioManager, mTestThread.getLooper(), mFeatureFlags);
+                mAudioManager, mTestThread.getLooper(), mFeatureFlags, mCommunicationDeviceTracker);
         sm.setCallAudioManager(mCallAudioManager);
         sm.sendMessage(mParams.initialAudioState);
         waitForHandlerAction(sm.getHandler(), TEST_TIMEOUT);
diff --git a/tests/src/com/android/server/telecom/tests/CallAudioRouteStateMachineTest.java b/tests/src/com/android/server/telecom/tests/CallAudioRouteStateMachineTest.java
index 9055dafaf..6dd93e557 100644
--- a/tests/src/com/android/server/telecom/tests/CallAudioRouteStateMachineTest.java
+++ b/tests/src/com/android/server/telecom/tests/CallAudioRouteStateMachineTest.java
@@ -143,6 +143,7 @@ public class CallAudioRouteStateMachineTest extends TelecomTestCase {
         doNothing().when(mockConnectionServiceWrapper).onCallAudioStateChanged(any(Call.class),
                 any(CallAudioState.class));
         when(mFeatureFlags.ignoreAutoRouteToWatchDevice()).thenReturn(false);
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(false);
     }
 
     @Override
@@ -708,6 +709,7 @@ public class CallAudioRouteStateMachineTest extends TelecomTestCase {
     @SmallTest
     @Test
     public void testCallDisconnectedWhenAudioRoutedToBluetooth() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         CallAudioRouteStateMachine stateMachine = new CallAudioRouteStateMachine(
                 mContext,
                 mockCallsManager,
@@ -896,6 +898,7 @@ public class CallAudioRouteStateMachineTest extends TelecomTestCase {
     @SmallTest
     @Test
     public void testSetAndClearEarpieceCommunicationDevice() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         CallAudioRouteStateMachine stateMachine = new CallAudioRouteStateMachine(
                 mContext,
                 mockCallsManager,
@@ -953,18 +956,21 @@ public class CallAudioRouteStateMachineTest extends TelecomTestCase {
     @SmallTest
     @Test
     public void testSetAndClearWiredHeadsetCommunicationDevice() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         verifySetAndClearHeadsetCommunicationDevice(AudioDeviceInfo.TYPE_WIRED_HEADSET);
     }
 
     @SmallTest
     @Test
     public void testSetAndClearUsbHeadsetCommunicationDevice() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         verifySetAndClearHeadsetCommunicationDevice(AudioDeviceInfo.TYPE_USB_HEADSET);
     }
 
     @SmallTest
     @Test
     public void testActiveFocusRouteSwitchFromQuiescentBluetooth() {
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         CallAudioRouteStateMachine stateMachine = new CallAudioRouteStateMachine(
                 mContext,
                 mockCallsManager,
@@ -1168,6 +1174,7 @@ public class CallAudioRouteStateMachineTest extends TelecomTestCase {
     @Test
     public void testIgnoreImplicitBTSwitchWhenDeviceIsWatch() {
         when(mFeatureFlags.ignoreAutoRouteToWatchDevice()).thenReturn(true);
+        when(mFeatureFlags.callAudioCommunicationDeviceRefactor()).thenReturn(true);
         CallAudioRouteStateMachine stateMachine = new CallAudioRouteStateMachine(
                 mContext,
                 mockCallsManager,
diff --git a/tests/src/com/android/server/telecom/tests/CallsManagerTest.java b/tests/src/com/android/server/telecom/tests/CallsManagerTest.java
index 36488b59e..fe94da4f3 100644
--- a/tests/src/com/android/server/telecom/tests/CallsManagerTest.java
+++ b/tests/src/com/android/server/telecom/tests/CallsManagerTest.java
@@ -310,7 +310,7 @@ public class CallsManagerTest extends TelecomTestCase {
                 mCallEndpointController);
         when(mCallAudioRouteStateMachineFactory.create(any(), any(), any(), any(), any(), any(),
                 anyInt(), any(), any(), any())).thenReturn(mCallAudioRouteStateMachine);
-        when(mCallAudioModeStateMachineFactory.create(any(), any(), any()))
+        when(mCallAudioModeStateMachineFactory.create(any(), any(), any(), any()))
                 .thenReturn(mCallAudioModeStateMachine);
         when(mClockProxy.currentTimeMillis()).thenReturn(System.currentTimeMillis());
         when(mClockProxy.elapsedRealtime()).thenReturn(SystemClock.elapsedRealtime());
diff --git a/tests/src/com/android/server/telecom/tests/TelecomSystemTest.java b/tests/src/com/android/server/telecom/tests/TelecomSystemTest.java
index 4a5f0036c..d7ba6e8fb 100644
--- a/tests/src/com/android/server/telecom/tests/TelecomSystemTest.java
+++ b/tests/src/com/android/server/telecom/tests/TelecomSystemTest.java
@@ -546,9 +546,12 @@ public class TelecomSystemTest extends TelecomTestCase{
                 new CallAudioModeStateMachine.Factory() {
                     @Override
                     public CallAudioModeStateMachine create(SystemStateHelper systemStateHelper,
-                            AudioManager am, FeatureFlags featureFlags) {
+                            AudioManager am, FeatureFlags featureFlags,
+                            CallAudioCommunicationDeviceTracker callAudioCommunicationDeviceTracker
+                    ) {
                         return new CallAudioModeStateMachine(systemStateHelper, am,
-                                mHandlerThread.getLooper(), featureFlags);
+                                mHandlerThread.getLooper(), featureFlags,
+                                callAudioCommunicationDeviceTracker);
                     }
                 },
                 mClockProxy,