Change onboarding flow.

This CL incorporates a few changes to the Wifi Wake onboarding flow.

Wifi Wake now operates as normal even if the user has not interacted
with the onboarding notification.

The onboarding notification will be shown when WiFi is disabled if the
user is not onboarded (3 times max, at most once every 24 hours).

The user is considered onboarded if they have interacted with the
notification, seen 3 notifications, or manually changed enabled/disabled
the feature from Wifi Settings page.

Bug: 72399908
Test: runtest, manual

Change-Id: I5fd02510b39ae89223ad28352f2045af33c13a5e

7 files changed, 227 insertions, 50 deletions

diff --git a/service/java/com/android/server/wifi/WakeupConfigStoreData.java b/service/java/com/android/server/wifi/WakeupConfigStoreData.java
index b936a4c6c..d98567766 100644
--- a/service/java/com/android/server/wifi/WakeupConfigStoreData.java
+++ b/service/java/com/android/server/wifi/WakeupConfigStoreData.java
@@ -39,12 +39,14 @@ public class WakeupConfigStoreData implements StoreData {
     private static final String XML_TAG_FEATURE_STATE_SECTION = "FeatureState";
     private static final String XML_TAG_IS_ACTIVE = "IsActive";
     private static final String XML_TAG_IS_ONBOARDED = "IsOnboarded";
+    private static final String XML_TAG_NOTIFICATIONS_SHOWN = "NotificationsShown";
     private static final String XML_TAG_NETWORK_SECTION = "Network";
     private static final String XML_TAG_SSID = "SSID";
     private static final String XML_TAG_SECURITY = "Security";
 
     private final DataSource<Boolean> mIsActiveDataSource;
     private final DataSource<Boolean> mIsOnboardedDataSource;
+    private final DataSource<Integer> mNotificationsDataSource;
     private final DataSource<Set<ScanResultMatchInfo>> mNetworkDataSource;
     private boolean mHasBeenRead = false;
 
@@ -76,9 +78,11 @@ public class WakeupConfigStoreData implements StoreData {
     public WakeupConfigStoreData(
             DataSource<Boolean> isActiveDataSource,
             DataSource<Boolean> isOnboardedDataSource,
+            DataSource<Integer> notificationsDataSource,
             DataSource<Set<ScanResultMatchInfo>> networkDataSource) {
         mIsActiveDataSource = isActiveDataSource;
         mIsOnboardedDataSource = isOnboardedDataSource;
+        mNotificationsDataSource = notificationsDataSource;
         mNetworkDataSource = networkDataSource;
     }
 
@@ -116,6 +120,8 @@ public class WakeupConfigStoreData implements StoreData {
 
         XmlUtil.writeNextValue(out, XML_TAG_IS_ACTIVE, mIsActiveDataSource.getData());
         XmlUtil.writeNextValue(out, XML_TAG_IS_ONBOARDED, mIsOnboardedDataSource.getData());
+        XmlUtil.writeNextValue(out, XML_TAG_NOTIFICATIONS_SHOWN,
+                mNotificationsDataSource.getData());
 
         XmlUtil.writeNextSectionEnd(out, XML_TAG_FEATURE_STATE_SECTION);
     }
@@ -185,6 +191,7 @@ public class WakeupConfigStoreData implements StoreData {
             throws IOException, XmlPullParserException {
         boolean isActive = false;
         boolean isOnboarded = false;
+        int notificationsShown = 0;
 
         while (!XmlUtil.isNextSectionEnd(in, outerTagDepth)) {
             String[] valueName = new String[1];
@@ -199,6 +206,9 @@ public class WakeupConfigStoreData implements StoreData {
                 case XML_TAG_IS_ONBOARDED:
                     isOnboarded = (Boolean) value;
                     break;
+                case XML_TAG_NOTIFICATIONS_SHOWN:
+                    notificationsShown = (Integer) value;
+                    break;
                 default:
                     throw new XmlPullParserException("Unknown value found: " + valueName[0]);
             }
@@ -206,6 +216,7 @@ public class WakeupConfigStoreData implements StoreData {
 
         mIsActiveDataSource.setData(isActive);
         mIsOnboardedDataSource.setData(isOnboarded);
+        mNotificationsDataSource.setData(notificationsShown);
     }
 
     /**
@@ -248,6 +259,7 @@ public class WakeupConfigStoreData implements StoreData {
             mNetworkDataSource.setData(Collections.emptySet());
             mIsActiveDataSource.setData(false);
             mIsOnboardedDataSource.setData(false);
+            mNotificationsDataSource.setData(0);
         }
     }
 
diff --git a/service/java/com/android/server/wifi/WakeupController.java b/service/java/com/android/server/wifi/WakeupController.java
index 97919ca72..e8adbb27c 100644
--- a/service/java/com/android/server/wifi/WakeupController.java
+++ b/service/java/com/android/server/wifi/WakeupController.java
@@ -127,24 +127,30 @@ public class WakeupController {
         mContentObserver = new ContentObserver(mHandler) {
             @Override
             public void onChange(boolean selfChange) {
-                mWifiWakeupEnabled = mFrameworkFacade.getIntegerSetting(
-                                mContext, Settings.Global.WIFI_WAKEUP_ENABLED, 0) == 1;
-                Log.d(TAG, "WifiWake " + (mWifiWakeupEnabled ? "enabled" : "disabled"));
+                readWifiWakeupEnabledFromSettings();
+                mWakeupOnboarding.setOnboarded();
             }
         };
         mFrameworkFacade.registerContentObserver(mContext, Settings.Global.getUriFor(
                 Settings.Global.WIFI_WAKEUP_ENABLED), true, mContentObserver);
-        mContentObserver.onChange(false /* selfChange */);
+        readWifiWakeupEnabledFromSettings();
 
         // registering the store data here has the effect of reading the persisted value of the
         // data sources after system boot finishes
         mWakeupConfigStoreData = new WakeupConfigStoreData(
                 new IsActiveDataSource(),
-                mWakeupOnboarding.getDataSource(),
+                mWakeupOnboarding.getIsOnboadedDataSource(),
+                mWakeupOnboarding.getNotificationsDataSource(),
                 mWakeupLock.getDataSource());
         wifiConfigStore.registerStoreData(mWakeupConfigStoreData);
     }
 
+    private void readWifiWakeupEnabledFromSettings() {
+        mWifiWakeupEnabled = mFrameworkFacade.getIntegerSetting(
+                mContext, Settings.Global.WIFI_WAKEUP_ENABLED, 0) == 1;
+        Log.d(TAG, "WifiWake " + (mWifiWakeupEnabled ? "enabled" : "disabled"));
+    }
+
     private void setActive(boolean isActive) {
         if (mIsActive != isActive) {
             Log.d(TAG, "Setting active to " + isActive);
@@ -185,6 +191,7 @@ public class WakeupController {
             // TODO(b/77291248): request low latency scan here
         }
     }
+
     /**
      * Stops listening for scans.
      *
@@ -280,24 +287,17 @@ public class WakeupController {
             return;
         }
 
-        Set<ScanResult> filteredScanResults = filterScanResults(scanResults);
-
-        // need to show notification here in case user enables Wifi Wake when Wifi is off
-        // TODO(b/72399908) make onboarding not blocking
-        mWakeupOnboarding.maybeShowNotification();
-        if (!mWakeupOnboarding.isOnboarded()) {
-            if (mVerboseLoggingEnabled) {
-                Log.d(TAG, "handleScanResults: Scan not handled because user is not onboarded.");
-            }
-            return;
-        }
-
-        // only count scan as handled if isEnabled and user onboarded
+        // only count scan as handled if isEnabled
         mNumScansHandled++;
         if (mVerboseLoggingEnabled) {
             Log.d(TAG, "Incoming scan #" + mNumScansHandled);
         }
 
+        // need to show notification here in case user turns phone on while wifi is off
+        mWakeupOnboarding.maybeShowNotification();
+
+        Set<ScanResult> filteredScanResults = filterScanResults(scanResults);
+
         mWakeupLock.update(toMatchInfos(filteredScanResults));
         if (!mWakeupLock.isUnlocked()) {
             return;
diff --git a/service/java/com/android/server/wifi/WakeupNotificationFactory.java b/service/java/com/android/server/wifi/WakeupNotificationFactory.java
index 42ae46707..23f31a7db 100644
--- a/service/java/com/android/server/wifi/WakeupNotificationFactory.java
+++ b/service/java/com/android/server/wifi/WakeupNotificationFactory.java
@@ -22,6 +22,7 @@ import android.content.Context;
 import android.content.Intent;
 
 import com.android.internal.R;
+import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
 import com.android.internal.notification.SystemNotificationChannels;
 
 
@@ -37,6 +38,9 @@ public class WakeupNotificationFactory {
     public static final String ACTION_TURN_OFF_WIFI_WAKE =
             "com.android.server.wifi.wakeup.TURN_OFF_WIFI_WAKE";
 
+    /** Notification channel ID for onboarding messages. */
+    public static final int ONBOARD_ID = SystemMessage.NOTE_WIFI_WAKE_ONBOARD;
+
     private final Context mContext;
     private final FrameworkFacade mFrameworkFacade;
 
diff --git a/service/java/com/android/server/wifi/WakeupOnboarding.java b/service/java/com/android/server/wifi/WakeupOnboarding.java
index d4caa0fdb..b6bcbc3c0 100644
--- a/service/java/com/android/server/wifi/WakeupOnboarding.java
+++ b/service/java/com/android/server/wifi/WakeupOnboarding.java
@@ -27,22 +27,31 @@ import android.content.Intent;
 import android.content.IntentFilter;
 import android.os.Handler;
 import android.os.Looper;
+import android.os.SystemClock;
 import android.provider.Settings;
+import android.text.format.DateUtils;
 import android.util.Log;
 
-import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
+import com.android.internal.annotations.VisibleForTesting;
 
 /**
  * Manages the WiFi Wake onboarding notification.
  *
  * <p>If a user disables wifi with Wifi Wake enabled, this notification is shown to explain that
- * wifi may turn back on automatically. Wifi will not automatically turn back on until after the
- * user interacts with the onboarding notification in some way (e.g. dismiss, tap).
+ * wifi may turn back on automatically. It will be displayed up to 3 times, or until the
+ * user either interacts with the onboarding notification in some way (e.g. dismiss, tap) or
+ * manually enables/disables the feature in WifiSettings.
  */
 public class WakeupOnboarding {
 
     private static final String TAG = "WakeupOnboarding";
 
+    @VisibleForTesting
+    static final int NOTIFICATIONS_UNTIL_ONBOARDED = 3;
+    @VisibleForTesting
+    static final long REQUIRED_NOTIFICATION_DELAY = DateUtils.DAY_IN_MILLIS;
+    private static final long NOT_SHOWN_TIMESTAMP = -1;
+
     private final Context mContext;
     private final WakeupNotificationFactory mWakeupNotificationFactory;
     private NotificationManager mNotificationManager;
@@ -52,6 +61,8 @@ public class WakeupOnboarding {
     private final FrameworkFacade mFrameworkFacade;
 
     private boolean mIsOnboarded;
+    private int mTotalNotificationsShown;
+    private long mLastShownTimestamp = NOT_SHOWN_TIMESTAMP;
     private boolean mIsNotificationShowing;
 
     private final BroadcastReceiver mBroadcastReceiver = new BroadcastReceiver() {
@@ -104,17 +115,46 @@ public class WakeupOnboarding {
 
     /** Shows the onboarding notification if applicable. */
     public void maybeShowNotification() {
-        if (isOnboarded() || mIsNotificationShowing) {
+        maybeShowNotification(SystemClock.elapsedRealtime());
+    }
+
+    @VisibleForTesting
+    void maybeShowNotification(long timestamp) {
+        if (!shouldShowNotification(timestamp)) {
             return;
         }
-
         Log.d(TAG, "Showing onboarding notification.");
 
+        incrementTotalNotificationsShown();
+        mIsNotificationShowing = true;
+        mLastShownTimestamp = timestamp;
+
         mContext.registerReceiver(mBroadcastReceiver, mIntentFilter,
                 null /* broadcastPermission */, mHandler);
-        getNotificationManager().notify(SystemMessage.NOTE_WIFI_WAKE_ONBOARD,
+        getNotificationManager().notify(WakeupNotificationFactory.ONBOARD_ID,
                 mWakeupNotificationFactory.createOnboardingNotification());
-        mIsNotificationShowing = true;
+    }
+
+    /**
+     * Increment the total number of shown notifications and onboard the user if reached the
+     * required amount.
+     */
+    private void incrementTotalNotificationsShown() {
+        mTotalNotificationsShown++;
+        if (mTotalNotificationsShown >= NOTIFICATIONS_UNTIL_ONBOARDED) {
+            setOnboarded();
+        } else {
+            mWifiConfigManager.saveToStore(false /* forceWrite */);
+        }
+    }
+
+    private boolean shouldShowNotification(long timestamp) {
+        if (isOnboarded() || mIsNotificationShowing) {
+            return false;
+        }
+
+        return mLastShownTimestamp == NOT_SHOWN_TIMESTAMP
+                || (timestamp - mLastShownTimestamp) > REQUIRED_NOTIFICATION_DELAY;
     }
 
     /** Handles onboarding cleanup on stop. */
@@ -132,11 +172,15 @@ public class WakeupOnboarding {
         }
 
         mContext.unregisterReceiver(mBroadcastReceiver);
-        getNotificationManager().cancel(SystemMessage.NOTE_WIFI_WAKE_ONBOARD);
+        getNotificationManager().cancel(WakeupNotificationFactory.ONBOARD_ID);
         mIsNotificationShowing = false;
     }
 
-    private void setOnboarded() {
+    /** Sets the user as onboarded and persists to store. */
+    public void setOnboarded() {
+        if (mIsOnboarded) {
+            return;
+        }
         Log.d(TAG, "Setting user as onboarded.");
         mIsOnboarded = true;
         mWifiConfigManager.saveToStore(false /* forceWrite */);
@@ -150,12 +194,17 @@ public class WakeupOnboarding {
         return mNotificationManager;
     }
 
-    /** Returns the {@link WakeupConfigStoreData.DataSource} for the {@link WifiConfigStore}. */
-    public WakeupConfigStoreData.DataSource<Boolean> getDataSource() {
-        return new OnboardingDataSource();
+    /** Returns the {@link WakeupConfigStoreData.DataSource} for the onboarded status. */
+    public WakeupConfigStoreData.DataSource<Boolean> getIsOnboadedDataSource() {
+        return new IsOnboardedDataSource();
     }
 
-    private class OnboardingDataSource implements WakeupConfigStoreData.DataSource<Boolean> {
+    /** Returns the {@link WakeupConfigStoreData.DataSource} for the notification status. */
+    public WakeupConfigStoreData.DataSource<Integer> getNotificationsDataSource() {
+        return new NotificationsDataSource();
+    }
+
+    private class IsOnboardedDataSource implements WakeupConfigStoreData.DataSource<Boolean> {
 
         @Override
         public Boolean getData() {
@@ -167,4 +216,17 @@ public class WakeupOnboarding {
             mIsOnboarded = data;
         }
     }
+
+    private class NotificationsDataSource implements WakeupConfigStoreData.DataSource<Integer> {
+
+        @Override
+        public Integer getData() {
+            return mTotalNotificationsShown;
+        }
+
+        @Override
+        public void setData(Integer data) {
+            mTotalNotificationsShown = data;
+        }
+    }
 }
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 3c7517bf3..1aedae286 100644
--- a/tests/wifitests/src/com/android/server/wifi/WakeupControllerTest.java
+++ b/tests/wifitests/src/com/android/server/wifi/WakeupControllerTest.java
@@ -27,6 +27,7 @@ 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;
@@ -420,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 =
@@ -437,9 +439,8 @@ public class WakeupControllerTest {
         // incoming scan results
         scanListener.onResults(mTestScanDatas);
 
-        verify(mWakeupLock, never()).isUnlocked();
-        verify(mWakeupLock, never()).update(any());
-
+        verify(mWakeupLock).update(any());
+        verify(mWakeupLock).isUnlocked();
         verifyDoesNotEnableWifi();
     }
 
@@ -493,4 +494,22 @@ public class WakeupControllerTest {
         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/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());
     }
 }