Merge changes I5fd02510,I062f0005 into pi-dev

* changes:
  Change onboarding flow.
  Fix WifiWake locking behavior.

4 files changed, 293 insertions, 79 deletions

diff --git a/tests/wifitests/src/com/android/server/wifi/WakeupConfigStoreDataTest.java b/tests/wifitests/src/com/android/server/wifi/WakeupConfigStoreDataTest.java
index 0f0bfcf7c..e59175c93 100644
--- a/tests/wifitests/src/com/android/server/wifi/WakeupConfigStoreDataTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WakeupConfigStoreDataTest.java
@@ -49,6 +49,7 @@ public class WakeupConfigStoreDataTest {
 
     @Mock private WakeupConfigStoreData.DataSource<Boolean> mActiveDataSource;
     @Mock private WakeupConfigStoreData.DataSource<Boolean> mIsOnboardedDataSource;
+    @Mock private WakeupConfigStoreData.DataSource<Integer> mNotificationsDataSource;
     @Mock private WakeupConfigStoreData.DataSource<Set<ScanResultMatchInfo>> mNetworkDataSource;
 
     private WakeupConfigStoreData mWakeupConfigData;
@@ -58,7 +59,7 @@ public class WakeupConfigStoreDataTest {
         MockitoAnnotations.initMocks(this);
 
         mWakeupConfigData = new WakeupConfigStoreData(mActiveDataSource, mIsOnboardedDataSource,
-                mNetworkDataSource);
+                mNotificationsDataSource, mNetworkDataSource);
     }
 
     /**
@@ -115,9 +116,11 @@ public class WakeupConfigStoreDataTest {
         Set<ScanResultMatchInfo> networks = Collections.emptySet();
         boolean isActive = false;
         boolean isOnboarded = false;
+        int notificationsShown = 0;
 
         when(mActiveDataSource.getData()).thenReturn(isActive);
         when(mIsOnboardedDataSource.getData()).thenReturn(isOnboarded);
+        when(mNotificationsDataSource.getData()).thenReturn(notificationsShown);
         when(mNetworkDataSource.getData()).thenReturn(networks);
 
         byte[] bytes = serializeData(false /* shared */);
@@ -125,6 +128,7 @@ public class WakeupConfigStoreDataTest {
 
         verify(mActiveDataSource).setData(eq(isActive));
         verify(mIsOnboardedDataSource).setData(eq(isOnboarded));
+        verify(mNotificationsDataSource).setData(notificationsShown);
         verify(mNetworkDataSource).setData(eq(networks));
     }
 
@@ -148,9 +152,11 @@ public class WakeupConfigStoreDataTest {
         Set<ScanResultMatchInfo> networks = Sets.newArraySet(network1, network2, network3);
         boolean isActive = true;
         boolean isOnboarded = false;
+        int notificationsShown = 1;
 
         when(mActiveDataSource.getData()).thenReturn(isActive);
         when(mIsOnboardedDataSource.getData()).thenReturn(isOnboarded);
+        when(mNotificationsDataSource.getData()).thenReturn(notificationsShown);
         when(mNetworkDataSource.getData()).thenReturn(networks);
 
         byte[] bytes = serializeData(false /* shared */);
@@ -158,6 +164,7 @@ public class WakeupConfigStoreDataTest {
 
         verify(mActiveDataSource).setData(eq(isActive));
         verify(mIsOnboardedDataSource).setData(eq(isOnboarded));
+        verify(mNotificationsDataSource).setData(notificationsShown);
         verify(mNetworkDataSource).setData(eq(networks));
     }
 
@@ -170,6 +177,7 @@ public class WakeupConfigStoreDataTest {
 
         verify(mActiveDataSource).setData(false);
         verify(mIsOnboardedDataSource).setData(false);
+        verify(mNotificationsDataSource).setData(0);
         verify(mNetworkDataSource).setData(eq(Collections.emptySet()));
     }
 
diff --git a/tests/wifitests/src/com/android/server/wifi/WakeupControllerTest.java b/tests/wifitests/src/com/android/server/wifi/WakeupControllerTest.java
index f246c4884..1aedae286 100644
--- a/tests/wifitests/src/com/android/server/wifi/WakeupControllerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WakeupControllerTest.java
@@ -23,11 +23,11 @@ import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.never;
-import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.content.Context;
+import android.database.ContentObserver;
 import android.net.wifi.ScanResult;
 import android.net.wifi.WifiConfiguration;
 import android.net.wifi.WifiScanner;
@@ -49,14 +49,16 @@ import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.util.Arrays;
-import java.util.Collection;
 import java.util.Collections;
+import java.util.Set;
 
 /**
  * Unit tests for {@link WakeupController}.
  */
 public class WakeupControllerTest {
 
+    private static final int DFS_CHANNEL = 5540;
+
     @Mock private Context mContext;
     @Mock private WakeupLock mWakeupLock;
     @Mock private WakeupEvaluator mWakeupEvaluator;
@@ -69,6 +71,7 @@ public class WakeupControllerTest {
     @Mock private WifiSettingsStore mWifiSettingsStore;
     @Mock private WifiWakeMetrics mWifiWakeMetrics;
     @Mock private WifiController mWifiController;
+    @Mock private WifiNative mWifiNative;
 
     private TestLooper mLooper;
     private WakeupController mWakeupController;
@@ -84,6 +87,10 @@ public class WakeupControllerTest {
         when(mWifiInjector.getWifiScanner()).thenReturn(mWifiScanner);
         when(mWifiInjector.getWifiSettingsStore()).thenReturn(mWifiSettingsStore);
         when(mWifiInjector.getWifiController()).thenReturn(mWifiController);
+        when(mWifiInjector.getWifiNative()).thenReturn(mWifiNative);
+        when(mWifiNative.getChannelsForBand(WifiScanner.WIFI_BAND_5_GHZ_DFS_ONLY))
+                .thenReturn(new int[]{DFS_CHANNEL});
+
         when(mWifiSettingsStore.handleWifiToggled(anyBoolean())).thenReturn(true);
 
         mLooper = new TestLooper();
@@ -92,6 +99,7 @@ public class WakeupControllerTest {
         mTestScanResult = new ScanResult();
         mTestScanResult.SSID = "open ssid 1";
         mTestScanResult.capabilities = "";
+        mTestScanResult.frequency = 2412;
         ScanResult[] scanResults = new ScanResult[1];
         scanResults[0] = mTestScanResult;
         mTestScanDatas = new WifiScanner.ScanData[1];
@@ -191,43 +199,43 @@ public class WakeupControllerTest {
     }
 
     /**
-     * Verify that start initializes the wakeup lock.
+     * Verify that start sets the wakeup lock.
      */
     @Test
-    public void startInitializesWakeupLock() {
+    public void startSetsWakeupLock() {
         initializeWakeupController(true /* enabled */);
         mWakeupController.start();
-        verify(mWakeupLock).initialize(any());
+        verify(mWakeupLock).setLock(any());
         verify(mWifiWakeMetrics).recordStartEvent(anyInt());
     }
 
     /**
-     * Verify that start does not initialize the wakeup lock when feature is disabled.
+     * Verify that start does not set the wakeup lock when feature is disabled.
      */
     @Test
-    public void startDoesNotInitializeWakeupLockWhenDisabled() {
+    public void startDoesNotSetWakeupLockWhenDisabled() {
         initializeWakeupController(false /* enabled */);
         mWakeupController.start();
-        verify(mWakeupLock, never()).initialize(any());
+        verify(mWakeupLock, never()).setLock(any());
         verify(mWifiWakeMetrics, never()).recordStartEvent(anyInt());
     }
 
     /**
-     * Verify that start does not re-initialize the wakeup lock if the controller is already active.
+     * Verify that start does not set the wakeup lock if the controller is already active.
      */
     @Test
-    public void startDoesNotInitializeWakeupLockIfAlreadyActive() {
+    public void startDoesNotSetWakeupLockIfAlreadyActive() {
         initializeWakeupController(true /* enabled */);
         InOrder lockInOrder = Mockito.inOrder(mWakeupLock);
         InOrder metricsInOrder = Mockito.inOrder(mWifiWakeMetrics);
 
         mWakeupController.start();
-        lockInOrder.verify(mWakeupLock).initialize(any());
+        lockInOrder.verify(mWakeupLock).setLock(any());
         metricsInOrder.verify(mWifiWakeMetrics).recordStartEvent(anyInt());
 
         mWakeupController.stop();
         mWakeupController.start();
-        lockInOrder.verify(mWakeupLock, never()).initialize(any());
+        lockInOrder.verify(mWakeupLock, never()).setLock(any());
         metricsInOrder.verify(mWifiWakeMetrics, never()).recordStartEvent(anyInt());
     }
 
@@ -256,7 +264,7 @@ public class WakeupControllerTest {
      * Verify that reset sets active to false.
      *
      * <p>This is accomplished by initiating another call to start and verifying that the wakeup
-     * lock is re-initialized.
+     * lock is re-set.
      */
     @Test
     public void resetSetsActiveToFalse() {
@@ -265,7 +273,7 @@ public class WakeupControllerTest {
         InOrder metricsInOrder = Mockito.inOrder(mWifiWakeMetrics);
 
         mWakeupController.start();
-        lockInOrder.verify(mWakeupLock).initialize(any());
+        lockInOrder.verify(mWakeupLock).setLock(any());
         metricsInOrder.verify(mWifiWakeMetrics).recordStartEvent(anyInt());
 
         mWakeupController.stop();
@@ -273,7 +281,7 @@ public class WakeupControllerTest {
         metricsInOrder.verify(mWifiWakeMetrics).recordResetEvent(0 /* numScans */);
 
         mWakeupController.start();
-        lockInOrder.verify(mWakeupLock).initialize(any());
+        lockInOrder.verify(mWakeupLock).setLock(any());
         metricsInOrder.verify(mWifiWakeMetrics).recordStartEvent(anyInt());
     }
 
@@ -297,17 +305,49 @@ public class WakeupControllerTest {
 
         // scan results from most recent scan
         ScanResult savedScanResult = createOpenScanResult(ssid1);
+        savedScanResult.frequency = 2412;
         ScanResult unsavedScanResult = createOpenScanResult(ssid2);
+        savedScanResult.frequency = 2412;
+
         when(mWifiScanner.getSingleScanResults())
                 .thenReturn(Arrays.asList(savedScanResult, unsavedScanResult));
 
         // intersection of most recent scan + saved configs
-        Collection<ScanResultMatchInfo> expectedMatchInfos =
+        Set<ScanResultMatchInfo> expectedMatchInfos =
                 Collections.singleton(ScanResultMatchInfo.fromScanResult(savedScanResult));
 
         initializeWakeupController(true /* enabled */);
         mWakeupController.start();
-        verify(mWakeupLock).initialize(eq(expectedMatchInfos));
+        verify(mWakeupLock).setLock(eq(expectedMatchInfos));
+        verify(mWifiWakeMetrics).recordStartEvent(expectedMatchInfos.size());
+    }
+
+    /**
+     * Verify that start filters out DFS channels.
+     */
+    @Test
+    public void startFiltersOutDfsScanResults() {
+        String ssid = "test_ssid";
+        String quotedSsid = ScanResultUtil.createQuotedSSID(ssid);
+
+        // saved config
+        WifiConfiguration openNetwork = WifiConfigurationTestUtil.createOpenNetwork(quotedSsid);
+        openNetwork.getNetworkSelectionStatus().setHasEverConnected(true);
+        when(mWifiConfigManager.getSavedNetworks())
+                .thenReturn(Collections.singletonList(openNetwork));
+
+        // scan result from most recent scan
+        ScanResult scanResult = createOpenScanResult(ssid);
+        scanResult.frequency = DFS_CHANNEL;
+
+        when(mWifiScanner.getSingleScanResults()).thenReturn(Collections.singletonList(scanResult));
+
+        // intersection of most recent scan + saved configs
+        Set<ScanResultMatchInfo> expectedMatchInfos = Collections.emptySet();
+
+        initializeWakeupController(true /* enabled */);
+        mWakeupController.start();
+        verify(mWakeupLock).setLock(eq(expectedMatchInfos));
         verify(mWifiWakeMetrics).recordStartEvent(expectedMatchInfos.size());
     }
 
@@ -333,16 +373,13 @@ public class WakeupControllerTest {
     }
 
     /**
-     * Verify that WifiWakeMetrics logs the unlock event when the WakeupLock is emptied.
+     * Verify that onResults filters out DFS channels.
      */
     @Test
-    public void recordMetricsWhenWakeupLockIsEmptied() {
+    public void onResultsFiltersOutDfsScanResults() {
         initializeWakeupController(true /* enabled */);
         mWakeupController.start();
 
-        // Simulates emptying the lock: first returns false then returns true
-        when(mWakeupLock.isEmpty()).thenReturn(false).thenReturn(true);
-
         ArgumentCaptor<WifiScanner.ScanListener> scanListenerArgumentCaptor =
                 ArgumentCaptor.forClass(WifiScanner.ScanListener.class);
 
@@ -350,20 +387,20 @@ public class WakeupControllerTest {
         WifiScanner.ScanListener scanListener = scanListenerArgumentCaptor.getValue();
 
         // incoming scan results
+        mTestScanResult.frequency = DFS_CHANNEL;
         scanListener.onResults(mTestScanDatas);
 
-        verify(mWakeupLock, times(2)).isEmpty();
-        verify(mWifiWakeMetrics).recordUnlockEvent(1 /* numScans */);
+        verify(mWakeupLock).update(eq(Collections.emptySet()));
     }
 
     /**
      * Verify that the controller searches for viable networks during onResults when WakeupLock is
-     * empty.
+     * unlocked.
      */
     @Test
-    public void onResultsSearchesForViableNetworkWhenWakeupLockIsEmpty() {
-        // empty wakeup lock
-        when(mWakeupLock.isEmpty()).thenReturn(true);
+    public void onResultsSearchesForViableNetworkWhenWakeupLockIsUnlocked() {
+        // unlock wakeup lock
+        when(mWakeupLock.isUnlocked()).thenReturn(true);
         // do not find viable network
         when(mWakeupEvaluator.findViableNetwork(any(), any())).thenReturn(null);
 
@@ -384,12 +421,13 @@ public class WakeupControllerTest {
     }
 
     /**
-     * Verify that the controller only updates the WakeupLock if the user is onboarded.
+     * Verify that the controller updates the WakeupLock even if the user is not onboarded.
      */
     @Test
-    public void onResultsDoesNotUpdateIfNotOnboarded() {
+    public void onResultsUpdatesIfNotOnboarded() {
         initializeWakeupController(true /* enabled */);
         when(mWakeupOnboarding.isOnboarded()).thenReturn(false);
+        when(mWakeupLock.isUnlocked()).thenReturn(false);
         mWakeupController.start();
 
         ArgumentCaptor<WifiScanner.ScanListener> scanListenerArgumentCaptor =
@@ -401,9 +439,8 @@ public class WakeupControllerTest {
         // incoming scan results
         scanListener.onResults(mTestScanDatas);
 
-        verify(mWakeupLock, never()).isEmpty();
-        verify(mWakeupLock, never()).update(any());
-
+        verify(mWakeupLock).update(any());
+        verify(mWakeupLock).isUnlocked();
         verifyDoesNotEnableWifi();
     }
 
@@ -412,8 +449,8 @@ public class WakeupControllerTest {
      */
     @Test
     public void onResultsEnablesWifi() {
-        // empty wakeup lock
-        when(mWakeupLock.isEmpty()).thenReturn(true);
+        // unlock wakeup lock
+        when(mWakeupLock.isUnlocked()).thenReturn(true);
         // find viable network
         when(mWakeupEvaluator.findViableNetwork(any(), any())).thenReturn(mTestScanResult);
 
@@ -451,10 +488,28 @@ public class WakeupControllerTest {
         // incoming scan results
         scanListener.onResults(mTestScanDatas);
 
-        verify(mWakeupLock, never()).initialize(any());
+        verify(mWakeupLock, never()).setLock(any());
         verify(mWakeupLock, never()).update(any());
-        verify(mWakeupLock, never()).isEmpty();
+        verify(mWakeupLock, never()).isUnlocked();
         verify(mWakeupOnboarding, never()).maybeShowNotification();
         verify(mWakeupEvaluator, never()).findViableNetwork(any(), any());
     }
+
+    @Test
+    public void userIsNotOnboardedByInitialization() {
+        initializeWakeupController(true /* enabled */);
+        verify(mWakeupOnboarding, never()).setOnboarded();
+    }
+
+    @Test
+    public void userIsOnboardedBySettingChange() {
+        initializeWakeupController(true /* enabled */);
+        ArgumentCaptor<ContentObserver> argumentCaptor =
+                ArgumentCaptor.forClass(ContentObserver.class);
+        verify(mFrameworkFacade).registerContentObserver(any(), any(), eq(true),
+                argumentCaptor.capture());
+        ContentObserver contentObserver = argumentCaptor.getValue();
+        contentObserver.onChange(false /* selfChange */);
+        verify(mWakeupOnboarding).setOnboarded();
+    }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WakeupLockTest.java b/tests/wifitests/src/com/android/server/wifi/WakeupLockTest.java
index 7144ecf3f..f3126c7a9 100644
--- a/tests/wifitests/src/com/android/server/wifi/WakeupLockTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WakeupLockTest.java
@@ -40,6 +40,8 @@ public class WakeupLockTest {
     private static final String SSID_2 = "ssid2";
 
     @Mock private WifiConfigManager mWifiConfigManager;
+    @Mock private WifiWakeMetrics mWifiWakeMetrics;
+    @Mock private Clock mClock;
 
     private ScanResultMatchInfo mNetwork1;
     private ScanResultMatchInfo mNetwork2;
@@ -52,6 +54,8 @@ public class WakeupLockTest {
     public void setUp() {
         MockitoAnnotations.initMocks(this);
 
+        when(mClock.getElapsedSinceBootMillis()).thenReturn(0L);
+
         mNetwork1 = new ScanResultMatchInfo();
         mNetwork1.networkSsid = SSID_1;
         mNetwork1.networkType = ScanResultMatchInfo.NETWORK_TYPE_OPEN;
@@ -60,7 +64,7 @@ public class WakeupLockTest {
         mNetwork2.networkSsid = SSID_2;
         mNetwork2.networkType = ScanResultMatchInfo.NETWORK_TYPE_EAP;
 
-        mWakeupLock = new WakeupLock(mWifiConfigManager);
+        mWakeupLock = new WakeupLock(mWifiConfigManager, mWifiWakeMetrics, mClock);
     }
 
     /**
@@ -71,7 +75,7 @@ public class WakeupLockTest {
      */
     private void updateEnoughTimesToEvictWithAsserts(Collection<ScanResultMatchInfo> networks) {
         for (int i = 0; i < WakeupLock.CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT; i++) {
-            assertFalse("Lock empty after " + i + " scans", mWakeupLock.isEmpty());
+            assertFalse("Lock empty after " + i + " scans", mWakeupLock.isUnlocked());
             mWakeupLock.update(networks);
         }
     }
@@ -89,33 +93,82 @@ public class WakeupLockTest {
     }
 
     /**
+     * Verify that calling update {@link WakeupLock#CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT}
+     * times sets the lock to be initialized.
+     */
+    @Test
+    public void verifyInitializingLockByScans() {
+        List<ScanResultMatchInfo> networks = Collections.singletonList(mNetwork1);
+        mWakeupLock.setLock(networks);
+        assertFalse(mWakeupLock.isInitialized());
+
+        mWakeupLock.update(networks);
+        assertFalse(mWakeupLock.isInitialized());
+        mWakeupLock.update(networks);
+        assertFalse(mWakeupLock.isInitialized());
+        mWakeupLock.update(networks);
+        assertTrue(mWakeupLock.isInitialized());
+    }
+
+    /**
+     * Verify that calling update after {@link WakeupLock#MAX_LOCK_TIME_MILLIS} milliseconds sets
+     * the lock to be initialized and does not add the scans to the lock.
+     */
+    @Test
+    public void verifyInitializingLockByTimeout() {
+        when(mClock.getElapsedSinceBootMillis())
+                .thenReturn(0L, WakeupLock.MAX_LOCK_TIME_MILLIS + 1);
+        mWakeupLock.setLock(Collections.emptyList());
+        assertFalse(mWakeupLock.isInitialized());
+
+        List<ScanResultMatchInfo> networks = Collections.singletonList(mNetwork1);
+        mWakeupLock.update(networks);
+        assertTrue(mWakeupLock.isInitialized());
+        assertTrue(mWakeupLock.isUnlocked());
+    }
+
+    /**
+     * Verify that addToLock saves to the store if it changes the contents of the lock.
+     */
+    @Test
+    public void addToLockSavesToStore() {
+        mWakeupLock.setLock(Collections.emptyList());
+
+        List<ScanResultMatchInfo> networks = Collections.singletonList(mNetwork1);
+        mWakeupLock.update(networks);
+
+        // want 2 invocations, once for setLock(), once for addToLock
+        verify(mWifiConfigManager, times(2)).saveToStore(false);
+    }
+
+    /**
      * Verify that the WakeupLock is not empty immediately after being initialized with networks.
      */
     @Test
-    public void verifyNotEmptyWhenInitializedWithNetworkList() {
-        mWakeupLock.initialize(Arrays.asList(mNetwork1, mNetwork2));
-        assertFalse(mWakeupLock.isEmpty());
+    public void verifyNotEmptyWhenSetWithNetworkList() {
+        setLockAndInitializeByTimeout(Arrays.asList(mNetwork1, mNetwork2));
+        assertFalse(mWakeupLock.isUnlocked());
     }
 
     /**
-     * Verify that the WakeupLock is empty when initialized with an empty list.
+     * Verify that the WakeupLock is unlocked when initialized with an empty list.
      */
     @Test
     public void isEmptyWhenInitializedWithEmptyList() {
-        mWakeupLock.initialize(Collections.emptyList());
-        assertTrue(mWakeupLock.isEmpty());
+        setLockAndInitializeByTimeout(Collections.emptyList());
+        assertTrue(mWakeupLock.isUnlocked());
     }
 
     /**
-     * Verify that initializing the WakeupLock clears out previous entries.
+     * Verify that setting the lock clears out previous entries.
      */
     @Test
-    public void initializingLockClearsPreviousNetworks() {
-        mWakeupLock.initialize(Collections.singletonList(mNetwork1));
-        assertFalse(mWakeupLock.isEmpty());
+    public void setLockClearsPreviousNetworks() {
+        setLockAndInitializeByTimeout(Collections.singletonList(mNetwork1));
+        assertFalse(mWakeupLock.isUnlocked());
 
-        mWakeupLock.initialize(Collections.emptyList());
-        assertTrue(mWakeupLock.isEmpty());
+        setLockAndInitializeByTimeout(Collections.emptyList());
+        assertTrue(mWakeupLock.isUnlocked());
     }
 
     /**
@@ -124,11 +177,11 @@ public class WakeupLockTest {
      */
     @Test
     public void updateShouldRemoveNetworksAfterConsecutiveMissedScans() {
-        mWakeupLock.initialize(Collections.singletonList(mNetwork1));
+        setLockAndInitializeByTimeout(Collections.singletonList(mNetwork1));
 
         updateEnoughTimesToEvictWithAsserts(Collections.singletonList(mNetwork2));
 
-        assertTrue(mWakeupLock.isEmpty());
+        assertTrue(mWakeupLock.isUnlocked());
     }
 
     /**
@@ -139,7 +192,7 @@ public class WakeupLockTest {
         List<ScanResultMatchInfo> lockedNetworks = Collections.singletonList(mNetwork1);
         List<ScanResultMatchInfo> updateNetworks = Collections.singletonList(mNetwork2);
 
-        mWakeupLock.initialize(lockedNetworks);
+        setLockAndInitializeByTimeout(lockedNetworks);
 
         // one update without network
         mWakeupLock.update(updateNetworks);
@@ -148,7 +201,7 @@ public class WakeupLockTest {
 
         updateEnoughTimesToEvictWithAsserts(updateNetworks);
 
-        assertTrue(mWakeupLock.isEmpty());
+        assertTrue(mWakeupLock.isUnlocked());
     }
 
     /**
@@ -157,13 +210,13 @@ public class WakeupLockTest {
     @Test
     public void updateWithLockedNetworkAfterItIsRemovedDoesNotReset() {
         List<ScanResultMatchInfo> lockedNetworks = Collections.singletonList(mNetwork1);
-        mWakeupLock.initialize(lockedNetworks);
+        setLockAndInitializeByTimeout(lockedNetworks);
 
         updateEnoughTimesToEvictWithAsserts(Collections.emptyList());
 
-        assertTrue(mWakeupLock.isEmpty());
+        assertTrue(mWakeupLock.isUnlocked());
         mWakeupLock.update(lockedNetworks);
-        assertTrue(mWakeupLock.isEmpty());
+        assertTrue(mWakeupLock.isUnlocked());
     }
 
     /**
@@ -173,13 +226,13 @@ public class WakeupLockTest {
     @Test
     public void networksCanBeRemovedIncrementallyFromLock() {
         List<ScanResultMatchInfo> lockedNetworks = Arrays.asList(mNetwork1, mNetwork2);
-        mWakeupLock.initialize(lockedNetworks);
+        setLockAndInitializeByTimeout(lockedNetworks);
 
         updateEnoughTimesToEvictWithAsserts(Collections.singletonList(mNetwork1));
-        assertFalse(mWakeupLock.isEmpty());
+        assertFalse(mWakeupLock.isUnlocked());
 
         updateEnoughTimesToEvictWithAsserts(Collections.singletonList(mNetwork2));
-        assertTrue(mWakeupLock.isEmpty());
+        assertTrue(mWakeupLock.isUnlocked());
     }
 
     /**
@@ -187,7 +240,7 @@ public class WakeupLockTest {
      */
     @Test
     public void initializeShouldSaveSsidsToStore() {
-        mWakeupLock.initialize(Collections.singletonList(mNetwork1));
+        setLockAndInitializeByTimeout(Collections.singletonList(mNetwork1));
         verify(mWifiConfigManager).saveToStore(eq(false));
     }
 
@@ -196,7 +249,7 @@ public class WakeupLockTest {
      */
     @Test
     public void updateShouldOnlySaveIfLockChanges() {
-        mWakeupLock.initialize(Collections.singletonList(mNetwork1));
+        setLockAndInitializeByTimeout(Collections.singletonList(mNetwork1));
         updateEnoughTimesToEvictWithoutAsserts(Collections.emptyList());
 
         // need exactly 2 invocations: 1 for initialize, 1 for successful update
@@ -208,7 +261,33 @@ public class WakeupLockTest {
      */
     @Test
     public void updateShouldNotSaveIfLockDoesNotChange() {
-        mWakeupLock.update(Collections.singletonList(mNetwork1));
-        verify(mWifiConfigManager, never()).saveToStore(anyBoolean());
+        List<ScanResultMatchInfo> networks = Collections.singletonList(mNetwork1);
+        setLockAndInitializeByTimeout(networks);
+        verify(mWifiConfigManager, times(1)).saveToStore(anyBoolean());
+        mWakeupLock.update(networks);
+    }
+
+    /**
+     * Verify that on unlock, records the unlock event with WifiWakeMetrics with the correct number
+     * of scans.
+     */
+    @Test
+    public void unlockingShouldRecordEventInMetrics() {
+        when(mClock.getElapsedSinceBootMillis())
+                .thenReturn(0L, WakeupLock.MAX_LOCK_TIME_MILLIS + 1);
+        List<ScanResultMatchInfo> networks = Collections.singletonList(mNetwork1);
+        mWakeupLock.setLock(networks);
+        for (int i = 0; i < WakeupLock.CONSECUTIVE_MISSED_SCANS_REQUIRED_TO_EVICT; i++) {
+            mWakeupLock.update(Collections.emptyList());
+        }
+        verify(mWifiWakeMetrics).recordUnlockEvent(3 /* numScans */);
+    }
+
+    private void setLockAndInitializeByTimeout(Collection<ScanResultMatchInfo> networks) {
+        when(mClock.getElapsedSinceBootMillis())
+                .thenReturn(0L, WakeupLock.MAX_LOCK_TIME_MILLIS + 1);
+        mWakeupLock.setLock(networks);
+        mWakeupLock.update(networks);
+        assertTrue(mWakeupLock.isInitialized());
     }
 }
diff --git a/tests/wifitests/src/com/android/server/wifi/WakeupOnboardingTest.java b/tests/wifitests/src/com/android/server/wifi/WakeupOnboardingTest.java
index 1e98b1455..0fe0e413f 100644
--- a/tests/wifitests/src/com/android/server/wifi/WakeupOnboardingTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WakeupOnboardingTest.java
@@ -25,6 +25,7 @@ import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -37,8 +38,6 @@ import android.os.Handler;
 import android.os.test.TestLooper;
 import android.provider.Settings;
 
-import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
-
 import org.junit.Before;
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
@@ -61,7 +60,11 @@ public class WakeupOnboardingTest {
 
     // convenience method for resetting onboarded status
     private void setOnboardedStatus(boolean isOnboarded) {
-        mWakeupOnboarding.getDataSource().setData(isOnboarded);
+        mWakeupOnboarding.getIsOnboadedDataSource().setData(isOnboarded);
+    }
+
+    private void setNotificationsShown(int numNotifications) {
+        mWakeupOnboarding.getNotificationsDataSource().setData(numNotifications);
     }
 
     @Before
@@ -84,7 +87,7 @@ public class WakeupOnboardingTest {
         setOnboardedStatus(false);
         mWakeupOnboarding.maybeShowNotification();
 
-        verify(mNotificationManager).notify(eq(SystemMessage.NOTE_WIFI_WAKE_ONBOARD), any());
+        verify(mNotificationManager).notify(eq(WakeupNotificationFactory.ONBOARD_ID), any());
     }
 
     /**
@@ -96,7 +99,7 @@ public class WakeupOnboardingTest {
         mWakeupOnboarding.maybeShowNotification();
 
         verify(mNotificationManager, never())
-                .notify(eq(SystemMessage.NOTE_WIFI_WAKE_ONBOARD), any());
+                .notify(eq(WakeupNotificationFactory.ONBOARD_ID), any());
     }
 
     /**
@@ -110,7 +113,7 @@ public class WakeupOnboardingTest {
 
         InOrder inOrder = Mockito.inOrder(mNotificationManager);
         inOrder.verify(mNotificationManager)
-                .notify(eq(SystemMessage.NOTE_WIFI_WAKE_ONBOARD), any());
+                .notify(eq(WakeupNotificationFactory.ONBOARD_ID), any());
         inOrder.verifyNoMoreInteractions();
     }
 
@@ -130,8 +133,7 @@ public class WakeupOnboardingTest {
 
         broadcastReceiver.onReceive(mContext, new Intent(ACTION_DISMISS_NOTIFICATION));
 
-        verify(mNotificationManager).cancel(SystemMessage.NOTE_WIFI_WAKE_ONBOARD);
-        verify(mWifiConfigManager).saveToStore(false);
+        verify(mNotificationManager).cancel(WakeupNotificationFactory.ONBOARD_ID);
         assertTrue(mWakeupOnboarding.isOnboarded());
     }
 
@@ -155,8 +157,7 @@ public class WakeupOnboardingTest {
         verify(mFrameworkFacade).setIntegerSetting(mContext,
                 Settings.Global.WIFI_WAKEUP_ENABLED, 0);
 
-        verify(mNotificationManager).cancel(SystemMessage.NOTE_WIFI_WAKE_ONBOARD);
-        verify(mWifiConfigManager).saveToStore(false);
+        verify(mNotificationManager).cancel(WakeupNotificationFactory.ONBOARD_ID);
         assertTrue(mWakeupOnboarding.isOnboarded());
     }
 
@@ -179,8 +180,7 @@ public class WakeupOnboardingTest {
 
         verify(mContext).startActivity(any());
 
-        verify(mNotificationManager).cancel(SystemMessage.NOTE_WIFI_WAKE_ONBOARD);
-        verify(mWifiConfigManager).saveToStore(false);
+        verify(mNotificationManager).cancel(WakeupNotificationFactory.ONBOARD_ID);
         assertTrue(mWakeupOnboarding.isOnboarded());
     }
 
@@ -195,7 +195,79 @@ public class WakeupOnboardingTest {
         mWakeupOnboarding.maybeShowNotification();
         mWakeupOnboarding.onStop();
 
-        verify(mNotificationManager).cancel(SystemMessage.NOTE_WIFI_WAKE_ONBOARD);
+        verify(mNotificationManager).cancel(WakeupNotificationFactory.ONBOARD_ID);
+        assertFalse(mWakeupOnboarding.isOnboarded());
+    }
+
+    /**
+     * Verify that incrementing the notification count saves to store.
+     */
+    @Test
+    public void setOnboardedSavesToStore() {
+        setOnboardedStatus(false);
+        mWakeupOnboarding.setOnboarded();
+        verify(mWifiConfigManager).saveToStore(false /* forceWrite */);
+        assertTrue(mWakeupOnboarding.isOnboarded());
+    }
+
+    /**
+     * Verify that incrementing the notification count saves to store.
+     */
+    @Test
+    public void incrementingNotificationCountSavesToStore() {
+        setOnboardedStatus(false);
+        setNotificationsShown(0);
+        mWakeupOnboarding.maybeShowNotification();
+        verify(mWifiConfigManager).saveToStore(false /* forceWrite */);
+    }
+
+    /**
+     * Verify that the notification does not show multiple times within 24 hours.
+     */
+    @Test
+    public void doesNotShowMultipleNotificationsWithin24Hours() {
+        setOnboardedStatus(false);
+        setNotificationsShown(0);
+
+        mWakeupOnboarding.maybeShowNotification(0 /* timestamp */);
+        mWakeupOnboarding.onStop();
+        mWakeupOnboarding.maybeShowNotification(0 /* timestamp */);
+
+        InOrder inOrder = Mockito.inOrder(mNotificationManager);
+        inOrder.verify(mNotificationManager)
+                .notify(eq(WakeupNotificationFactory.ONBOARD_ID), any());
+        inOrder.verify(mNotificationManager).cancel(WakeupNotificationFactory.ONBOARD_ID);
+        inOrder.verifyNoMoreInteractions();
+    }
+
+    /**
+     * Verify that notification reappears after 24 hours if not onboarded.
+     */
+    @Test
+    public void showsNotificationsOutsideOf24Hours() {
+        setOnboardedStatus(false);
+        setNotificationsShown(0);
+
+        mWakeupOnboarding.maybeShowNotification(0 /* timestamp */);
         assertFalse(mWakeupOnboarding.isOnboarded());
+
+        mWakeupOnboarding.onStop();
+        mWakeupOnboarding.maybeShowNotification(WakeupOnboarding.REQUIRED_NOTIFICATION_DELAY + 1);
+
+        verify(mNotificationManager, times(2))
+                .notify(eq(WakeupNotificationFactory.ONBOARD_ID), any());
+    }
+
+    /**
+     * Verify that the user is onboarded after
+     * {@link WakeupOnboarding#NOTIFICATIONS_UNTIL_ONBOARDED} notifications are shown.
+     */
+    @Test
+    public void onboardsUserAfterThreeNotifications() {
+        setOnboardedStatus(false);
+        setNotificationsShown(WakeupOnboarding.NOTIFICATIONS_UNTIL_ONBOARDED - 1);
+
+        mWakeupOnboarding.maybeShowNotification(0 /* timestamp */);
+        assertTrue(mWakeupOnboarding.isOnboarded());
     }
 }