diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index c958123..5874e9d 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -34,6 +34,7 @@
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel
 import com.android.customization.picker.notifications.data.repository.NotificationsRepository
 import com.android.customization.picker.notifications.domain.interactor.NotificationsInteractor
+import com.android.customization.picker.notifications.domain.interactor.NotificationsSnapshotRestorer
 import com.android.customization.picker.notifications.ui.viewmodel.NotificationSectionViewModel
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
@@ -71,6 +72,7 @@
     private var fragmentFactory: FragmentFactory? = null
     private var keyguardQuickAffordanceSnapshotRestorer: KeyguardQuickAffordanceSnapshotRestorer? =
         null
+    private var notificationsSnapshotRestorer: NotificationsSnapshotRestorer? = null
     private var clockRegistryProvider: ClockRegistryProvider? = null
     private var clockPickerInteractor: ClockPickerInteractor? = null
     private var clockSectionViewModel: ClockSectionViewModel? = null
@@ -146,6 +148,7 @@
             this[KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER] =
                 getKeyguardQuickAffordanceSnapshotRestorer(context)
             this[KEY_WALLPAPER_SNAPSHOT_RESTORER] = getWallpaperSnapshotRestorer(context)
+            this[KEY_NOTIFICATIONS_SNAPSHOT_RESTORER] = getNotificationsSnapshotRestorer(context)
         }
     }
 
@@ -225,6 +228,17 @@
                 .also { keyguardQuickAffordanceSnapshotRestorer = it }
     }
 
+    private fun getNotificationsSnapshotRestorer(context: Context): NotificationsSnapshotRestorer {
+        return notificationsSnapshotRestorer
+            ?: NotificationsSnapshotRestorer(
+                    interactor =
+                        getNotificationsInteractor(
+                            context = context,
+                        ),
+                )
+                .also { notificationsSnapshotRestorer = it }
+    }
+
     override fun getClockRegistryProvider(context: Context): ClockRegistryProvider {
         return clockRegistryProvider
             ?: ClockRegistryProvider(context).also { clockRegistryProvider = it }
@@ -265,7 +279,8 @@
                             scope = GlobalScope,
                             backgroundDispatcher = Dispatchers.IO,
                             secureSettingsRepository = getSecureSettingsRepository(context),
-                        )
+                        ),
+                    snapshotRestorer = { getNotificationsSnapshotRestorer(context) },
                 )
                 .also { notificationsInteractor = it }
     }
@@ -297,11 +312,13 @@
             WallpaperPicker2Injector.MIN_SNAPSHOT_RESTORER_KEY
         @JvmStatic
         private val KEY_WALLPAPER_SNAPSHOT_RESTORER = KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER + 1
+        @JvmStatic
+        private val KEY_NOTIFICATIONS_SNAPSHOT_RESTORER = KEY_WALLPAPER_SNAPSHOT_RESTORER + 1
 
         /**
          * When this injector is overridden, this is the minimal value that should be used by
          * restorers returns in [getSnapshotRestorers].
          */
-        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_WALLPAPER_SNAPSHOT_RESTORER + 1
+        @JvmStatic protected val MIN_SNAPSHOT_RESTORER_KEY = KEY_NOTIFICATIONS_SNAPSHOT_RESTORER + 1
     }
 }
diff --git a/src/com/android/customization/picker/notifications/data/repository/NotificationsRepository.kt b/src/com/android/customization/picker/notifications/data/repository/NotificationsRepository.kt
index 5b03cbe..c75ddce 100644
--- a/src/com/android/customization/picker/notifications/data/repository/NotificationsRepository.kt
+++ b/src/com/android/customization/picker/notifications/data/repository/NotificationsRepository.kt
@@ -51,22 +51,24 @@
                 replay = 1,
             )
 
-    /** Updates the setting to show or hide notifications on the lock screen. */
-    suspend fun setShowNotificationsOnLockScreenEnabled(isEnabled: Boolean) {
-        withContext(backgroundDispatcher) {
-            secureSettingsRepository.set(
-                name = Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS,
-                value = if (isEnabled) 1 else 0,
+    suspend fun getSettings(): NotificationSettingsModel {
+        return withContext(backgroundDispatcher) {
+            NotificationSettingsModel(
+                isShowNotificationsOnLockScreenEnabled =
+                    secureSettingsRepository.get(
+                        name = Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS,
+                        defaultValue = 0,
+                    ) == 1
             )
         }
     }
 
-    suspend fun isShowNotificationsOnLockScreenEnabled(): Boolean {
-        return withContext(backgroundDispatcher) {
-            secureSettingsRepository.get(
+    suspend fun setSettings(model: NotificationSettingsModel) {
+        withContext(backgroundDispatcher) {
+            secureSettingsRepository.set(
                 name = Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS,
-                defaultValue = 0,
-            ) == 1
+                value = if (model.isShowNotificationsOnLockScreenEnabled) 1 else 0,
+            )
         }
     }
 }
diff --git a/src/com/android/customization/picker/notifications/domain/interactor/NotificationsInteractor.kt b/src/com/android/customization/picker/notifications/domain/interactor/NotificationsInteractor.kt
index 24860fb..1f892f0 100644
--- a/src/com/android/customization/picker/notifications/domain/interactor/NotificationsInteractor.kt
+++ b/src/com/android/customization/picker/notifications/domain/interactor/NotificationsInteractor.kt
@@ -19,19 +19,34 @@
 
 import com.android.customization.picker.notifications.data.repository.NotificationsRepository
 import com.android.customization.picker.notifications.shared.model.NotificationSettingsModel
+import javax.inject.Provider
 import kotlinx.coroutines.flow.Flow
 
 /** Encapsulates business logic for interacting with notifications. */
 class NotificationsInteractor(
     private val repository: NotificationsRepository,
+    private val snapshotRestorer: Provider<NotificationsSnapshotRestorer>,
 ) {
     /** The current state of the notification setting. */
     val settings: Flow<NotificationSettingsModel> = repository.settings
 
     /** Toggles the setting to show or hide notifications on the lock screen. */
     suspend fun toggleShowNotificationsOnLockScreenEnabled() {
-        repository.setShowNotificationsOnLockScreenEnabled(
-            isEnabled = !repository.isShowNotificationsOnLockScreenEnabled(),
+        val currentModel = repository.getSettings()
+        setSettings(
+            currentModel.copy(
+                isShowNotificationsOnLockScreenEnabled =
+                    !currentModel.isShowNotificationsOnLockScreenEnabled,
+            )
         )
     }
+
+    suspend fun setSettings(model: NotificationSettingsModel) {
+        repository.setSettings(model)
+        snapshotRestorer.get().storeSnapshot(model)
+    }
+
+    suspend fun getSettings(): NotificationSettingsModel {
+        return repository.getSettings()
+    }
 }
diff --git a/src/com/android/customization/picker/notifications/domain/interactor/NotificationsSnapshotRestorer.kt b/src/com/android/customization/picker/notifications/domain/interactor/NotificationsSnapshotRestorer.kt
new file mode 100644
index 0000000..1fc205c
--- /dev/null
+++ b/src/com/android/customization/picker/notifications/domain/interactor/NotificationsSnapshotRestorer.kt
@@ -0,0 +1,66 @@
+/*
+ * 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.customization.picker.notifications.domain.interactor
+
+import com.android.customization.picker.notifications.shared.model.NotificationSettingsModel
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
+import com.android.wallpaper.picker.undo.domain.interactor.SnapshotStore
+import com.android.wallpaper.picker.undo.shared.model.RestorableSnapshot
+
+/** Handles state restoration for notification settings. */
+class NotificationsSnapshotRestorer(
+    private val interactor: NotificationsInteractor,
+) : SnapshotRestorer {
+
+    private lateinit var snapshotStore: SnapshotStore
+
+    fun storeSnapshot(model: NotificationSettingsModel) {
+        snapshotStore.store(snapshot(model))
+    }
+
+    override suspend fun setUpSnapshotRestorer(
+        store: SnapshotStore,
+    ): RestorableSnapshot {
+        snapshotStore = store
+        return snapshot(interactor.getSettings())
+    }
+
+    override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
+        val isShowNotificationsOnLockScreenEnabled =
+            snapshot.args[KEY_IS_SHOW_NOTIFICATIONS_ON_LOCK_SCREEN_ENABLED]?.toBoolean() ?: false
+        interactor.setSettings(
+            NotificationSettingsModel(
+                isShowNotificationsOnLockScreenEnabled = isShowNotificationsOnLockScreenEnabled,
+            )
+        )
+    }
+
+    private fun snapshot(model: NotificationSettingsModel): RestorableSnapshot {
+        return RestorableSnapshot(
+            mapOf(
+                KEY_IS_SHOW_NOTIFICATIONS_ON_LOCK_SCREEN_ENABLED to
+                    model.isShowNotificationsOnLockScreenEnabled.toString(),
+            )
+        )
+    }
+
+    companion object {
+        private const val KEY_IS_SHOW_NOTIFICATIONS_ON_LOCK_SCREEN_ENABLED =
+            "is_show_notifications_on_lock_screen_enabled"
+    }
+}
diff --git a/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt b/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
index e29bd33..be799db 100644
--- a/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
+++ b/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
@@ -79,18 +79,16 @@
         }
 
     @Test
-    fun setShowNotificationsOnLockScreenEnabled() =
+    fun setSettings() =
         testScope.runTest {
             val settings = collectLastValue(underTest.settings)
 
-            underTest.setShowNotificationsOnLockScreenEnabled(isEnabled = true)
-            assertThat(settings())
-                .isEqualTo(NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = true))
+            val model1 = NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = true)
+            underTest.setSettings(model1)
+            assertThat(settings()).isEqualTo(model1)
 
-            underTest.setShowNotificationsOnLockScreenEnabled(isEnabled = false)
-            assertThat(settings())
-                .isEqualTo(
-                    NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = false)
-                )
+            val model2 = NotificationSettingsModel(isShowNotificationsOnLockScreenEnabled = false)
+            underTest.setSettings(model2)
+            assertThat(settings()).isEqualTo(model2)
         }
 }
diff --git a/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt b/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
index c74f848..6442609 100644
--- a/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
+++ b/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
@@ -20,13 +20,16 @@
 import androidx.test.filters.SmallTest
 import com.android.customization.picker.notifications.data.repository.NotificationsRepository
 import com.android.customization.picker.notifications.domain.interactor.NotificationsInteractor
+import com.android.customization.picker.notifications.domain.interactor.NotificationsSnapshotRestorer
 import com.android.wallpaper.testing.FakeSecureSettingsRepository
+import com.android.wallpaper.testing.FakeSnapshotStore
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.resetMain
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.test.setMain
@@ -44,23 +47,32 @@
     private lateinit var underTest: NotificationSectionViewModel
 
     private lateinit var testScope: TestScope
+    private lateinit var interactor: NotificationsInteractor
 
     @Before
     fun setUp() {
-        val testDispatcher = StandardTestDispatcher()
+        val testDispatcher = UnconfinedTestDispatcher()
         Dispatchers.setMain(testDispatcher)
         testScope = TestScope(testDispatcher)
+        interactor =
+            NotificationsInteractor(
+                repository =
+                    NotificationsRepository(
+                        scope = testScope.backgroundScope,
+                        backgroundDispatcher = testDispatcher,
+                        secureSettingsRepository = FakeSecureSettingsRepository(),
+                    ),
+                snapshotRestorer = {
+                    NotificationsSnapshotRestorer(
+                            interactor = interactor,
+                        )
+                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
+                },
+            )
+
         underTest =
             NotificationSectionViewModel(
-                interactor =
-                    NotificationsInteractor(
-                        repository =
-                            NotificationsRepository(
-                                scope = testScope.backgroundScope,
-                                backgroundDispatcher = testDispatcher,
-                                secureSettingsRepository = FakeSecureSettingsRepository(),
-                            )
-                    )
+                interactor = interactor,
             )
     }