diff --git a/DeviceLockController/src/com/android/devicelockcontroller/activities/DeviceLockNotificationManager.java b/DeviceLockController/src/com/android/devicelockcontroller/activities/DeviceLockNotificationManager.java
index b5d5d33..a3957fd 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/activities/DeviceLockNotificationManager.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/activities/DeviceLockNotificationManager.java
@@ -39,6 +39,7 @@
 import com.android.devicelockcontroller.storage.UserParameters;
 import com.android.devicelockcontroller.util.LogUtil;
 import com.android.devicelockcontroller.util.StringUtil;
+import com.android.internal.annotations.VisibleForTesting;
 
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -60,8 +61,10 @@
     private static final String TAG = "DeviceLockNotificationManager";
 
     private static final String PROVISION_NOTIFICATION_CHANNEL_ID_BASE = "devicelock-provision";
-    private static final String DEVICE_RESET_NOTIFICATION_TAG = "devicelock-device-reset";
-    private static final int DEVICE_RESET_NOTIFICATION_ID = 0;
+    @VisibleForTesting
+    public static final String DEVICE_RESET_NOTIFICATION_TAG = "devicelock-device-reset";
+    @VisibleForTesting
+    public static final int DEVICE_RESET_NOTIFICATION_ID = 0;
     private static final int DEFER_PROVISIONING_NOTIFICATION_ID = 1;
 
     private static final ListeningExecutorService sListeningExecutorService =
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorker.java b/DeviceLockController/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorker.java
index 35b0bd1..b286034 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorker.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorker.java
@@ -16,6 +16,12 @@
 
 package com.android.devicelockcontroller.provision.worker;
 
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_DISMISSIBLE_UI;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_FACTORY_RESET;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_PERSISTENT_UI;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_RETRY;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_SUCCESS;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_UNSPECIFIED;
 import static com.android.devicelockcontroller.common.DeviceLockConstants.ProvisionFailureReason.DEADLINE_PASSED;
 
 import android.content.Context;
@@ -32,6 +38,7 @@
 import androidx.work.WorkManager;
 import androidx.work.WorkerParameters;
 
+import com.android.devicelockcontroller.activities.DeviceLockNotificationManager;
 import com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState;
 import com.android.devicelockcontroller.common.DeviceLockConstants.ProvisionFailureReason;
 import com.android.devicelockcontroller.provision.grpc.DeviceCheckInClient;
@@ -57,6 +64,8 @@
     public static final String KEY_IS_PROVISION_SUCCESSFUL = "is-provision-successful";
     public static final String KEY_PROVISION_FAILURE_REASON = "provision-failure-reason";
     public static final String REPORT_PROVISION_STATE_WORK_NAME = "report-provision-state";
+    @VisibleForTesting
+    static final String UNEXPECTED_PROVISION_STATE_ERROR_MESSAGE = "Unexpected provision state!";
 
     private final StatsLogger mStatsLogger;
 
@@ -170,7 +179,7 @@
             if (response.hasFatalError()) {
                 LogUtil.e(TAG,
                         "Report provision state failed: " + response + "\nRetry current step");
-                scheduler.scheduleNextProvisionFailedStepAlarm(/* shouldGoOffImmediately= */ false);
+                scheduler.scheduleNextProvisionFailedStepAlarm();
                 return Result.failure();
             }
             int daysLeftUntilReset = response.getDaysLeftUntilReset();
@@ -179,21 +188,36 @@
             }
             int nextState = response.getNextClientProvisionState();
             Futures.getUnchecked(globalParametersClient.setLastReceivedProvisionState(nextState));
-            scheduler.scheduleNextProvisionFailedStepAlarm(
-                    shouldRunNextStepImmediately(Futures.getDone(lastState), nextState));
+            onNextProvisionStateReceived(nextState, daysLeftUntilReset);
+            if (nextState == PROVISION_STATE_FACTORY_RESET) {
+                scheduler.scheduleResetDeviceAlarm();
+            } else if (nextState != PROVISION_STATE_SUCCESS) {
+                scheduler.scheduleNextProvisionFailedStepAlarm();
+            }
             mStatsLogger.logReportDeviceProvisionState();
             return Result.success();
         }, mExecutorService);
     }
 
-    @VisibleForTesting
-    static boolean shouldRunNextStepImmediately(@DeviceProvisionState int lastState,
-            @DeviceProvisionState int nextState) {
-        // Always wait before performing a retry;
-        if (nextState == DeviceProvisionState.PROVISION_STATE_RETRY) return false;
-        // Otherwise, when the user just goes through the provision UI, we should
-        // perform next step immediately.
-        return lastState == DeviceProvisionState.PROVISION_STATE_UNSPECIFIED
-                || lastState == DeviceProvisionState.PROVISION_STATE_RETRY;
+    private void onNextProvisionStateReceived(@DeviceProvisionState int provisionState,
+            int daysLeftUntilReset) {
+        switch (provisionState) {
+            case PROVISION_STATE_RETRY:
+            case PROVISION_STATE_SUCCESS:
+            case PROVISION_STATE_UNSPECIFIED:
+            case PROVISION_STATE_FACTORY_RESET:
+                // no-op
+                break;
+            case PROVISION_STATE_DISMISSIBLE_UI:
+                DeviceLockNotificationManager.sendDeviceResetNotification(mContext,
+                        daysLeftUntilReset);
+                break;
+            case PROVISION_STATE_PERSISTENT_UI:
+                DeviceLockNotificationManager
+                        .sendDeviceResetInOneDayOngoingNotification(mContext);
+                break;
+            default:
+                throw new IllegalStateException(UNEXPECTED_PROVISION_STATE_ERROR_MESSAGE);
+        }
     }
 }
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/receivers/NextProvisionFailedStepReceiver.java b/DeviceLockController/src/com/android/devicelockcontroller/receivers/NextProvisionFailedStepReceiver.java
index 76fd9c4..618aa50 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/receivers/NextProvisionFailedStepReceiver.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/receivers/NextProvisionFailedStepReceiver.java
@@ -16,9 +16,6 @@
 
 package com.android.devicelockcontroller.receivers;
 
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_DISMISSIBLE_UI;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_FACTORY_RESET;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_PERSISTENT_UI;
 import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_RETRY;
 import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_SUCCESS;
 import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_UNSPECIFIED;
@@ -31,14 +28,10 @@
 import androidx.annotation.VisibleForTesting;
 import androidx.work.WorkManager;
 
-import com.android.devicelockcontroller.activities.DeviceLockNotificationManager;
 import com.android.devicelockcontroller.policy.PolicyObjectsProvider;
 import com.android.devicelockcontroller.policy.ProvisionStateController;
 import com.android.devicelockcontroller.provision.worker.ReportDeviceProvisionStateWorker;
-import com.android.devicelockcontroller.schedule.DeviceLockControllerScheduler;
-import com.android.devicelockcontroller.schedule.DeviceLockControllerSchedulerProvider;
 import com.android.devicelockcontroller.storage.GlobalParametersClient;
-import com.android.devicelockcontroller.storage.UserParameters;
 import com.android.devicelockcontroller.util.LogUtil;
 
 import com.google.common.util.concurrent.FutureCallback;
@@ -53,8 +46,6 @@
  * A broadcast receiver to perform the next step in the provision failure flow.
  */
 public final class NextProvisionFailedStepReceiver extends BroadcastReceiver {
-    @VisibleForTesting
-    static final String UNEXPECTED_PROVISION_STATE_ERROR_MESSAGE = "Unexpected provision state!";
     public static final String TAG = "NextProvisionFailedStepReceiver";
     private final Executor mExecutor;
 
@@ -78,43 +69,21 @@
                 ((PolicyObjectsProvider) applicationContext)
                         .getProvisionStateController();
 
-        DeviceLockControllerSchedulerProvider schedulerProvider =
-                (DeviceLockControllerSchedulerProvider) applicationContext;
-        DeviceLockControllerScheduler scheduler =
-                schedulerProvider.getDeviceLockControllerScheduler();
         GlobalParametersClient globalParameters = GlobalParametersClient.getInstance();
         ListenableFuture<Boolean> needToReportFuture = Futures.transform(
                 globalParameters.getLastReceivedProvisionState(),
                 provisionState -> {
-                    switch (provisionState) {
-                        case PROVISION_STATE_RETRY:
-                            provisionStateController.postSetNextStateForEventRequest(
-                                    PROVISION_RETRY);
-                            // We can not report the state here, because we do not know the
-                            // result of retry. It will be reported after the retry finishes, no
-                            // matter it succeeds or fails.
-                            return false;
-                        case PROVISION_STATE_DISMISSIBLE_UI:
-                            int daysLeftUntilReset = UserParameters.getDaysLeftUntilReset(context);
-                            DeviceLockNotificationManager.sendDeviceResetNotification(context,
-                                    daysLeftUntilReset);
-                            return true;
-                        case PROVISION_STATE_PERSISTENT_UI:
-                            DeviceLockNotificationManager
-                                    .sendDeviceResetInOneDayOngoingNotification(context);
-                            return true;
-                        case PROVISION_STATE_FACTORY_RESET:
-                            scheduler.scheduleResetDeviceAlarm();
-                            return true;
-                        case PROVISION_STATE_SUCCESS:
-                        case PROVISION_STATE_UNSPECIFIED:
-                            return false;
-                        default:
-                            throw new IllegalStateException(
-                                    UNEXPECTED_PROVISION_STATE_ERROR_MESSAGE);
+                    if (provisionState == PROVISION_STATE_RETRY) {
+                        // We cannot report the state here because we do not know the
+                        // result of the retry. It will be reported after the retry finishes no
+                        // matter whether it succeeds or fails.
+                        provisionStateController.postSetNextStateForEventRequest(
+                                PROVISION_RETRY);
+                        return false;
                     }
+                    return !(provisionState == PROVISION_STATE_SUCCESS
+                            || provisionState == PROVISION_STATE_UNSPECIFIED);
                 }, mExecutor);
-
         Futures.addCallback(needToReportFuture, new FutureCallback<>() {
             @Override
             public void onSuccess(Boolean needToReport) {
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerScheduler.java b/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerScheduler.java
index ed98290..6e2511c 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerScheduler.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerScheduler.java
@@ -66,11 +66,8 @@
 
     /**
      * Schedule an alarm to perform next provision failed step.
-     *
-     * @param shouldRunImmediately true, if the alarm should go off immediately; false, if the
-     *                             alarm should go off after a default delay.
      */
-    void scheduleNextProvisionFailedStepAlarm(boolean shouldRunImmediately);
+    void scheduleNextProvisionFailedStepAlarm();
 
     /**
      * Notify the scheduler that device reboot when provision has failed.
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImpl.java b/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImpl.java
index 56934d4..aeaf3b8 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImpl.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImpl.java
@@ -16,13 +16,13 @@
 
 package com.android.devicelockcontroller.schedule;
 
+import static com.android.devicelockcontroller.WorkManagerExceptionHandler.AlarmReason;
 import static com.android.devicelockcontroller.common.DeviceLockConstants.MANDATORY_PROVISION_DEVICE_RESET_COUNTDOWN_MINUTE;
 import static com.android.devicelockcontroller.common.DeviceLockConstants.NON_MANDATORY_PROVISION_DEVICE_RESET_COUNTDOWN_MINUTE;
 import static com.android.devicelockcontroller.policy.ProvisionStateController.ProvisionState.PROVISION_FAILED;
 import static com.android.devicelockcontroller.policy.ProvisionStateController.ProvisionState.PROVISION_PAUSED;
 import static com.android.devicelockcontroller.policy.ProvisionStateController.ProvisionState.UNPROVISIONED;
 import static com.android.devicelockcontroller.provision.worker.AbstractCheckInWorker.BACKOFF_DELAY;
-import static com.android.devicelockcontroller.WorkManagerExceptionHandler.AlarmReason;
 
 import android.app.AlarmManager;
 import android.app.PendingIntent;
@@ -366,10 +366,9 @@
     }
 
     @Override
-    public void scheduleNextProvisionFailedStepAlarm(boolean shouldRunImmediately) {
+    public void scheduleNextProvisionFailedStepAlarm() {
         LogUtil.d(TAG,
-                "Scheduling next provision failed step alarm. Run immediately: "
-                        + shouldRunImmediately);
+                "Scheduling next provision failed step alarm");
         long lastTimestamp = UserParameters.getNextProvisionFailedStepTimeMills(mContext);
         long nextTimestamp;
         if (lastTimestamp == 0) {
@@ -379,7 +378,7 @@
                 DEBUG_DEVICELOCK_REPORT_INTERVAL_MINUTES,
                 PROVISION_STATE_REPORT_INTERVAL_DEFAULT_MINUTES)
                 : PROVISION_STATE_REPORT_INTERVAL_DEFAULT_MINUTES;
-        Duration delay = shouldRunImmediately ? Duration.ZERO : Duration.ofMinutes(minutes);
+        Duration delay = Duration.ofMinutes(minutes);
         nextTimestamp = lastTimestamp + delay.toMillis();
         scheduleNextProvisionFailedStepAlarm(
                 Duration.between(Instant.now(mClock), Instant.ofEpochMilli(nextTimestamp)));
diff --git a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorkerResponseTest.java b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorkerResponseTest.java
deleted file mode 100644
index ef72a12..0000000
--- a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorkerResponseTest.java
+++ /dev/null
@@ -1,79 +0,0 @@
-/*
- * 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 com.android.devicelockcontroller.provision.worker;
-
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_DISMISSIBLE_UI;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_FACTORY_RESET;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_PERSISTENT_UI;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_RETRY;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_SUCCESS;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_UNSPECIFIED;
-
-import com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState;
-
-import com.google.common.truth.Truth;
-
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.robolectric.ParameterizedRobolectricTestRunner;
-
-import java.util.Arrays;
-import java.util.List;
-
-@RunWith(ParameterizedRobolectricTestRunner.class)
-public final class ReportDeviceProvisionStateWorkerResponseTest {
-    @ParameterizedRobolectricTestRunner.Parameter
-    public boolean mWithDelay;
-
-    @ParameterizedRobolectricTestRunner.Parameter(1)
-    @DeviceProvisionState
-    public int mLastState;
-
-    @ParameterizedRobolectricTestRunner.Parameter(2)
-    @DeviceProvisionState
-    public int mNextState;
-
-    /**
-     * Parameters
-     */
-    @ParameterizedRobolectricTestRunner.Parameters(name =
-            "Schedule next step with delay: {0}, when last step is {1} and next step is {2}.")
-    public static List<Object[]> parameters() {
-        return Arrays.asList(new Object[][]{
-                {false, PROVISION_STATE_UNSPECIFIED, PROVISION_STATE_RETRY},
-                {true, PROVISION_STATE_UNSPECIFIED, PROVISION_STATE_DISMISSIBLE_UI},
-                {true, PROVISION_STATE_UNSPECIFIED, PROVISION_STATE_PERSISTENT_UI},
-                {true, PROVISION_STATE_UNSPECIFIED, PROVISION_STATE_FACTORY_RESET},
-                {true, PROVISION_STATE_UNSPECIFIED, PROVISION_STATE_SUCCESS},
-                {false, PROVISION_STATE_RETRY, PROVISION_STATE_RETRY},
-                {true, PROVISION_STATE_RETRY, PROVISION_STATE_DISMISSIBLE_UI},
-                {true, PROVISION_STATE_RETRY, PROVISION_STATE_PERSISTENT_UI},
-                {true, PROVISION_STATE_RETRY, PROVISION_STATE_FACTORY_RESET},
-                {true, PROVISION_STATE_RETRY, PROVISION_STATE_SUCCESS},
-                {false, PROVISION_STATE_DISMISSIBLE_UI, PROVISION_STATE_DISMISSIBLE_UI},
-                {false, PROVISION_STATE_DISMISSIBLE_UI, PROVISION_STATE_PERSISTENT_UI},
-                {false, PROVISION_STATE_DISMISSIBLE_UI, PROVISION_STATE_FACTORY_RESET},
-                {false, PROVISION_STATE_PERSISTENT_UI, PROVISION_STATE_FACTORY_RESET},
-        });
-    }
-
-    @Test
-    public void shouldScheduleWithDelay_returnCorrectValue() {
-        Truth.assertThat(ReportDeviceProvisionStateWorker.shouldRunNextStepImmediately(mLastState,
-                mNextState)).isEqualTo(mWithDelay);
-    }
-}
diff --git a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorkerTest.java b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorkerTest.java
index 33400b2..aa1b854 100644
--- a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorkerTest.java
+++ b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/provision/worker/ReportDeviceProvisionStateWorkerTest.java
@@ -16,19 +16,25 @@
 
 package com.android.devicelockcontroller.provision.worker;
 
+import static com.android.devicelockcontroller.activities.DeviceLockNotificationManager.DEVICE_RESET_NOTIFICATION_ID;
+import static com.android.devicelockcontroller.activities.DeviceLockNotificationManager.DEVICE_RESET_NOTIFICATION_TAG;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_DISMISSIBLE_UI;
 import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_FACTORY_RESET;
-import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_SUCCESS;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_PERSISTENT_UI;
+import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceProvisionState.PROVISION_STATE_RETRY;
 
 import static com.google.common.truth.Truth.assertThat;
 
 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.verify;
 import static org.mockito.Mockito.when;
 
+import android.app.NotificationManager;
 import android.content.Context;
+import android.os.Looper;
+import android.service.notification.StatusBarNotification;
 
 import androidx.annotation.NonNull;
 import androidx.test.core.app.ApplicationProvider;
@@ -47,9 +53,11 @@
 import com.android.devicelockcontroller.stats.StatsLogger;
 import com.android.devicelockcontroller.stats.StatsLoggerProvider;
 import com.android.devicelockcontroller.storage.GlobalParametersClient;
+import com.android.devicelockcontroller.storage.SetupParametersClient;
 import com.android.devicelockcontroller.storage.UserParameters;
 
 import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
 
 import org.junit.Before;
@@ -60,7 +68,10 @@
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 import org.robolectric.RobolectricTestRunner;
+import org.robolectric.Shadows;
+import org.robolectric.shadows.ShadowNotificationManager;
 
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
 
 @RunWith(RobolectricTestRunner.class)
@@ -75,6 +86,8 @@
     private StatsLogger mStatsLogger;
     private ReportDeviceProvisionStateWorker mWorker;
     private TestDeviceLockControllerApplication mTestApp;
+    private ListeningExecutorService mExecutorService =
+            MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor());
 
     @Before
     public void setUp() throws Exception {
@@ -107,6 +120,7 @@
         StatsLoggerProvider loggerProvider =
                 (StatsLoggerProvider) mTestApp.getApplicationContext();
         mStatsLogger = loggerProvider.getStatsLogger();
+        SetupParametersClient.getInstance(mTestApp, mExecutorService);
     }
 
     @Test
@@ -123,14 +137,96 @@
         when(mResponse.hasFatalError()).thenReturn(true);
 
         assertThat(Futures.getUnchecked(mWorker.startWork())).isEqualTo(Result.failure());
-        verify(mTestApp.getDeviceLockControllerScheduler()).scheduleNextProvisionFailedStepAlarm(
-                /* shouldGoOffImmediately= */ eq(false));
+        verify(mTestApp.getDeviceLockControllerScheduler()).scheduleNextProvisionFailedStepAlarm();
         // THEN report provisioning state or complete was NOT logged
         verify(mStatsLogger, never()).logReportDeviceProvisionState();
     }
 
     @Test
-    public void doWork_responseIsSuccessful_globalParametersShouldBeSet_returnSuccessAndLogReportState()
+    public void doWork_responseIsSuccessful_globalParametersSetAndEventLogged() throws Exception {
+        when(mResponse.isSuccessful()).thenReturn(true);
+        when(mResponse.getNextClientProvisionState()).thenReturn(PROVISION_STATE_RETRY);
+        when(mResponse.getDaysLeftUntilReset()).thenReturn(TEST_DAYS_LEFT_UNTIL_RESET);
+
+        assertThat(Futures.getUnchecked(mWorker.startWork())).isEqualTo(Result.success());
+
+        // THEN parameters are saved and report state was logged
+        GlobalParametersClient globalParameters = GlobalParametersClient.getInstance();
+        assertThat(globalParameters.getLastReceivedProvisionState().get()).isEqualTo(
+                PROVISION_STATE_RETRY);
+        Executors.newSingleThreadExecutor().submit(
+                () -> assertThat(UserParameters.getDaysLeftUntilReset(mTestApp)).isEqualTo(
+                        TEST_DAYS_LEFT_UNTIL_RESET)).get();
+        verify(mStatsLogger).logReportDeviceProvisionState();
+    }
+
+    @Test
+    public void doWork_retryState_schedulesAlarm() {
+        when(mResponse.isSuccessful()).thenReturn(true);
+        when(mResponse.getNextClientProvisionState()).thenReturn(PROVISION_STATE_RETRY);
+        when(mResponse.getDaysLeftUntilReset()).thenReturn(TEST_DAYS_LEFT_UNTIL_RESET);
+
+        assertThat(Futures.getUnchecked(mWorker.startWork())).isEqualTo(Result.success());
+
+        // THEN we schedule to try again later
+        verify(mTestApp.getDeviceLockControllerScheduler()).scheduleNextProvisionFailedStepAlarm();
+    }
+
+    @Test
+    public void doWork_dismissibleUiState_schedulesAlarmAndSendsNotification() throws Exception {
+        when(mResponse.isSuccessful()).thenReturn(true);
+        when(mResponse.getNextClientProvisionState()).thenReturn(PROVISION_STATE_DISMISSIBLE_UI);
+        when(mResponse.getDaysLeftUntilReset()).thenReturn(TEST_DAYS_LEFT_UNTIL_RESET);
+
+        assertThat(Futures.getUnchecked(mWorker.startWork())).isEqualTo(Result.success());
+
+        CountDownLatch latch = new CountDownLatch(1);
+        Futures.getUnchecked(mExecutorService.submit(latch::countDown));
+        latch.await();
+        Shadows.shadowOf(Looper.getMainLooper()).idle();
+
+        // THEN we schedule to try again later and send a notification to the user
+        verify(mTestApp.getDeviceLockControllerScheduler()).scheduleNextProvisionFailedStepAlarm();
+
+        ShadowNotificationManager shadowNotificationManager = Shadows.shadowOf(
+                mTestApp.getSystemService(NotificationManager.class));
+        StatusBarNotification[] activeNotifs = shadowNotificationManager.getActiveNotifications();
+        assertThat(activeNotifs.length).isGreaterThan(0);
+        StatusBarNotification notif = activeNotifs[0];
+        assertThat(notif.getTag()).isEqualTo(DEVICE_RESET_NOTIFICATION_TAG);
+        assertThat(notif.getId()).isEqualTo(DEVICE_RESET_NOTIFICATION_ID);
+        assertThat(notif.isOngoing()).isFalse();
+    }
+
+    @Test
+    public void doWork_persistentUiState_schedulesAlarmAndSendsOngoingNotification()
+            throws Exception {
+        when(mResponse.isSuccessful()).thenReturn(true);
+        when(mResponse.getNextClientProvisionState()).thenReturn(PROVISION_STATE_PERSISTENT_UI);
+        when(mResponse.getDaysLeftUntilReset()).thenReturn(TEST_DAYS_LEFT_UNTIL_RESET);
+
+        assertThat(Futures.getUnchecked(mWorker.startWork())).isEqualTo(Result.success());
+
+        CountDownLatch latch = new CountDownLatch(1);
+        Futures.getUnchecked(mExecutorService.submit(latch::countDown));
+        latch.await();
+        Shadows.shadowOf(Looper.getMainLooper()).idle();
+
+        // THEN we schedule to try again later and send an undismissable notification to the user
+        verify(mTestApp.getDeviceLockControllerScheduler()).scheduleNextProvisionFailedStepAlarm();
+
+        ShadowNotificationManager shadowNotificationManager = Shadows.shadowOf(
+                mTestApp.getSystemService(NotificationManager.class));
+        StatusBarNotification[] activeNotifs = shadowNotificationManager.getActiveNotifications();
+        assertThat(activeNotifs.length).isGreaterThan(0);
+        StatusBarNotification notif = activeNotifs[0];
+        assertThat(notif.getTag()).isEqualTo(DEVICE_RESET_NOTIFICATION_TAG);
+        assertThat(notif.getId()).isEqualTo(DEVICE_RESET_NOTIFICATION_ID);
+        assertThat(notif.isOngoing()).isTrue();
+    }
+
+    @Test
+    public void doWork_factoryResetState_schedulesResetDeviceAlarm()
             throws Exception {
         when(mResponse.isSuccessful()).thenReturn(true);
         when(mResponse.getNextClientProvisionState()).thenReturn(PROVISION_STATE_FACTORY_RESET);
@@ -138,16 +234,13 @@
 
         assertThat(Futures.getUnchecked(mWorker.startWork())).isEqualTo(Result.success());
 
-        GlobalParametersClient globalParameters = GlobalParametersClient.getInstance();
-        assertThat(globalParameters.getLastReceivedProvisionState().get()).isEqualTo(
-                PROVISION_STATE_FACTORY_RESET);
-        Executors.newSingleThreadExecutor().submit(
-                () -> assertThat(UserParameters.getDaysLeftUntilReset(mTestApp)).isEqualTo(
-                        TEST_DAYS_LEFT_UNTIL_RESET)).get();
+        CountDownLatch latch = new CountDownLatch(1);
+        Futures.getUnchecked(mExecutorService.submit(latch::countDown));
+        latch.await();
+        Shadows.shadowOf(Looper.getMainLooper()).idle();
 
-        verify(mTestApp.getDeviceLockControllerScheduler()).scheduleNextProvisionFailedStepAlarm(
-                /* shouldGoOffImmediately= */ eq(true));
-        // THEN report provisioning state was logged
-        verify(mStatsLogger).logReportDeviceProvisionState();
+        // THEN we schedule reset.
+        // Note that the scheduler class sends the notification
+        verify(mTestApp.getDeviceLockControllerScheduler()).scheduleResetDeviceAlarm();
     }
 }
diff --git a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImplTest.java b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImplTest.java
index b5269ba..849d016 100644
--- a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImplTest.java
+++ b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/schedule/DeviceLockControllerSchedulerImplTest.java
@@ -39,7 +39,6 @@
 import com.android.devicelockcontroller.TestDeviceLockControllerApplication;
 import com.android.devicelockcontroller.common.DeviceLockConstants;
 import com.android.devicelockcontroller.provision.worker.DeviceCheckInWorker;
-import com.android.devicelockcontroller.shadows.ShadowBuild;
 import com.android.devicelockcontroller.storage.UserParameters;
 import com.android.devicelockcontroller.util.ThreadUtils;
 
@@ -385,7 +384,7 @@
     }
 
     @Test
-    public void scheduleNextProvisionFailedStepAlarm_initialStep_defaultDelay() {
+    public void scheduleNextProvisionFailedStepAlarm_initialStep_setsAlarm() {
         // GIVEN no alarm is scheduled
         ShadowAlarmManager alarmManager = Shadows.shadowOf(
                 mTestApp.getSystemService(AlarmManager.class));
@@ -396,9 +395,7 @@
                 UserParameters.getNextProvisionFailedStepTimeMills(mTestApp)).isEqualTo(0));
 
         // WHEN schedule next provision failed step alarm
-        runBySequentialExecutor(
-                () -> mScheduler.scheduleNextProvisionFailedStepAlarm(
-                        /* shouldGoOffImmediately= */ false));
+        runBySequentialExecutor(() -> mScheduler.scheduleNextProvisionFailedStepAlarm());
 
         // THEN correct alarm should be scheduled
         PendingIntent actualPendingIntent = alarmManager.peekNextScheduledAlarm().operation;
@@ -416,70 +413,7 @@
     }
 
     @Test
-    public void scheduleNextProvisionFailedStepAlarm_initialStep_noDelay() {
-        // GIVEN no alarm is scheduled
-        ShadowAlarmManager alarmManager = Shadows.shadowOf(
-                mTestApp.getSystemService(AlarmManager.class));
-        assertThat(alarmManager.peekNextScheduledAlarm()).isNull();
-
-        // GIVEN no existing timestamp
-        runBySequentialExecutor(() -> assertThat(
-                UserParameters.getNextProvisionFailedStepTimeMills(mTestApp)).isEqualTo(0));
-
-        // WHEN schedule next provision failed step alarm
-        runBySequentialExecutor(
-                () -> mScheduler.scheduleNextProvisionFailedStepAlarm(
-                        /* shouldGoOffImmediately= */ true));
-
-        // THEN correct alarm should be scheduled
-        PendingIntent actualPendingIntent = alarmManager.peekNextScheduledAlarm().operation;
-        assertThat(actualPendingIntent.isBroadcast()).isTrue();
-
-        // THEN alarm should be scheduled at correct time
-        long actualTriggerTime = alarmManager.peekNextScheduledAlarm().triggerAtTime;
-        assertThat(actualTriggerTime).isEqualTo(SystemClock.elapsedRealtime());
-
-        // THEN expected trigger time should be stored in storage
-        runBySequentialExecutor(() -> assertThat(
-                UserParameters.getNextProvisionFailedStepTimeMills(mTestApp)).isEqualTo(
-                TEST_CURRENT_TIME_MILLIS));
-    }
-
-    @Test
-    public void scheduleNextProvisionFailedStepAlarm_debuggable_initialStep_noDelay() {
-        // GIVEN build is debuggable build
-        ShadowBuild.setIsDebuggable(true);
-
-        // GIVEN no alarm is scheduled
-        ShadowAlarmManager alarmManager = Shadows.shadowOf(
-                mTestApp.getSystemService(AlarmManager.class));
-        assertThat(alarmManager.peekNextScheduledAlarm()).isNull();
-
-        // GIVEN no existing timestamp
-        runBySequentialExecutor(() -> assertThat(
-                UserParameters.getNextProvisionFailedStepTimeMills(mTestApp)).isEqualTo(0));
-
-        // WHEN schedule next provision failed step alarm
-        runBySequentialExecutor(
-                () -> mScheduler.scheduleNextProvisionFailedStepAlarm(
-                        /* shouldGoOffImmediately= */ true));
-
-        // THEN correct alarm should be scheduled
-        PendingIntent actualPendingIntent = alarmManager.peekNextScheduledAlarm().operation;
-        assertThat(actualPendingIntent.isBroadcast()).isTrue();
-
-        // THEN alarm should be scheduled at correct time
-        long actualTriggerTime = alarmManager.peekNextScheduledAlarm().triggerAtTime;
-        assertThat(actualTriggerTime).isEqualTo(SystemClock.elapsedRealtime());
-
-        // THEN expected trigger time should be stored in storage
-        runBySequentialExecutor(() -> assertThat(
-                UserParameters.getNextProvisionFailedStepTimeMills(mTestApp)).isEqualTo(
-                TEST_CURRENT_TIME_MILLIS));
-    }
-
-    @Test
-    public void scheduleNextProvisionFailedStepAlarm_followUpStep_defaultDelay() {
+    public void scheduleNextProvisionFailedStepAlarm_followUpStep_setsAlarm() {
         // GIVEN no alarm is scheduled
         ShadowAlarmManager alarmManager = Shadows.shadowOf(
                 mTestApp.getSystemService(AlarmManager.class));
@@ -490,9 +424,7 @@
                 TEST_NEXT_PROVISION_FAILED_STEP_TIME_MILLIS);
 
         // WHEN schedule next provision failed step alarm
-        runBySequentialExecutor(
-                () -> mScheduler.scheduleNextProvisionFailedStepAlarm(
-                        /* shouldGoOffImmediately= */ false));
+        runBySequentialExecutor(() -> mScheduler.scheduleNextProvisionFailedStepAlarm());
 
         // THEN correct alarm should be scheduled
         PendingIntent actualPendingIntent = alarmManager.peekNextScheduledAlarm().operation;
@@ -514,40 +446,6 @@
     }
 
     @Test
-    public void scheduleNextProvisionFailedStepAlarm_followUpStep_noDelay() {
-        // GIVEN no alarm is scheduled
-        ShadowAlarmManager alarmManager = Shadows.shadowOf(
-                mTestApp.getSystemService(AlarmManager.class));
-        assertThat(alarmManager.peekNextScheduledAlarm()).isNull();
-
-        // GIVEN timestamp exists for last performed step.
-        UserParameters.setNextProvisionFailedStepTimeMills(mTestApp,
-                TEST_NEXT_PROVISION_FAILED_STEP_TIME_MILLIS);
-
-        // WHEN schedule next provision failed step alarm
-        runBySequentialExecutor(
-                () -> mScheduler.scheduleNextProvisionFailedStepAlarm(
-                        /* shouldGoOffImmediately= */ true));
-
-        // THEN correct alarm should be scheduled
-        PendingIntent actualPendingIntent = alarmManager.peekNextScheduledAlarm().operation;
-        assertThat(actualPendingIntent.isBroadcast()).isTrue();
-
-        // THEN alarm should be scheduled at correct time
-        long actualTriggerTime = alarmManager.peekNextScheduledAlarm().triggerAtTime;
-        long expectedTriggerTime =
-                TEST_NEXT_PROVISION_FAILED_STEP_TIME_MILLIS - TEST_CURRENT_TIME_MILLIS
-                        + SystemClock.elapsedRealtime();
-        assertThat(actualTriggerTime).isEqualTo(expectedTriggerTime);
-
-
-        // THEN expected trigger time should be stored in storage
-        runBySequentialExecutor(() -> assertThat(
-                UserParameters.getNextProvisionFailedStepTimeMills(mTestApp)).isEqualTo(
-                TEST_NEXT_PROVISION_FAILED_STEP_TIME_MILLIS));
-    }
-
-    @Test
     public void scheduleResetDeviceAlarm() {
         // GIVEN no alarm is scheduled
         ShadowAlarmManager alarmManager = Shadows.shadowOf(