WifiStateMachine: update new mode in initial state

When WifiController calls setOperationalMode a CMD_SET_OPERATIONAL_MODE
message is sent to WifiStateMachine.  When this message is received in
the InitialState, it is not handled where it is then dropped in the
default state.  Dropping this message causes WifiController and
WifiStateMachine to get out of sync on the operating mode.

BUG:29938263
Change-Id: I8481af8d6826f3da18f5762833a7b145d81b74dd
TEST:runtest frameworks-wifi
diff --git a/service/java/com/android/server/wifi/WifiStateMachine.java b/service/java/com/android/server/wifi/WifiStateMachine.java
index 0e3ed54..d89594f 100644
--- a/service/java/com/android/server/wifi/WifiStateMachine.java
+++ b/service/java/com/android/server/wifi/WifiStateMachine.java
@@ -1884,6 +1884,14 @@
     }
 
     /**
+     * Allow tests to confirm the operational mode for WSM.
+     */
+    @VisibleForTesting
+    protected int getOperationalModeForTest() {
+        return mOperationalMode;
+    }
+
+    /**
      * TODO: doc
      */
     public List<ScanResult> syncGetScanResultsList() {
@@ -4334,6 +4342,9 @@
                         transitionTo(mInitialState);
                     }
                     break;
+                case CMD_SET_OPERATIONAL_MODE:
+                    mOperationalMode = message.arg1;
+                    break;
                 default:
                     return NOT_HANDLED;
             }
diff --git a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
index e0f94ad..53d8f46 100644
--- a/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WifiStateMachineTest.java
@@ -475,6 +475,75 @@
         assertEquals("InitialState", getCurrentState().getName());
     }
 
+    /**
+     * Test to check that mode changes from WifiController will be properly handled in the
+     * InitialState by WifiStateMachine.
+     */
+    @Test
+    public void checkOperationalModeInInitialState() throws Exception {
+        when(mWifiNative.loadDriver()).thenReturn(true);
+        when(mWifiNative.startHal()).thenReturn(true);
+        when(mWifiNative.startSupplicant(anyBoolean())).thenReturn(true);
+
+        mLooper.dispatchAll();
+        assertEquals("InitialState", getCurrentState().getName());
+        assertEquals(WifiStateMachine.CONNECT_MODE, mWsm.getOperationalModeForTest());
+
+        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_WITH_WIFI_OFF_MODE);
+        mLooper.dispatchAll();
+        assertEquals(WifiStateMachine.SCAN_ONLY_WITH_WIFI_OFF_MODE,
+                mWsm.getOperationalModeForTest());
+
+        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_MODE);
+        mLooper.dispatchAll();
+        assertEquals(WifiStateMachine.SCAN_ONLY_MODE, mWsm.getOperationalModeForTest());
+
+        mWsm.setOperationalMode(WifiStateMachine.CONNECT_MODE);
+        mLooper.dispatchAll();
+        assertEquals(WifiStateMachine.CONNECT_MODE, mWsm.getOperationalModeForTest());
+    }
+
+    /**
+     * Test that mode changes for WifiStateMachine in the InitialState are realized when supplicant
+     * is started.
+     */
+    @Test
+    public void checkStartInCorrectStateAfterChangingInitialState() throws Exception {
+        when(mWifiNative.loadDriver()).thenReturn(true);
+        when(mWifiNative.startHal()).thenReturn(true);
+        when(mWifiNative.startSupplicant(anyBoolean())).thenReturn(true);
+
+        // Check initial state
+        mLooper.dispatchAll();
+        assertEquals("InitialState", getCurrentState().getName());
+        assertEquals(WifiStateMachine.CONNECT_MODE, mWsm.getOperationalModeForTest());
+
+        // Update the mode
+        mWsm.setOperationalMode(WifiStateMachine.SCAN_ONLY_MODE);
+        mLooper.dispatchAll();
+        assertEquals(WifiStateMachine.SCAN_ONLY_MODE, mWsm.getOperationalModeForTest());
+
+        // Start supplicant so we move to the next state
+        mWsm.setSupplicantRunning(true);
+        mLooper.dispatchAll();
+        assertEquals("SupplicantStartingState", getCurrentState().getName());
+        when(mWifiNative.setBand(anyInt())).thenReturn(true);
+        when(mWifiNative.setDeviceName(anyString())).thenReturn(true);
+        when(mWifiNative.setManufacturer(anyString())).thenReturn(true);
+        when(mWifiNative.setModelName(anyString())).thenReturn(true);
+        when(mWifiNative.setModelNumber(anyString())).thenReturn(true);
+        when(mWifiNative.setSerialNumber(anyString())).thenReturn(true);
+        when(mWifiNative.setConfigMethods(anyString())).thenReturn(true);
+        when(mWifiNative.setDeviceType(anyString())).thenReturn(true);
+        when(mWifiNative.setSerialNumber(anyString())).thenReturn(true);
+        when(mWifiNative.setScanningMacOui(any(byte[].class))).thenReturn(true);
+
+        mWsm.sendMessage(WifiMonitor.SUP_CONNECTION_EVENT);
+        mLooper.dispatchAll();
+
+        assertEquals("ScanModeState", getCurrentState().getName());
+    }
+
     private void addNetworkAndVerifySuccess() throws Exception {
         addNetworkAndVerifySuccess(false);
     }