diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index d4dde59..296d2ca 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -24,6 +24,7 @@
         "junit",
         "kotlinx_coroutines_test",
         "truth",
+        "SystemUICustomizationTestUtils",
     ],
 
     libs: [
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 1ff1fc8..8966de5 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
@@ -20,16 +20,13 @@
 import androidx.test.filters.SmallTest
 import com.android.customization.module.logging.TestThemesUserEventLogger
 import com.android.customization.module.logging.ThemesUserEventLogger
-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.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.collectLastValue
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
 import kotlinx.coroutines.test.resetMain
@@ -51,7 +48,7 @@
     private lateinit var underTest: NotificationSectionViewModel
 
     private lateinit var testScope: TestScope
-    private lateinit var interactor: NotificationsInteractor
+    private lateinit var interactor: NotificationSettingsInteractor
 
     @Before
     fun setUp() {
@@ -59,19 +56,13 @@
         Dispatchers.setMain(testDispatcher)
         testScope = TestScope(testDispatcher)
         interactor =
-            NotificationsInteractor(
+            NotificationSettingsInteractor(
                 repository =
-                    NotificationsRepository(
+                    NotificationSettingsRepository(
                         scope = testScope.backgroundScope,
                         backgroundDispatcher = testDispatcher,
                         secureSettingsRepository = FakeSecureSettingsRepository(),
                     ),
-                snapshotRestorer = {
-                    NotificationsSnapshotRestorer(
-                            interactor = interactor,
-                        )
-                        .apply { runBlocking { setUpSnapshotRestorer(FakeSnapshotStore()) } }
-                },
             )
 
         underTest =
diff --git a/tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt b/tests/robotests/src/com/android/customization/picker/repository/NotificationSettingsRepositoryTest.kt
similarity index 85%
rename from tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
rename to tests/robotests/src/com/android/customization/picker/repository/NotificationSettingsRepositoryTest.kt
index be799db..bb6f292 100644
--- a/tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/repository/NotificationSettingsRepositoryTest.kt
@@ -12,18 +12,17 @@
  * 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.data.repository
+package com.android.customization.picker.repository
 
 import android.provider.Settings
 import androidx.test.filters.SmallTest
-import com.android.customization.picker.notifications.shared.model.NotificationSettingsModel
-import com.android.wallpaper.testing.FakeSecureSettingsRepository
+import com.android.systemui.shared.notifications.data.repository.NotificationSettingsRepository
+import com.android.systemui.shared.notifications.shared.model.NotificationSettingsModel
+import com.android.systemui.shared.settings.data.repository.FakeSecureSettingsRepository
 import com.android.wallpaper.testing.collectLastValue
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
@@ -32,12 +31,11 @@
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
 
-@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(JUnit4::class)
-class NotificationsRepositoryTest {
+class NotificationSettingsRepositoryTest {
 
-    private lateinit var underTest: NotificationsRepository
+    private lateinit var underTest: NotificationSettingsRepository
 
     private lateinit var testScope: TestScope
     private lateinit var secureSettingsRepository: FakeSecureSettingsRepository
@@ -49,7 +47,7 @@
         secureSettingsRepository = FakeSecureSettingsRepository()
 
         underTest =
-            NotificationsRepository(
+            NotificationSettingsRepository(
                 scope = testScope.backgroundScope,
                 backgroundDispatcher = testDispatcher,
                 secureSettingsRepository = secureSettingsRepository,