WifiController: use callbacks instead of bcasts

Instead of depending on system broadcasts to communicate between
internal wifi system components we now use callbacks.  This allows the
modes to communicate back to WifiController about the current state of
the operating mode.

Bug: 31346104
Test: manually verify events are received in WifiController (via logs)
Test: frameworks/opt/net/wifi/tests/wifitests/runtests.sh
Test: will submit for integration tests

Change-Id: I45486c95e094776b429ab7755af264eabd79aa77

6 files changed, 148 insertions, 31 deletions

diff --git a/service/java/com/android/server/wifi/WifiController.java b/service/java/com/android/server/wifi/WifiController.java
index 607d86341..4fb948e94 100644
--- a/service/java/com/android/server/wifi/WifiController.java
+++ b/service/java/com/android/server/wifi/WifiController.java
@@ -30,6 +30,7 @@ import android.os.Message;
 import android.os.SystemClock;
 import android.os.WorkSource;
 import android.provider.Settings;
+import android.util.Log;
 
 import com.android.internal.util.Protocol;
 import com.android.internal.util.State;
@@ -93,6 +94,8 @@ public class WifiController extends StateMachine {
     private static final int CMD_RECOVERY_RESTART_WIFI_CONTINUE = BASE + 18;
     // Command to disable wifi when SelfRecovery is throttled or otherwise not doing full recovery
     static final int CMD_RECOVERY_DISABLE_WIFI                  = BASE + 19;
+    static final int CMD_STA_STOPPED                            = BASE + 20;
+    static final int CMD_SCANNING_STOPPED                       = BASE + 21;
 
     private DefaultState mDefaultState = new DefaultState();
     private StaEnabledState mStaEnabledState = new StaEnabledState();
@@ -102,6 +105,9 @@ public class WifiController extends StateMachine {
     private DeviceActiveState mDeviceActiveState = new DeviceActiveState();
     private EcmState mEcmState = new EcmState();
 
+    private ScanOnlyModeManager.Listener mScanOnlyModeCallback = new ScanOnlyCallback();
+    private ClientModeManager.Listener mClientModeCallback = new ClientModeCallback();
+
     WifiController(Context context, WifiStateMachine wsm, Looper wifiStateMachineLooper,
                    WifiSettingsStore wss, Looper wifiServiceLooper, FrameworkFacade f,
                    WifiStateMachinePrime wsmp) {
@@ -113,6 +119,7 @@ public class WifiController extends StateMachine {
         mWifiStateMachinePrime = wsmp;
         mSettingsStore = wss;
 
+        // CHECKSTYLE:OFF IndentationCheck
         addState(mDefaultState);
             addState(mApStaDisabledState, mDefaultState);
             addState(mStaEnabledState, mDefaultState);
@@ -120,6 +127,7 @@ public class WifiController extends StateMachine {
             addState(mStaDisabledWithScanState, mDefaultState);
             addState(mApEnabledState, mDefaultState);
             addState(mEcmState, mDefaultState);
+        // CHECKSTYLE:ON IndentationCheck
 
         boolean isAirplaneModeOn = mSettingsStore.isAirplaneModeOn();
         boolean isWifiEnabled = mSettingsStore.isWifiToggleEnabled();
@@ -142,6 +150,10 @@ public class WifiController extends StateMachine {
         setLogRecSize(100);
         setLogOnlyTransitions(false);
 
+        // register for state updates via callbacks (vs the intents registered below)
+        mWifiStateMachinePrime.registerScanOnlyCallback(mScanOnlyModeCallback);
+        mWifiStateMachinePrime.registerClientModeCallback(mClientModeCallback);
+
         IntentFilter filter = new IntentFilter();
         filter.addAction(WifiManager.NETWORK_STATE_CHANGED_ACTION);
         filter.addAction(WifiManager.WIFI_AP_STATE_CHANGED_ACTION);
@@ -160,19 +172,11 @@ public class WifiController extends StateMachine {
                                     WifiManager.EXTRA_WIFI_AP_STATE,
                                     WifiManager.WIFI_AP_STATE_FAILED);
                             if (state == WifiManager.WIFI_AP_STATE_FAILED) {
-                                loge(TAG + "SoftAP start failed");
+                                Log.e(TAG, "SoftAP start failed");
                                 sendMessage(CMD_AP_START_FAILURE);
                             } else if (state == WifiManager.WIFI_AP_STATE_DISABLED) {
                                 sendMessage(CMD_AP_STOPPED);
                             }
-                        } else if (action.equals(WifiManager.WIFI_STATE_CHANGED_ACTION)) {
-                            int state = intent.getIntExtra(
-                                    WifiManager.EXTRA_WIFI_STATE,
-                                    WifiManager.WIFI_STATE_UNKNOWN);
-                            if (state == WifiManager.WIFI_STATE_UNKNOWN) {
-                                loge(TAG + "Wifi turn on failed");
-                                sendMessage(CMD_STA_START_FAILURE);
-                            }
                         } else if (action.equals(LocationManager.MODE_CHANGED_ACTION)) {
                             // Location mode has been toggled...  trigger with the scan change
                             // update to make sure we are in the correct mode
@@ -194,6 +198,48 @@ public class WifiController extends StateMachine {
         return mSettingsStore.isScanAlwaysAvailable();
     }
 
+    /**
+     * Listener used to receive scan mode updates - really needed for disabled updates to trigger
+     * mode changes.
+     */
+    private class ScanOnlyCallback implements ScanOnlyModeManager.Listener {
+        @Override
+        public void onStateChanged(int state) {
+            if (state == WifiManager.WIFI_STATE_UNKNOWN) {
+                Log.d(TAG, "ScanOnlyMode unexpected failure: state unknown");
+            } else if (state == WifiManager.WIFI_STATE_DISABLED) {
+                Log.d(TAG, "ScanOnlyMode stopped");
+                sendMessage(CMD_SCANNING_STOPPED);
+            } else if (state == WifiManager.WIFI_STATE_ENABLED) {
+                // scan mode is ready to go
+                Log.d(TAG, "scan mode active");
+            } else {
+                Log.d(TAG, "unexpected state update: " + state);
+            }
+        }
+    }
+
+    /**
+     * Listener used to receive client mode updates
+     */
+    private class ClientModeCallback implements ClientModeManager.Listener {
+        @Override
+        public void onStateChanged(int state) {
+            if (state == WifiManager.WIFI_STATE_UNKNOWN) {
+                logd("ClientMode unexpected failure: state unknown");
+                sendMessage(CMD_STA_START_FAILURE);
+            } else if (state == WifiManager.WIFI_STATE_DISABLED) {
+                logd("ClientMode stopped");
+                sendMessage(CMD_STA_STOPPED);
+            } else if (state == WifiManager.WIFI_STATE_ENABLED) {
+                // scan mode is ready to go
+                logd("client mode active");
+            } else {
+                logd("unexpected state update: " + state);
+            }
+        }
+    }
+
     private void readWifiReEnableDelay() {
         mReEnableDelayMillis = mFacade.getLongSetting(mContext,
                 Settings.Global.WIFI_REENABLE_DELAY_MS, DEFAULT_REENABLE_DELAY_MS);
@@ -216,6 +262,8 @@ public class WifiController extends StateMachine {
                 case CMD_EMERGENCY_CALL_STATE_CHANGED:
                 case CMD_AP_START_FAILURE:
                 case CMD_AP_STOPPED:
+                case CMD_SCANNING_STOPPED:
+                case CMD_STA_STOPPED:
                 case CMD_STA_START_FAILURE:
                 case CMD_RECOVERY_RESTART_WIFI_CONTINUE:
                 case CMD_RECOVERY_DISABLE_WIFI:
@@ -247,7 +295,7 @@ public class WifiController extends StateMachine {
 
         @Override
         public void enter() {
-            mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE);
+            mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE, null);
             mWifiStateMachinePrime.disableWifi();
             // Supplicant can't restart right away, so note the time we switched off
             mDisabledTimestamp = SystemClock.elapsedRealtime();
@@ -273,7 +321,8 @@ public class WifiController extends StateMachine {
                         // are headed to connect mode before going to the DeviceActiveState
                         // since that will start supplicant and WifiStateMachine may not know
                         // what state to head to (it might go to scan mode).
-                        mWifiStateMachine.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+                        mWifiStateMachine.setOperationalMode(WifiStateMachine.CONNECT_MODE,
+                                                             mClientModeCallback);
                         transitionTo(mDeviceActiveState);
                     } else if (checkScanOnlyModeAvailable()) {
                         transitionTo(mStaDisabledWithScanState);
@@ -284,7 +333,7 @@ public class WifiController extends StateMachine {
                         transitionTo(mStaDisabledWithScanState);
                         break;
                     }
-                    mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE);
+                    mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE, null);
                     break;
                 case CMD_SET_AP:
                     if (msg.arg1 == 1) {
@@ -412,7 +461,8 @@ public class WifiController extends StateMachine {
         @Override
         public void enter() {
             // first send the message to WSM to trigger the transition and act as a shadow
-            mWifiStateMachine.setOperationalMode(WifiStateMachine.SCAN_ONLY_WITH_WIFI_OFF_MODE);
+            mWifiStateMachine.setOperationalMode(
+                    WifiStateMachine.SCAN_ONLY_WITH_WIFI_OFF_MODE, null);
 
             // now trigger the actual mode switch in WifiStateMachinePrime
             mWifiStateMachinePrime.enterScanOnlyMode();
@@ -438,7 +488,7 @@ public class WifiController extends StateMachine {
                             break;
                         }
                         // transition from scan mode to initial state in WifiStateMachine
-                        mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE);
+                        mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE, null);
                         transitionTo(mDeviceActiveState);
                     }
                     break;
@@ -599,7 +649,7 @@ public class WifiController extends StateMachine {
         private int mEcmEntryCount;
         @Override
         public void enter() {
-            mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE);
+            mWifiStateMachine.setOperationalMode(WifiStateMachine.DISABLED_MODE, null);
             mWifiStateMachinePrime.disableWifi();
             mWifiStateMachine.clearANQPCache();
             mEcmEntryCount = 1;
@@ -661,7 +711,8 @@ public class WifiController extends StateMachine {
     class DeviceActiveState extends State {
         @Override
         public void enter() {
-            mWifiStateMachine.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+            mWifiStateMachine.setOperationalMode(WifiStateMachine.CONNECT_MODE,
+                                                 mClientModeCallback);
             mWifiStateMachinePrime.enterClientMode();
             mWifiStateMachine.setHighPerfModeEnabled(false);
         }
diff --git a/service/java/com/android/server/wifi/WifiStateMachine.java b/service/java/com/android/server/wifi/WifiStateMachine.java
index c44384f40..08b08d55b 100644
--- a/service/java/com/android/server/wifi/WifiStateMachine.java
+++ b/service/java/com/android/server/wifi/WifiStateMachine.java
@@ -296,6 +296,8 @@ public class WifiStateMachine extends StateMachine {
     // variable indicating we are expecting a mode switch - do not attempt recovery for failures
     private boolean mModeChange = false;
 
+    private ClientModeManager.Listener mClientModeCallback = null;
+
     private boolean mBluetoothConnectionActive = false;
 
     private PowerManager.WakeLock mSuspendWakeLock;
@@ -1544,9 +1546,12 @@ public class WifiStateMachine extends StateMachine {
     /**
      * TODO: doc
      */
-    public void setOperationalMode(int mode) {
+    public void setOperationalMode(int mode, ClientModeManager.Listener callback) {
         if (mVerboseLoggingEnabled) log("setting operational mode to " + String.valueOf(mode));
         mModeChange = true;
+        if (callback != null) {
+            mClientModeCallback = callback;
+        }
         sendMessage(CMD_SET_OPERATIONAL_MODE, mode, 0);
     }
 
@@ -2540,6 +2545,17 @@ public class WifiStateMachine extends StateMachine {
 
         if (mVerboseLoggingEnabled) log("setWifiState: " + syncGetWifiStateByName());
 
+        // first let WifiController know what is going on
+        if (mClientModeCallback != null) {
+            mClientModeCallback.onStateChanged(wifiState);
+            // once this instance of client mode is complete, remove the callback so we don't
+            // confuse ourselves
+            if (wifiState == WifiManager.WIFI_STATE_UNKNOWN
+                    || wifiState == WifiManager.WIFI_STATE_DISABLED) {
+                mClientModeCallback = null;
+            }
+        }
+
         final Intent intent = new Intent(WifiManager.WIFI_STATE_CHANGED_ACTION);
         intent.addFlags(Intent.FLAG_RECEIVER_REGISTERED_ONLY_BEFORE_BOOT);
         intent.putExtra(WifiManager.EXTRA_WIFI_STATE, wifiState);
diff --git a/service/java/com/android/server/wifi/WifiStateMachinePrime.java b/service/java/com/android/server/wifi/WifiStateMachinePrime.java
index 8099066e5..add602f4b 100644
--- a/service/java/com/android/server/wifi/WifiStateMachinePrime.java
+++ b/service/java/com/android/server/wifi/WifiStateMachinePrime.java
@@ -98,14 +98,31 @@ public class WifiStateMachinePrime {
     private StatusListener mWifiNativeStatusListener;
 
     private WifiManager.SoftApCallback mSoftApCallback;
+    private ScanOnlyModeManager.Listener mScanOnlyCallback;
+    private ClientModeManager.Listener mClientModeCallback;
 
     /**
      * Called from WifiServiceImpl to register a callback for notifications from SoftApManager
      */
-    public void registerSoftApCallback(WifiManager.SoftApCallback callback) {
+    public void registerSoftApCallback(@NonNull WifiManager.SoftApCallback callback) {
         mSoftApCallback = callback;
     }
 
+    /**
+     * Called from WifiController to register a callback for notifications from ScanOnlyModeManager
+     */
+    public void registerScanOnlyCallback(@NonNull ScanOnlyModeManager.Listener callback) {
+        mScanOnlyCallback = callback;
+    }
+
+    /**
+     * Called from WifiController to register a callback for notifications from ClientModeManager
+     */
+    public void registerClientModeCallback(@NonNull ClientModeManager.Listener callback) {
+        // uncomment when client mode moves here
+        // mClientModeCallback = callback;
+    }
+
     WifiStateMachinePrime(WifiInjector wifiInjector,
                           Looper looper,
                           WifiNative wifiNative,
@@ -376,10 +393,14 @@ public class WifiStateMachinePrime {
                         break;
                     case CMD_SCAN_ONLY_MODE_FAILED:
                         Log.d(TAG, "ScanOnlyMode failed, return to WifiDisabledState.");
+                        // notify WifiController that ScanOnlyMode failed
+                        mScanOnlyCallback.onStateChanged(WifiManager.WIFI_STATE_UNKNOWN);
                         mModeStateMachine.transitionTo(mWifiDisabledState);
                         break;
                     case CMD_SCAN_ONLY_MODE_STOPPED:
                         Log.d(TAG, "ScanOnlyMode stopped, return to WifiDisabledState.");
+                        // notify WifiController that ScanOnlyMode stopped
+                        mScanOnlyCallback.onStateChanged(WifiManager.WIFI_STATE_DISABLED);
                         mModeStateMachine.transitionTo(mWifiDisabledState);
                         break;
                     default:
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiControllerTest.java b/tests/wifitests/src/com/android/server/wifi/WifiControllerTest.java
index 5211964a6..0b7bf38eb 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiControllerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiControllerTest.java
@@ -370,7 +370,8 @@ public class WifiControllerTest {
         mLooper.dispatchAll();
 
         InOrder inOrder = inOrder(mWifiStateMachine);
-        inOrder.verify(mWifiStateMachine).setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        inOrder.verify(mWifiStateMachine).setOperationalMode(eq(WifiStateMachine.CONNECT_MODE),
+                                                             any());
         assertEquals("DeviceActiveState", getCurrentState().getName());
     }
 
@@ -396,7 +397,8 @@ public class WifiControllerTest {
         mLooper.dispatchAll();
 
         InOrder inOrder = inOrder(mWifiStateMachine);
-        inOrder.verify(mWifiStateMachine).setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        inOrder.verify(mWifiStateMachine).setOperationalMode(eq(WifiStateMachine.CONNECT_MODE),
+                                                             any());
         assertEquals("DeviceActiveState", getCurrentState().getName());
     }
 
@@ -427,7 +429,7 @@ public class WifiControllerTest {
         reset(mWifiStateMachine);
         mWifiController.sendMessage(CMD_RECOVERY_DISABLE_WIFI);
         mLooper.dispatchAll();
-        verify(mWifiStateMachine).setOperationalMode(WifiStateMachine.DISABLED_MODE);
+        verify(mWifiStateMachine).setOperationalMode(eq(WifiStateMachine.DISABLED_MODE), any());
     }
 
     /**
@@ -513,7 +515,8 @@ public class WifiControllerTest {
         mWifiController.sendMessage(CMD_RECOVERY_RESTART_WIFI);
         mLooper.dispatchAll();
         InOrder inOrder = inOrder(mWifiStateMachine);
-        inOrder.verify(mWifiStateMachine).setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        inOrder.verify(mWifiStateMachine).setOperationalMode(eq(WifiStateMachine.CONNECT_MODE),
+                                                             any());
         assertEquals("DeviceActiveState", getCurrentState().getName());
     }
 
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiStateMachinePrimeTest.java b/tests/wifitests/src/com/android/server/wifi/WifiStateMachinePrimeTest.java
index 9deb628eb..2df6e9a9f 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiStateMachinePrimeTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiStateMachinePrimeTest.java
@@ -60,6 +60,8 @@ public class WifiStateMachinePrimeTest {
     @Mock SelfRecovery mSelfRecovery;
     @Mock BaseWifiDiagnostics mWifiDiagnostics;
     ScanOnlyModeManager.Listener mScanOnlyListener;
+    ScanOnlyModeCallback mScanOnlyCallback = new ScanOnlyModeCallback();
+    ClientModeCallback mClientModeCallback = new ClientModeCallback();
     WifiManager.SoftApCallback mSoftApManagerCallback;
     @Mock WifiManager.SoftApCallback mSoftApStateMachineCallback;
     WifiNative.StatusListener mWifiNativeStatusListener;
@@ -88,6 +90,8 @@ public class WifiStateMachinePrimeTest {
         mWifiNativeStatusListener = mStatusListenerCaptor.getValue();
 
         mWifiStateMachinePrime.registerSoftApCallback(mSoftApStateMachineCallback);
+        mWifiStateMachinePrime.registerScanOnlyCallback(mScanOnlyCallback);
+        mWifiStateMachinePrime.registerClientModeCallback(mClientModeCallback);
     }
 
     private WifiStateMachinePrime createWifiStateMachinePrime() {
@@ -106,6 +110,24 @@ public class WifiStateMachinePrimeTest {
         mWifiStateMachinePrime = null;
     }
 
+    private class ClientModeCallback implements ClientModeManager.Listener {
+        public int currentState = WifiManager.WIFI_STATE_UNKNOWN;
+
+        @Override
+        public void onStateChanged(int state) {
+            currentState = state;
+        }
+    }
+
+    private class ScanOnlyModeCallback implements ScanOnlyModeManager.Listener {
+        public int currentState = WifiManager.WIFI_STATE_UNKNOWN;
+
+        @Override
+        public void onStateChanged(int state) {
+            currentState = state;
+        }
+    }
+
     private void enterSoftApActiveMode() throws Exception {
         enterSoftApActiveMode(
                 new SoftApModeConfiguration(WifiManager.IFACE_IP_MODE_TETHERED, null));
@@ -127,6 +149,9 @@ public class WifiStateMachinePrimeTest {
                         any(ScanOnlyModeManager.Listener.class));
         mWifiStateMachinePrime.enterScanOnlyMode();
         mLooper.dispatchAll();
+        mScanOnlyListener.onStateChanged(WifiManager.WIFI_STATE_ENABLED);
+        mLooper.dispatchAll();
+
         assertEquals(SCAN_ONLY_MODE_STATE_STRING, mWifiStateMachinePrime.getCurrentMode());
         verify(mScanOnlyModeManager).start();
         verify(mBatteryStats).noteWifiOn();
@@ -321,6 +346,7 @@ public class WifiStateMachinePrimeTest {
         verify(mScanOnlyModeManager).stop();
         verify(mWifiNative, times(3)).teardownAllInterfaces();
         verify(mBatteryStats).noteWifiOff();
+        assertEquals(WifiManager.WIFI_STATE_UNKNOWN, mScanOnlyCallback.currentState);
     }
 
     /**
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
index 28353074d..8a83b3d8b 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
@@ -597,7 +597,7 @@ public class WifiStateMachineTest {
     @Test
     public void shouldRequireSupplicantStartupToLeaveInitialState() throws Exception {
         when(mWifiNative.setupInterfaceForClientMode(anyBoolean(), any())).thenReturn(null);
-        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE, null);
         mLooper.dispatchAll();
         assertEquals("DefaultState", getCurrentState().getName());
         // we should be sending a wifi enabling update
@@ -609,7 +609,7 @@ public class WifiStateMachineTest {
     public void loadComponentsFailure() throws Exception {
         when(mWifiNative.setupInterfaceForClientMode(anyBoolean(), any())).thenReturn(null);
 
-        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE, null);
         mLooper.dispatchAll();
         assertEquals("DefaultState", getCurrentState().getName());
     }
@@ -620,7 +620,7 @@ public class WifiStateMachineTest {
         assertEquals("DefaultState", getCurrentState().getName());
         assertEquals(WifiStateMachine.DISABLED_MODE, mWsm.getOperationalModeForTest());
 
-        mWsm.setOperationalMode(WifiStateMachine.DISABLED_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.DISABLED_MODE, null);
         mLooper.dispatchAll();
         assertEquals(WifiStateMachine.DISABLED_MODE, mWsm.getOperationalModeForTest());
         assertEquals("DefaultState", getCurrentState().getName());
@@ -634,7 +634,7 @@ public class WifiStateMachineTest {
         assertEquals(WifiStateMachine.DISABLED_MODE, mWsm.getOperationalModeForTest());
 
         // But if someone tells us to enter connect mode, we start up supplicant
-        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE, null);
         mLooper.dispatchAll();
         assertEquals("DisconnectedState", getCurrentState().getName());
     }
@@ -649,7 +649,7 @@ public class WifiStateMachineTest {
         assertEquals("DefaultState", getCurrentState().getName());
         assertEquals(WifiManager.WIFI_STATE_DISABLED, mWsm.syncGetWifiState());
 
-        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_MODE, null);
         mLooper.dispatchAll();
         assertEquals(WifiStateMachine.SCAN_ONLY_MODE, mWsm.getOperationalModeForTest());
         assertEquals("ScanModeState", getCurrentState().getName());
@@ -671,7 +671,7 @@ public class WifiStateMachineTest {
         reset(mContext);
 
         // now go back to scan mode with "wifi disabled" to verify the reported wifi state.
-        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_WITH_WIFI_OFF_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_WITH_WIFI_OFF_MODE, null);
         mLooper.dispatchAll();
         assertEquals(WifiStateMachine.SCAN_ONLY_WITH_WIFI_OFF_MODE,
                      mWsm.getOperationalModeForTest());
@@ -854,7 +854,7 @@ public class WifiStateMachineTest {
      * Helper method to move through startup states.
      */
     private void startSupplicantAndDispatchMessages() throws Exception {
-        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE, null);
         mLooper.dispatchAll();
 
         // this will be removed when interface management is dynamic
@@ -1043,7 +1043,7 @@ public class WifiStateMachineTest {
     public void verifyWifiStateTrackerUpdatedWhenDisabled() throws Exception {
         connect();
 
-        mWsm.setOperationalMode(WifiStateMachine.DISABLED_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.DISABLED_MODE, null);
         mLooper.dispatchAll();
         verify(mWifiStateTracker).updateState(eq(WifiStateTracker.DISCONNECTED));
     }
@@ -2015,7 +2015,7 @@ public class WifiStateMachineTest {
 
         // Set WSM to SCAN_ONLY_MODE, verify state and wifi disabled in ConnectivityManager, and
         // WifiInfo is reset() and state set to DISCONNECTED
-        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_MODE);
+        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_MODE, null);
         mLooper.dispatchAll();
 
         assertEquals(WifiStateMachine.SCAN_ONLY_MODE, mWsm.getOperationalModeForTest());