diff --git a/src/com/android/customization/picker/notifications/domain/interactor/NotificationsSnapshotRestorer.kt b/src/com/android/customization/picker/notifications/domain/interactor/NotificationsSnapshotRestorer.kt
index 09d2d29..e82d00a 100644
--- a/src/com/android/customization/picker/notifications/domain/interactor/NotificationsSnapshotRestorer.kt
+++ b/src/com/android/customization/picker/notifications/domain/interactor/NotificationsSnapshotRestorer.kt
@@ -40,16 +40,19 @@
         store: SnapshotStore,
     ): RestorableSnapshot {
         snapshotStore = store
-        backgroundScope.launch {
-            interactor.isShowNotificationsOnLockScreenEnabled.collect {
-                storeSnapshot(
-                    NotificationSnapshotModel(isShowNotificationsOnLockScreenEnabled = it)
-                )
-            }
-        }
+        // The initial snapshot should be returned and stored before storing additional snapshots.
         return snapshot(
-            NotificationSnapshotModel(interactor.isShowNotificationsOnLockScreenEnabled.value)
-        )
+                NotificationSnapshotModel(interactor.isShowNotificationsOnLockScreenEnabled().value)
+            )
+            .also {
+                backgroundScope.launch {
+                    interactor.isShowNotificationsOnLockScreenEnabled().collect {
+                        storeSnapshot(
+                            NotificationSnapshotModel(isShowNotificationsOnLockScreenEnabled = it)
+                        )
+                    }
+                }
+            }
     }
 
     override suspend fun restoreToSnapshot(snapshot: RestorableSnapshot) {
diff --git a/src/com/android/customization/picker/notifications/ui/binder/NotificationSectionBinder.kt b/src/com/android/customization/picker/notifications/ui/binder/NotificationSectionBinder.kt
index 8f2727c..baf7bbf 100644
--- a/src/com/android/customization/picker/notifications/ui/binder/NotificationSectionBinder.kt
+++ b/src/com/android/customization/picker/notifications/ui/binder/NotificationSectionBinder.kt
@@ -44,7 +44,7 @@
 
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                launch { viewModel.isSwitchOn.collect { switch.isChecked = it } }
+                launch { viewModel.isSwitchOn().collect { switch.isChecked = it } }
             }
         }
     }
diff --git a/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt b/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt
index 5db7626..022dc9b 100644
--- a/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt
+++ b/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModel.kt
@@ -35,14 +35,14 @@
 ) : ViewModel() {
 
     /** Whether the switch should be on. */
-    val isSwitchOn: Flow<Boolean> = interactor.isShowNotificationsOnLockScreenEnabled
+    suspend fun isSwitchOn(): Flow<Boolean> = interactor.isShowNotificationsOnLockScreenEnabled()
 
     /** Notifies that the section has been clicked. */
     fun onClicked() {
         viewModelScope.launch {
             interactor.toggleShowNotificationsOnLockscreenEnabled()
             logger.logLockScreenNotificationApplied(
-                interactor.isShowNotificationsOnLockScreenEnabled.value
+                interactor.isShowNotificationsOnLockScreenEnabled().value
             )
         }
     }
diff --git a/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt b/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
index e8bac6a..c61267c 100644
--- a/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
+++ b/src/com/android/wallpaper/customization/ui/util/ThemePickerCustomizationOptionUtil.kt
@@ -174,8 +174,8 @@
         layoutInflater: LayoutInflater,
     ): View =
         when (option) {
-            ThemePickerLockCustomizationOption.CLOCK -> R.layout.bottom_sheet_clock
-            ThemePickerLockCustomizationOption.SHORTCUTS -> R.layout.bottom_sheet_shortcut
+            ThemePickerLockCustomizationOption.CLOCK -> com.android.wallpaper.R.layout.bottom_sheet_clock
+            ThemePickerLockCustomizationOption.SHORTCUTS -> com.android.wallpaper.R.layout.bottom_sheet_shortcut
             else ->
                 throw IllegalStateException(
                     "Customization option $option does not have a bottom sheet view"
diff --git a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
index 5ed8962..a12190f 100644
--- a/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
+++ b/tests/module/src/com/android/wallpaper/ThemePickerTestModule.kt
@@ -36,11 +36,11 @@
 import com.android.wallpaper.picker.customization.ui.binder.CustomizationOptionsBinder
 import com.android.wallpaper.picker.customization.ui.binder.DefaultCustomizationOptionsBinder
 import com.android.wallpaper.picker.di.modules.EffectsModule
-import com.android.wallpaper.picker.preview.data.util.FakeLiveWallpaperDownloader
 import com.android.wallpaper.picker.preview.data.util.LiveWallpaperDownloader
 import com.android.wallpaper.picker.preview.ui.util.DefaultImageEffectDialogUtil
 import com.android.wallpaper.picker.preview.ui.util.ImageEffectDialogUtil
 import com.android.wallpaper.testing.FakeDefaultRequester
+import com.android.wallpaper.testing.FakeLiveWallpaperDownloader
 import com.android.wallpaper.testing.TestPartnerProvider
 import com.android.wallpaper.util.converter.DefaultWallpaperModelFactory
 import com.android.wallpaper.util.converter.WallpaperModelFactory
diff --git a/tests/robotests/src/com/android/customization/model/notifications/domain/interactor/NotificationsSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/notifications/domain/interactor/NotificationsSnapshotRestorerTest.kt
new file mode 100644
index 0000000..bf8cfda
--- /dev/null
+++ b/tests/robotests/src/com/android/customization/model/notifications/domain/interactor/NotificationsSnapshotRestorerTest.kt
@@ -0,0 +1,128 @@
+/*
+ * Copyright (C) 2024 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.model.notifications.domain.interactor
+
+import android.provider.Settings
+import androidx.test.filters.SmallTest
+import com.android.customization.picker.notifications.domain.interactor.NotificationsSnapshotRestorer
+import com.android.systemui.shared.notifications.data.repository.NotificationSettingsRepository
+import com.android.systemui.shared.notifications.domain.interactor.NotificationSettingsInteractor
+import com.android.systemui.shared.settings.data.repository.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.test.TestScope
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@OptIn(ExperimentalCoroutinesApi::class)
+@SmallTest
+@RunWith(RobolectricTestRunner::class)
+class NotificationsSnapshotRestorerTest {
+
+    private lateinit var underTest: NotificationsSnapshotRestorer
+    private lateinit var fakeSecureSettingsRepository: FakeSecureSettingsRepository
+    private lateinit var interactor: NotificationSettingsInteractor
+
+    private lateinit var testScope: TestScope
+
+    @Before
+    fun setUp() {
+        val testDispatcher = StandardTestDispatcher()
+        Dispatchers.setMain(testDispatcher)
+        testScope = TestScope(testDispatcher)
+        fakeSecureSettingsRepository = FakeSecureSettingsRepository()
+        interactor =
+            NotificationSettingsInteractor(
+                repository =
+                    NotificationSettingsRepository(
+                        scope = testScope.backgroundScope,
+                        backgroundDispatcher = testDispatcher,
+                        secureSettingsRepository = fakeSecureSettingsRepository,
+                    ),
+            )
+        underTest =
+            NotificationsSnapshotRestorer(
+                interactor = interactor,
+                backgroundScope = testScope.backgroundScope
+            )
+    }
+
+    @Test
+    fun setUpAndRestore_Active() =
+        testScope.runTest {
+            fakeSecureSettingsRepository.setInt(Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS, 1)
+            val showNotifs = collectLastValue(interactor.isShowNotificationsOnLockScreenEnabled())
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+
+            assertThat(showNotifs()).isTrue()
+        }
+
+    @Test
+    fun setUpAndRestore_Inactive() =
+        testScope.runTest {
+            fakeSecureSettingsRepository.setInt(Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS, 0)
+            val showNotifs = collectLastValue(interactor.isShowNotificationsOnLockScreenEnabled())
+
+            val store = FakeSnapshotStore()
+            store.store(underTest.setUpSnapshotRestorer(store = store))
+            val initialSnapshot = store.retrieve()
+            underTest.restoreToSnapshot(snapshot = initialSnapshot)
+
+            assertThat(showNotifs()).isFalse()
+        }
+
+    @Test
+    fun setUp_deactivate_restoreToActive() = runTest {
+        fakeSecureSettingsRepository.setInt(Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS, 1)
+        val showNotifs = collectLastValue(interactor.isShowNotificationsOnLockScreenEnabled())
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        fakeSecureSettingsRepository.setInt(Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS, 0)
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+
+        assertThat(showNotifs()).isTrue()
+    }
+
+    @Test
+    fun setUp_activate_restoreToInactive() = runTest {
+        fakeSecureSettingsRepository.setInt(Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS, 0)
+        val showNotifs = collectLastValue(interactor.isShowNotificationsOnLockScreenEnabled())
+        val store = FakeSnapshotStore()
+        store.store(underTest.setUpSnapshotRestorer(store = store))
+        val initialSnapshot = store.retrieve()
+
+        fakeSecureSettingsRepository.setInt(Settings.Secure.LOCK_SCREEN_SHOW_NOTIFICATIONS, 1)
+        underTest.restoreToSnapshot(snapshot = initialSnapshot)
+
+        assertThat(showNotifs()).isFalse()
+    }
+}
diff --git a/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
index 8966de5..e9f7ffd 100644
--- a/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
@@ -80,7 +80,7 @@
     @Test
     fun `toggles back and forth`() =
         testScope.runTest {
-            val isSwitchOn = collectLastValue(underTest.isSwitchOn)
+            val isSwitchOn = collectLastValue(underTest.isSwitchOn())
 
             val initialIsSwitchOn = isSwitchOn()