Merge "[DO NOT MERGE]wifi: scan enhancement" into udc-d1-dev

wifi/java/src/android/net/wifi/nl80211/InstantWifi.java

@@ -0,0 +1,318 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.net.wifi.nl80211;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.app.AlarmManager;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.net.ConnectivityManager;
+import android.net.ConnectivityManager.NetworkCallback;
+import android.net.Network;
+import android.net.NetworkCapabilities;
+import android.net.NetworkRequest;
+import android.net.wifi.WifiConfiguration;
+import android.net.wifi.WifiInfo;
+import android.net.wifi.WifiManager;
+import android.os.Handler;
+import android.os.PowerManager;
+import android.os.SystemClock;
+import android.util.Log;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * @hide
+ */
+public class InstantWifi {
+    private static final String INSTANT_WIFI_TAG = "InstantWifi";
+    private static final int OVERRIDED_SCAN_CONNECTION_TIMEOUT_MS = 1000;
+    private static final int WIFI_NETWORK_EXPIRED_MS = 7 * 24 * 60 * 60 * 1000; // a week
+    private static final String NO_CONNECTION_TIMEOUT_ALARM_TAG =
+            INSTANT_WIFI_TAG + " No Connection Timeout";
+
+    private Context mContext;
+    private AlarmManager mAlarmManager;
+    private Handler mEventHandler;
+    private ConnectivityManager mConnectivityManager;
+    private WifiManager mWifiManager;
+    private PowerManager mPowerManager;
+    private long mLastWifiOnSinceBootMs;
+    private long mLastScreenOnSinceBootMs;
+    private boolean mIsWifiConnected = false;
+    private boolean mScreenOn = false;
+    private boolean mWifiEnabled = false;
+    private boolean mIsNoConnectionAlarmSet = false;
+    private ArrayList<WifiNetwork> mConnectedWifiNetworkList = new ArrayList<>();
+    private AlarmManager.OnAlarmListener mNoConnectionTimeoutCallback =
+            new AlarmManager.OnAlarmListener() {
+                public void onAlarm() {
+                    Log.i(INSTANT_WIFI_TAG, "Timed out waiting for wifi connection");
+                    mIsNoConnectionAlarmSet = false;
+                    mWifiManager.startScan();
+                }
+            };
+
+    public InstantWifi(Context context, AlarmManager alarmManager, Handler eventHandler) {
+        mContext = context;
+        mAlarmManager = alarmManager;
+        mEventHandler = eventHandler;
+        mWifiManager = mContext.getSystemService(WifiManager.class);
+        mConnectivityManager = mContext.getSystemService(ConnectivityManager.class);
+        mConnectivityManager.registerNetworkCallback(
+                new NetworkRequest.Builder()
+                .addCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                .addTransportType(NetworkCapabilities.TRANSPORT_WIFI)
+                .build(), new WifiNetworkCallback());
+        // System power service was initialized before wifi nl80211 service.
+        mPowerManager = mContext.getSystemService(PowerManager.class);
+        IntentFilter screenEventfilter = new IntentFilter();
+        screenEventfilter.addAction(Intent.ACTION_SCREEN_ON);
+        screenEventfilter.addAction(Intent.ACTION_SCREEN_OFF);
+        mContext.registerReceiver(
+                new BroadcastReceiver() {
+                @Override
+                    public void onReceive(Context context, Intent intent) {
+                        String action = intent.getAction();
+                        if (action.equals(Intent.ACTION_SCREEN_ON)) {
+                            if (!mScreenOn) {
+                                mLastScreenOnSinceBootMs = getMockableElapsedRealtime();
+                            }
+                            mScreenOn = true;
+                        } else if (action.equals(Intent.ACTION_SCREEN_OFF)) {
+                            mScreenOn = false;
+                        }
+                        Log.d(INSTANT_WIFI_TAG, "mScreenOn is changed to " + mScreenOn);
+                    }
+                }, screenEventfilter, null, mEventHandler);
+        mScreenOn = mPowerManager.isInteractive();
+        mContext.registerReceiver(
+                new BroadcastReceiver() {
+                    @Override
+                    public void onReceive(Context context, Intent intent) {
+                        int state = intent.getIntExtra(WifiManager.EXTRA_WIFI_STATE,
+                                WifiManager.WIFI_STATE_UNKNOWN);
+                        mWifiEnabled = state == WifiManager.WIFI_STATE_ENABLED;
+                        if (mWifiEnabled) {
+                            mLastWifiOnSinceBootMs = getMockableElapsedRealtime();
+                        }
+                        Log.d(INSTANT_WIFI_TAG, "mWifiEnabled is changed to " + mWifiEnabled);
+                    }
+                },
+                new IntentFilter(WifiManager.WIFI_STATE_CHANGED_ACTION),
+                null, mEventHandler);
+    }
+
+    @VisibleForTesting
+    protected long getMockableElapsedRealtime() {
+        return SystemClock.elapsedRealtime();
+    }
+
+    private class WifiNetwork {
+        private final int mNetId;
+        private Set<Integer> mConnectedFrequencies = new HashSet<Integer>();
+        private int[] mLastTwoConnectedFrequencies = new int[2];
+        private long mLastConnectedTimeMillis;
+        WifiNetwork(int netId) {
+            mNetId = netId;
+        }
+
+        public int getNetId() {
+            return mNetId;
+        }
+
+        public boolean addConnectedFrequency(int channelFrequency) {
+            mLastConnectedTimeMillis = getMockableElapsedRealtime();
+            if (mLastTwoConnectedFrequencies[0] != channelFrequency
+                    && mLastTwoConnectedFrequencies[1] != channelFrequency) {
+                mLastTwoConnectedFrequencies[0] = mLastTwoConnectedFrequencies[1];
+                mLastTwoConnectedFrequencies[1] = channelFrequency;
+            }
+            return mConnectedFrequencies.add(channelFrequency);
+        }
+
+        public Set<Integer> getConnectedFrequencies() {
+            return mConnectedFrequencies;
+        }
+
+        public int[] getLastTwoConnectedFrequencies() {
+            if ((getMockableElapsedRealtime() - mLastConnectedTimeMillis)
+                    > WIFI_NETWORK_EXPIRED_MS) {
+                return new int[0];
+            }
+            return mLastTwoConnectedFrequencies;
+        }
+
+        public long getLastConnectedTimeMillis() {
+            return mLastConnectedTimeMillis;
+        }
+    }
+
+    private class WifiNetworkCallback extends NetworkCallback {
+        @Override
+        public void onAvailable(@NonNull Network network) {
+        }
+
+        @Override
+        public void onCapabilitiesChanged(Network network,
+                NetworkCapabilities networkCapabilities) {
+            if (networkCapabilities != null && network != null) {
+                WifiInfo wifiInfo = (WifiInfo) networkCapabilities.getTransportInfo();
+                if (wifiInfo == null || mWifiManager == null) {
+                    return;
+                }
+                WifiConfiguration config = mWifiManager.getPrivilegedConnectedNetwork();
+                if (config == null) {
+                    return;
+                }
+                final int currentNetworkId = config.networkId;
+                final int connectecFrequency = wifiInfo.getFrequency();
+                if (connectecFrequency < 0 || currentNetworkId < 0) {
+                    return;
+                }
+                mIsWifiConnected = true;
+                if (mIsNoConnectionAlarmSet) {
+                    mAlarmManager.cancel(mNoConnectionTimeoutCallback);
+                }
+                Log.d(INSTANT_WIFI_TAG, "Receive Wifi is connected, freq =  " + connectecFrequency
+                        + " and currentNetworkId : " + currentNetworkId
+                        + ", wifiinfo = " + wifiInfo);
+                boolean isExist = false;
+                for (WifiNetwork wifiNetwork : mConnectedWifiNetworkList) {
+                    if (wifiNetwork.getNetId() == currentNetworkId) {
+                        if (wifiNetwork.addConnectedFrequency(connectecFrequency)) {
+                            Log.d(INSTANT_WIFI_TAG, "Update connected frequency: "
+                                    + connectecFrequency + " to Network currentNetworkId : "
+                                    + currentNetworkId);
+                        }
+                        isExist = true;
+                    }
+                }
+                if (!isExist) {
+                    WifiNetwork currentNetwork = new WifiNetwork(currentNetworkId);
+                    currentNetwork.addConnectedFrequency(connectecFrequency);
+                    if (mConnectedWifiNetworkList.size() < 5) {
+                        mConnectedWifiNetworkList.add(currentNetwork);
+                    } else {
+                        ArrayList<WifiNetwork> lastConnectedWifiNetworkList = new ArrayList<>();
+                        WifiNetwork legacyNetwork = mConnectedWifiNetworkList.get(0);
+                        for (WifiNetwork connectedNetwork : mConnectedWifiNetworkList) {
+                            if (connectedNetwork.getNetId() == legacyNetwork.getNetId()) {
+                                continue;
+                            }
+                            // Keep the used recently network in the last connected list
+                            if (connectedNetwork.getLastConnectedTimeMillis()
+                                    > legacyNetwork.getLastConnectedTimeMillis()) {
+                                lastConnectedWifiNetworkList.add(connectedNetwork);
+                            } else {
+                                lastConnectedWifiNetworkList.add(legacyNetwork);
+                                legacyNetwork = connectedNetwork;
+                            }
+                        }
+                        mConnectedWifiNetworkList = lastConnectedWifiNetworkList;
+                    }
+                }
+            }
+        }
+
+        @Override
+        public void onLost(@NonNull Network network) {
+            mIsWifiConnected = false;
+        }
+    }
+
+    /**
+     * Returns whether or not the scan freqs should be overrided by using predicted channels.
+     */
+    public boolean isUsePredictedScanningChannels() {
+        if (mIsWifiConnected || mConnectedWifiNetworkList.size() == 0
+                || !mWifiManager.isWifiEnabled() || !mPowerManager.isInteractive()) {
+            return false;
+        }
+        if (!mWifiEnabled || !mScreenOn) {
+            Log.d(INSTANT_WIFI_TAG, "WiFi/Screen State mis-match, run instant Wifi anyway!");
+            return true;
+        }
+        return (((getMockableElapsedRealtime() - mLastWifiOnSinceBootMs)
+                        < OVERRIDED_SCAN_CONNECTION_TIMEOUT_MS)
+                || ((getMockableElapsedRealtime() - mLastScreenOnSinceBootMs)
+                        < OVERRIDED_SCAN_CONNECTION_TIMEOUT_MS));
+    }
+
+    /**
+     * Overrides the frequenies in SingleScanSetting
+     *
+     * @param settings the SingleScanSettings will be overrided.
+     * @param freqs new frequencies of SingleScanSettings
+     */
+    @Nullable
+    public void overrideFreqsForSingleScanSettingsIfNecessary(
+            @Nullable SingleScanSettings settings, @Nullable Set<Integer> freqs) {
+        if (!isUsePredictedScanningChannels() || settings == null || freqs == null
+                || freqs.size() == 0) {
+            return;
+        }
+        if (settings.channelSettings == null) {
+            settings.channelSettings = new ArrayList<>();
+        } else {
+            settings.channelSettings.clear();
+        }
+        for (int freq : freqs) {
+            if (freq > 0) {
+                ChannelSettings channel = new ChannelSettings();
+                channel.frequency = freq;
+                settings.channelSettings.add(channel);
+            }
+        }
+        // Monitor connection after last override scan request.
+        if (mIsNoConnectionAlarmSet) {
+            mAlarmManager.cancel(mNoConnectionTimeoutCallback);
+        }
+        mAlarmManager.set(AlarmManager.ELAPSED_REALTIME_WAKEUP,
+                getMockableElapsedRealtime() + OVERRIDED_SCAN_CONNECTION_TIMEOUT_MS,
+                NO_CONNECTION_TIMEOUT_ALARM_TAG, mNoConnectionTimeoutCallback, mEventHandler);
+        mIsNoConnectionAlarmSet = true;
+    }
+
+    /**
+     * Returns the predicted scanning chcnnels set.
+     */
+    @NonNull
+    public Set<Integer> getPredictedScanningChannels() {
+        Set<Integer> predictedScanChannels = new HashSet<>();
+        if (!isUsePredictedScanningChannels()) {
+            Log.d(INSTANT_WIFI_TAG, "Drop, size: " + mConnectedWifiNetworkList.size());
+            return predictedScanChannels;
+        }
+        for (WifiNetwork network : mConnectedWifiNetworkList) {
+            for (int connectedFrequency : network.getLastTwoConnectedFrequencies()) {
+                if (connectedFrequency > 0) {
+                    predictedScanChannels.add(connectedFrequency);
+                    Log.d(INSTANT_WIFI_TAG, "Add channel: " + connectedFrequency
+                            + " to predicted channel");
+                }
+            }
+        }
+        return predictedScanChannels;
+    }
+}

wifi/java/src/android/net/wifi/nl80211/WifiNl80211Manager.java

@@ -105,6 +105,8 @@ public class WifiNl80211Manager {
 
     // Cached wificond binder handlers.
     private IWificond mWificond;
+    private Context mContext;
+    private InstantWifi mInstantWifi;
     private WificondEventHandler mWificondEventHandler = new WificondEventHandler();
     private HashMap<String, IClientInterface> mClientInterfaces = new HashMap<>();
     private HashMap<String, IApInterface> mApInterfaces = new HashMap<>();
@@ -172,6 +174,12 @@ public class WifiNl80211Manager {
         void onPnoRequestFailed();
     }
 
+    /** @hide */
+    @VisibleForTesting
+    protected InstantWifi getInstantWifiMockable() {
+        return mInstantWifi;
+    }
+
     /** @hide */
     @VisibleForTesting
     public class WificondEventHandler extends IWificondEventCallback.Stub {
@@ -419,6 +427,7 @@ public class WifiNl80211Manager {
     public WifiNl80211Manager(Context context) {
         mAlarmManager = context.getSystemService(AlarmManager.class);
         mEventHandler = new Handler(context.getMainLooper());
+        mContext = context;
     }
 
     /**
@@ -434,6 +443,7 @@ public class WifiNl80211Manager {
         if (mWificond == null) {
             Log.e(TAG, "Failed to get reference to wificond");
         }
+        mContext = context;
     }
 
     /** @hide */
@@ -441,6 +451,7 @@ public class WifiNl80211Manager {
     public WifiNl80211Manager(Context context, IWificond wificond) {
         this(context);
         mWificond = wificond;
+        mContext = context;
     }
 
     /** @hide */
@@ -744,6 +755,9 @@ public class WifiNl80211Manager {
             Log.e(TAG, "Failed to refresh wificond scanner due to remote exception");
         }
 
+        if (getInstantWifiMockable() == null) {
+            mInstantWifi = new InstantWifi(mContext, mAlarmManager, mEventHandler);
+        }
         return true;
     }
 
@@ -1071,6 +1085,10 @@ public class WifiNl80211Manager {
         if (settings == null) {
             return false;
         }
+        if (getInstantWifiMockable() != null) {
+            getInstantWifiMockable().overrideFreqsForSingleScanSettingsIfNecessary(settings,
+                    getInstantWifiMockable().getPredictedScanningChannels());
+        }
         try {
             return scannerImpl.scan(settings);
         } catch (RemoteException e1) {
@@ -1115,6 +1133,10 @@ public class WifiNl80211Manager {
         if (settings == null) {
             return WifiScanner.REASON_INVALID_ARGS;
         }
+        if (getInstantWifiMockable() != null) {
+            getInstantWifiMockable().overrideFreqsForSingleScanSettingsIfNecessary(settings,
+                    getInstantWifiMockable().getPredictedScanningChannels());
+        }
         try {
             int status = scannerImpl.scanRequest(settings);
             return toFrameworkScanStatusCode(status);

wifi/tests/src/android/net/wifi/nl80211/InstantWifiTest.java

@@ -0,0 +1,256 @@
+/*
+ * Copyright (C) 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package android.net.wifi.nl80211;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.any;
+import static org.mockito.Mockito.anyInt;
+import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import android.app.AlarmManager;
+import android.app.test.TestAlarmManager;
+import android.content.BroadcastReceiver;
+import android.content.Context;
+import android.content.Intent;
+import android.content.IntentFilter;
+import android.net.ConnectivityManager;
+import android.net.ConnectivityManager.NetworkCallback;
+import android.net.Network;
+import android.net.NetworkCapabilities;
+import android.net.wifi.WifiConfiguration;
+import android.net.wifi.WifiInfo;
+import android.net.wifi.WifiManager;
+import android.os.Handler;
+import android.os.IPowerManager;
+import android.os.IThermalService;
+import android.os.PowerManager;
+import android.os.test.TestLooper;
+
+import androidx.test.filters.SmallTest;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.HashSet;
+import java.util.Set;
+
+/**
+ * Unit tests for {@link android.net.wifi.nl80211.InstantWifi}.
+ */
+@SmallTest
+public class InstantWifiTest {
+    @Mock private Context mContext;
+    @Mock private ConnectivityManager mMockConnectivityManager;
+    @Mock private WifiManager mMockWifiManager;
+    @Mock private Network mMockWifiNetwork;
+    @Mock private WifiInfo mMockWifiInfo;
+    @Mock private WifiConfiguration mMockWifiConfiguration;
+    @Mock private IPowerManager mPowerManagerService;
+    private InstantWifi mInstantWifi;
+    private TestLooper mLooper;
+    private Handler mHandler;
+    private TestAlarmManager mTestAlarmManager;
+    private AlarmManager mAlarmManager;
+    private PowerManager mMockPowerManager;
+
+    private final ArgumentCaptor<NetworkCallback> mWifiNetworkCallbackCaptor =
+            ArgumentCaptor.forClass(NetworkCallback.class);
+    private final ArgumentCaptor<BroadcastReceiver> mScreenBroadcastReceiverCaptor =
+            ArgumentCaptor.forClass(BroadcastReceiver.class);
+    private final ArgumentCaptor<BroadcastReceiver> mWifiStateBroadcastReceiverCaptor =
+            ArgumentCaptor.forClass(BroadcastReceiver.class);
+
+    private static final int TEST_NETWORK_ID = 1;
+    private static final int TEST_24G_FREQUENCY = 2412;
+    private static final int TEST_5G_FREQUENCY = 5745;
+    private long mTimeOffsetMs = 0;
+
+    private class InstantWifiSpy extends InstantWifi {
+        InstantWifiSpy(Context context, AlarmManager alarmManager, Handler handler) {
+            super(context, alarmManager, handler);
+        }
+
+        @Override
+        protected long getMockableElapsedRealtime() {
+            return mTimeOffsetMs;
+        }
+    }
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        mLooper = new TestLooper();
+        mHandler = new Handler(mLooper.getLooper());
+
+        mTestAlarmManager = new TestAlarmManager();
+        mAlarmManager = mTestAlarmManager.getAlarmManager();
+        when(mContext.getSystemServiceName(AlarmManager.class)).thenReturn(Context.ALARM_SERVICE);
+        when(mContext.getSystemService(AlarmManager.class)).thenReturn(mAlarmManager);
+        when(mContext.getSystemServiceName(WifiManager.class)).thenReturn(Context.WIFI_SERVICE);
+        when(mContext.getSystemService(WifiManager.class)).thenReturn(mMockWifiManager);
+        when(mContext.getSystemServiceName(ConnectivityManager.class))
+                .thenReturn(Context.CONNECTIVITY_SERVICE);
+        when(mContext.getSystemService(ConnectivityManager.class))
+                .thenReturn(mMockConnectivityManager);
+        mMockPowerManager = new PowerManager(mContext, mPowerManagerService,
+                mock(IThermalService.class), mHandler);
+        when(mContext.getSystemServiceName(PowerManager.class)).thenReturn(Context.POWER_SERVICE);
+        when(mContext.getSystemService(PowerManager.class)).thenReturn(mMockPowerManager);
+        when(mPowerManagerService.isInteractive()).thenReturn(true);
+
+        doReturn(mMockWifiInfo).when(mMockWifiInfo).makeCopy(anyLong());
+        mTimeOffsetMs = 0;
+        mInstantWifi = new InstantWifiSpy(mContext, mAlarmManager, mHandler);
+        verifyInstantWifiInitialization();
+    }
+
+    private void verifyInstantWifiInitialization() {
+        verify(mMockConnectivityManager).registerNetworkCallback(any(),
+                mWifiNetworkCallbackCaptor.capture());
+        verify(mContext).registerReceiver(mScreenBroadcastReceiverCaptor.capture(),
+                argThat((IntentFilter filter) ->
+                        filter.hasAction(Intent.ACTION_SCREEN_ON)
+                                && filter.hasAction(Intent.ACTION_SCREEN_OFF)), eq(null), any());
+
+        verify(mContext).registerReceiver(mWifiStateBroadcastReceiverCaptor.capture(),
+                argThat((IntentFilter filter) ->
+                        filter.hasAction(WifiManager.WIFI_STATE_CHANGED_ACTION)), eq(null), any());
+    }
+
+    private void mockWifiConnectedEvent(int networkId, int connectedFrequency) {
+        // Send wifi connected event
+        NetworkCapabilities mockWifiNetworkCapabilities =
+                new NetworkCapabilities.Builder().setTransportInfo(mMockWifiInfo).build();
+        mMockWifiConfiguration.networkId = networkId;
+        when(mMockWifiManager.getPrivilegedConnectedNetwork()).thenReturn(mMockWifiConfiguration);
+        when(mMockWifiInfo.getFrequency()).thenReturn(connectedFrequency);
+        mWifiNetworkCallbackCaptor.getValue().onCapabilitiesChanged(mMockWifiNetwork,
+                mockWifiNetworkCapabilities);
+        mLooper.dispatchAll();
+    }
+
+    private void mockWifiOnScreenOnBroadcast(boolean isWifiOn, boolean isScreenOn)
+            throws Exception {
+        // Send Wifi On broadcast
+        Intent wifiOnIntent = new Intent(WifiManager.WIFI_STATE_CHANGED_ACTION);
+        wifiOnIntent.putExtra(WifiManager.EXTRA_WIFI_STATE,
+                isWifiOn ? WifiManager.WIFI_STATE_ENABLED : WifiManager.WIFI_STATE_DISABLED);
+        mWifiStateBroadcastReceiverCaptor.getValue().onReceive(mContext, wifiOnIntent);
+        mLooper.dispatchAll();
+        // Send Screen On broadcast
+        Intent screenOnIntent =
+                new Intent(isScreenOn ? Intent.ACTION_SCREEN_ON : Intent.ACTION_SCREEN_OFF);
+        mScreenBroadcastReceiverCaptor.getValue().onReceive(mContext, screenOnIntent);
+        mLooper.dispatchAll();
+        when(mMockWifiManager.isWifiEnabled()).thenReturn(isWifiOn);
+        when(mPowerManagerService.isInteractive()).thenReturn(isScreenOn);
+    }
+
+    @Test
+    public void testisUsePredictedScanningChannels() throws Exception {
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+        mockWifiOnScreenOnBroadcast(true /* isWifiOn */, false /* isScreenOn */);
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+        mockWifiOnScreenOnBroadcast(false /* isWifiOn */, true /* isScreenOn */);
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+        mockWifiOnScreenOnBroadcast(true /* isWifiOn */, true /* isScreenOn */);
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+        // Send wifi connected event
+        mockWifiConnectedEvent(TEST_NETWORK_ID, TEST_24G_FREQUENCY);
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+        // Send wifi disconnect
+        mWifiNetworkCallbackCaptor.getValue().onLost(mMockWifiNetwork);
+        assertTrue(mInstantWifi.isUsePredictedScanningChannels());
+        // Shift time to make it expired
+        mTimeOffsetMs = 1100;
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+    }
+
+    @Test
+    public void testGetPredictedScanningChannels() throws Exception {
+        mockWifiOnScreenOnBroadcast(true /* isWifiOn */, true /* isScreenOn */);
+        // Send wifi connected event on T0
+        mockWifiConnectedEvent(TEST_NETWORK_ID, TEST_24G_FREQUENCY);
+        // Send wifi disconnect
+        mWifiNetworkCallbackCaptor.getValue().onLost(mMockWifiNetwork);
+        assertTrue(mInstantWifi.isUsePredictedScanningChannels());
+        assertTrue(mInstantWifi.getPredictedScanningChannels().contains(TEST_24G_FREQUENCY));
+        mTimeOffsetMs += 1000; // T1 = 1000 ms
+        // Send wifi connected event
+        mockWifiConnectedEvent(TEST_NETWORK_ID + 1, TEST_5G_FREQUENCY);
+        // Send wifi disconnect
+        mWifiNetworkCallbackCaptor.getValue().onLost(mMockWifiNetwork);
+        // isUsePredictedScanningChannels is false since wifi on & screen on is expired
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+        // Override the Wifi On & Screen on time
+        mockWifiOnScreenOnBroadcast(true /* isWifiOn */, true /* isScreenOn */);
+        assertTrue(mInstantWifi.getPredictedScanningChannels().contains(TEST_5G_FREQUENCY));
+        mTimeOffsetMs += 7 * 24 * 60 * 60 * 1000; // Make T0 expired
+        // Override the Wifi On & Screen on time
+        mockWifiOnScreenOnBroadcast(true /* isWifiOn */, true /* isScreenOn */);
+        assertFalse(mInstantWifi.getPredictedScanningChannels().contains(TEST_24G_FREQUENCY));
+        assertTrue(mInstantWifi.getPredictedScanningChannels().contains(TEST_5G_FREQUENCY));
+    }
+
+    @Test
+    public void testOverrideFreqsForSingleScanSettings() throws Exception {
+        mockWifiOnScreenOnBroadcast(true /* isWifiOn */, true /* isScreenOn */);
+        // Send wifi connected event
+        mockWifiConnectedEvent(TEST_NETWORK_ID, TEST_24G_FREQUENCY);
+        assertFalse(mInstantWifi.isUsePredictedScanningChannels());
+        // Send wifi disconnect
+        mWifiNetworkCallbackCaptor.getValue().onLost(mMockWifiNetwork);
+        assertTrue(mInstantWifi.isUsePredictedScanningChannels());
+
+        final ArgumentCaptor<AlarmManager.OnAlarmListener> alarmListenerCaptor =
+                ArgumentCaptor.forClass(AlarmManager.OnAlarmListener.class);
+        doNothing().when(mAlarmManager).set(anyInt(), anyLong(), any(),
+                alarmListenerCaptor.capture(), any());
+        Set<Integer> testFreqs = Set.of(
+                TEST_24G_FREQUENCY, TEST_5G_FREQUENCY);
+        SingleScanSettings testSingleScanSettings = new SingleScanSettings();
+        mInstantWifi.overrideFreqsForSingleScanSettingsIfNecessary(
+                testSingleScanSettings, new HashSet<Integer>());
+        mInstantWifi.overrideFreqsForSingleScanSettingsIfNecessary(
+                testSingleScanSettings, null);
+        mInstantWifi.overrideFreqsForSingleScanSettingsIfNecessary(null, null);
+        verify(mAlarmManager, never()).set(anyInt(), anyLong(), any(), any(), any());
+        mInstantWifi.overrideFreqsForSingleScanSettingsIfNecessary(testSingleScanSettings,
+                testFreqs);
+        verify(mAlarmManager).set(anyInt(), anyLong(), any(), any(), any());
+        Set<Integer> overridedFreqs = new HashSet<Integer>();
+        for (ChannelSettings channel : testSingleScanSettings.channelSettings) {
+            overridedFreqs.add(channel.frequency);
+        }
+        assertEquals(testFreqs, overridedFreqs);
+        alarmListenerCaptor.getValue().onAlarm();
+        verify(mMockWifiManager).startScan();
+    }
+}

wifi/tests/src/android/net/wifi/nl80211/WifiNl80211ManagerTest.java

@@ -18,6 +18,7 @@ package android.net.wifi.nl80211;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
@@ -26,6 +27,7 @@ import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Matchers.argThat;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyLong;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
@@ -37,6 +39,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 import android.app.AlarmManager;
+import android.app.test.MockAnswerUtil.AnswerWithArguments;
 import android.app.test.TestAlarmManager;
 import android.content.Context;
 import android.net.MacAddress;
@@ -104,6 +107,9 @@ public class WifiNl80211ManagerTest {
     private WifiNl80211Manager.CountryCodeChangedListener mCountryCodeChangedListener2;
     @Mock
     private Context mContext;
+    @Mock
+    private InstantWifi mMockInstantWifi;
+
     private TestLooper mLooper;
     private TestAlarmManager mTestAlarmManager;
     private AlarmManager mAlarmManager;
@@ -167,6 +173,17 @@ public class WifiNl80211ManagerTest {
             0x00, 0x00
     };
 
+    private class WifiNl80211ManagerSpy extends WifiNl80211Manager {
+        WifiNl80211ManagerSpy(Context context, IWificond wificond) {
+            super(context, wificond);
+        }
+
+        @Override
+        protected InstantWifi getInstantWifiMockable() {
+            return mMockInstantWifi;
+        }
+    }
+
     @Before
     public void setUp() throws Exception {
         // Setup mocks for successful WificondControl operation. Failure case mocks should be
@@ -181,6 +198,8 @@ public class WifiNl80211ManagerTest {
         mLooper = new TestLooper();
         when(mContext.getMainLooper()).thenReturn(mLooper.getLooper());
 
+        doNothing().when(mMockInstantWifi).overrideFreqsForSingleScanSettingsIfNecessary(
+                any(), any());
         when(mWificond.asBinder()).thenReturn(mWifiCondBinder);
         when(mClientInterface.getWifiScannerImpl()).thenReturn(mWifiScannerImpl);
         when(mWificond.createClientInterface(any())).thenReturn(mClientInterface);
@@ -189,7 +208,7 @@ public class WifiNl80211ManagerTest {
         when(mWificond.tearDownApInterface(any())).thenReturn(true);
         when(mClientInterface.getWifiScannerImpl()).thenReturn(mWifiScannerImpl);
         when(mClientInterface.getInterfaceName()).thenReturn(TEST_INTERFACE_NAME);
-        mWificondControl = new WifiNl80211Manager(mContext, mWificond);
+        mWificondControl = new WifiNl80211ManagerSpy(mContext, mWificond);
         mWificondEventHandler = mWificondControl.getWificondEventHandler();
         assertEquals(true,
                 mWificondControl.setupInterfaceForClientMode(TEST_INTERFACE_NAME, Runnable::run,
@@ -1159,6 +1178,40 @@ public class WifiNl80211ManagerTest {
         verify(mWificond).notifyCountryCodeChanged();
     }
 
+    @Test
+    public void testInstantWifi() throws Exception {
+        doAnswer(new AnswerWithArguments() {
+            public void answer(SingleScanSettings settings, Set<Integer> freqs) {
+                if (settings.channelSettings == null) {
+                    settings.channelSettings = new ArrayList<>();
+                } else {
+                    settings.channelSettings.clear();
+                }
+                for (int freq : freqs) {
+                    if (freq > 0) {
+                        ChannelSettings channel = new ChannelSettings();
+                        channel.frequency = freq;
+                        settings.channelSettings.add(channel);
+                    }
+                }
+            }
+        }).when(mMockInstantWifi).overrideFreqsForSingleScanSettingsIfNecessary(
+                any(), any());
+        Set<Integer> testPredictedChannelsSet = Set.of(2412, 5745);
+        assertNotEquals(testPredictedChannelsSet, SCAN_FREQ_SET);
+        when(mWifiScannerImpl.scan(any(SingleScanSettings.class))).thenReturn(true);
+        when(mMockInstantWifi.getPredictedScanningChannels()).thenReturn(testPredictedChannelsSet);
+
+        // Trigger scan to check scan settings are changed
+        assertTrue(mWificondControl.startScan(
+                TEST_INTERFACE_NAME, WifiScanner.SCAN_TYPE_LOW_POWER,
+                SCAN_FREQ_SET, SCAN_HIDDEN_NETWORK_SSID_LIST));
+        verify(mMockInstantWifi).getPredictedScanningChannels();
+        verify(mWifiScannerImpl).scan(argThat(new ScanMatcher(
+                IWifiScannerImpl.SCAN_TYPE_LOW_POWER,
+                testPredictedChannelsSet, SCAN_HIDDEN_NETWORK_SSID_LIST, false, null)));
+    }
+
     // Create a ArgumentMatcher which captures a SingleScanSettings parameter and checks if it
     // matches the provided frequency set and ssid set.
     private class ScanMatcher implements ArgumentMatcher<SingleScanSettings> {