diff --git a/packages/SettingsLib/DataStore/Android.bp b/packages/SettingsLib/DataStore/Android.bp
index 9fafcab..86c8f0da 100644
--- a/packages/SettingsLib/DataStore/Android.bp
+++ b/packages/SettingsLib/DataStore/Android.bp
@@ -2,12 +2,17 @@
     default_applicable_licenses: ["frameworks_base_license"],
 }
 
+filegroup {
+    name: "SettingsLibDataStore-srcs",
+    srcs: ["src/**/*"],
+}
+
 android_library {
     name: "SettingsLibDataStore",
     defaults: [
         "SettingsLintDefaults",
     ],
-    srcs: ["src/**/*"],
+    srcs: [":SettingsLibDataStore-srcs"],
     static_libs: [
         "androidx.annotation_annotation",
         "androidx.collection_collection-ktx",
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
index 9d3fb66..7644bc9 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreFileArchiver.kt
@@ -19,6 +19,7 @@
 import android.app.backup.BackupDataInputStream
 import android.content.Context
 import android.util.Log
+import androidx.annotation.VisibleForTesting
 import java.io.File
 import java.io.InputStream
 import java.io.OutputStream
@@ -33,11 +34,9 @@
  */
 internal class BackupRestoreFileArchiver(
     private val context: Context,
-    private val fileStorages: List<BackupRestoreFileStorage>,
+    @get:VisibleForTesting internal val fileStorages: List<BackupRestoreFileStorage>,
+    override val name: String,
 ) : BackupRestoreStorage() {
-    override val name: String
-        get() = "file_archiver"
-
     override fun createBackupRestoreEntities(): List<BackupRestoreEntity> =
         fileStorages.map { it.toBackupRestoreEntity() }
 
@@ -88,7 +87,8 @@
     }
 }
 
-private fun BackupRestoreFileStorage.toBackupRestoreEntity() =
+@VisibleForTesting
+internal fun BackupRestoreFileStorage.toBackupRestoreEntity() =
     object : BackupRestoreEntity {
         override val key: String
             get() = storageFilePath
@@ -107,7 +107,7 @@
                 Log.i(LOG_TAG, "[$name] $key not exist")
                 return EntityBackupResult.DELETE
             }
-            val codec = codec() ?: defaultCodec()
+            val codec = defaultCodec()
             // MUST close to flush the data
             wrapBackupOutputStream(codec, outputStream).use { stream ->
                 val bytesCopied = file.inputStream().use { it.copyTo(stream) }
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
index c4c00cb..935f9cc 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorage.kt
@@ -22,6 +22,7 @@
 import android.app.backup.BackupHelper
 import android.os.ParcelFileDescriptor
 import android.util.Log
+import androidx.annotation.VisibleForTesting
 import androidx.collection.MutableScatterMap
 import com.google.common.io.ByteStreams
 import java.io.ByteArrayOutputStream
@@ -60,10 +61,11 @@
      *
      * Map key is the entity key, map value is the checksum of backup data.
      */
-    protected val entityStates = MutableScatterMap<String, Long>()
+    @VisibleForTesting(otherwise = VisibleForTesting.PROTECTED)
+    val entityStates = MutableScatterMap<String, Long>()
 
     /** Entities created by [createBackupRestoreEntities]. This field is for restore only. */
-    private var entities: List<BackupRestoreEntity>? = null
+    @VisibleForTesting internal var entities: List<BackupRestoreEntity>? = null
 
     /** Entities to back up and restore. */
     abstract fun createBackupRestoreEntities(): List<BackupRestoreEntity>
@@ -76,7 +78,7 @@
         data: BackupDataOutput,
         newState: ParcelFileDescriptor,
     ) {
-        oldState.readEntityStates(entityStates)
+        readEntityStates(oldState, entityStates)
         val backupContext = BackupContext(data)
         if (!enableBackup(backupContext)) {
             Log.i(LOG_TAG, "[$name] Backup disabled")
@@ -94,7 +96,10 @@
             val codec = entity.codec() ?: defaultCodec()
             val result =
                 try {
-                    entity.backup(backupContext, wrapBackupOutputStream(codec, checkedOutputStream))
+                    // MUST close to flush all data
+                    wrapBackupOutputStream(codec, checkedOutputStream).use {
+                        entity.backup(backupContext, it)
+                    }
                 } catch (exception: Exception) {
                     Log.e(LOG_TAG, "[$name] Fail to backup entity $key", exception)
                     continue
@@ -191,9 +196,13 @@
     /** Callbacks when restore finished. */
     open fun onRestoreFinished() {}
 
-    private fun ParcelFileDescriptor?.readEntityStates(state: MutableScatterMap<String, Long>) {
+    @VisibleForTesting
+    internal fun readEntityStates(
+        parcelFileDescriptor: ParcelFileDescriptor?,
+        state: MutableScatterMap<String, Long>,
+    ) {
         state.clear()
-        if (this == null) return
+        val fileDescriptor = parcelFileDescriptor?.fileDescriptor ?: return
         // do not close the streams
         val fileInputStream = FileInputStream(fileDescriptor)
         val dataInputStream = DataInputStream(fileInputStream)
@@ -233,6 +242,7 @@
                 dataOutputStream.writeUTF(key)
                 dataOutputStream.writeLong(value)
             }
+            dataOutputStream.flush()
         } catch (exception: Exception) {
             Log.e(LOG_TAG, "[$name] Fail to write state file", exception)
         }
@@ -241,7 +251,7 @@
     }
 
     companion object {
-        private const val STATE_VERSION: Byte = 0
+        internal const val STATE_VERSION: Byte = 0
 
         /** Checksum for entity backup data. */
         fun createChecksum(): Checksum = CRC32()
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
index cfdcaff..8242347 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/BackupRestoreStorageManager.kt
@@ -21,23 +21,32 @@
 import android.app.backup.BackupManager
 import android.content.Context
 import android.util.Log
+import androidx.annotation.VisibleForTesting
 import com.google.common.util.concurrent.MoreExecutors
 import java.util.concurrent.ConcurrentHashMap
 
 /** Manager of [BackupRestoreStorage]. */
 class BackupRestoreStorageManager private constructor(private val application: Application) {
-    private val storageWrappers = ConcurrentHashMap<String, StorageWrapper>()
+    @VisibleForTesting internal val storageWrappers = ConcurrentHashMap<String, StorageWrapper>()
 
     private val executor = MoreExecutors.directExecutor()
 
     /**
      * Adds all the registered [BackupRestoreStorage] as the helpers of given [BackupAgentHelper].
      *
-     * All [BackupRestoreFileStorage]s will be wrapped as a single [BackupRestoreFileArchiver].
+     * All [BackupRestoreFileStorage]s will be wrapped as a single [BackupRestoreFileArchiver],
+     * specify [fileArchiverName] to avoid key prefix conflict if needed.
      *
+     * @param backupAgentHelper backup agent helper to add helpers
+     * @param fileArchiverName key prefix of the [BackupRestoreFileArchiver], the value must not be
+     *   changed in future
      * @see BackupAgentHelper.addHelper
      */
-    fun addBackupAgentHelpers(backupAgentHelper: BackupAgentHelper) {
+    @JvmOverloads
+    fun addBackupAgentHelpers(
+        backupAgentHelper: BackupAgentHelper,
+        fileArchiverName: String = "file_archiver",
+    ) {
         val fileStorages = mutableListOf<BackupRestoreFileStorage>()
         for ((keyPrefix, storageWrapper) in storageWrappers) {
             val storage = storageWrapper.storage
@@ -48,7 +57,7 @@
             }
         }
         // Always add file archiver even fileStorages is empty to handle forward compatibility
-        val fileArchiver = BackupRestoreFileArchiver(application, fileStorages)
+        val fileArchiver = BackupRestoreFileArchiver(application, fileStorages, fileArchiverName)
         backupAgentHelper.addHelper(fileArchiver.name, fileArchiver)
     }
 
@@ -106,7 +115,8 @@
     /** Returns storage with given name, exception is raised if not found. */
     fun getOrThrow(name: String): BackupRestoreStorage = storageWrappers[name]!!.storage
 
-    private inner class StorageWrapper(val storage: BackupRestoreStorage) :
+    @VisibleForTesting
+    internal inner class StorageWrapper(val storage: BackupRestoreStorage) :
         Observer, KeyedObserver<Any?> {
         init {
             when (storage) {
@@ -139,7 +149,7 @@
                 LOG_TAG,
                 "Notify BackupManager dataChanged: storage=$name key=$key reason=$reason"
             )
-            BackupManager.dataChanged(application.packageName)
+            BackupManager(application).dataChanged()
         }
 
         fun removeObserver() {
diff --git a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
index 0c1b417..9f9c0d8 100644
--- a/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
+++ b/packages/SettingsLib/DataStore/src/com/android/settingslib/datastore/SharedPreferencesStorage.kt
@@ -20,9 +20,11 @@
 import android.content.SharedPreferences
 import android.os.Build
 import android.util.Log
-import androidx.core.content.ContextCompat
+import androidx.annotation.VisibleForTesting
 import java.io.File
 
+private fun defaultVerbose() = Build.TYPE == "eng"
+
 /**
  * [SharedPreferences] based storage.
  *
@@ -43,24 +45,35 @@
  * @param verbose Verbose logging on key/value pairs during backup/restore. Enable for dev only!
  * @param filter Filter of key/value pairs for backup and restore.
  */
-class SharedPreferencesStorage
+open class SharedPreferencesStorage
 @JvmOverloads
 constructor(
     context: Context,
     override val name: String,
-    mode: Int,
-    private val verbose: Boolean = (Build.TYPE == "eng"),
+    @get:VisibleForTesting internal val sharedPreferences: SharedPreferences,
+    private val codec: BackupCodec? = null,
+    private val verbose: Boolean = defaultVerbose(),
     private val filter: (String, Any?) -> Boolean = { _, _ -> true },
 ) :
     BackupRestoreFileStorage(context, context.getSharedPreferencesFilePath(name)),
     KeyedObservable<String> by KeyedDataObservable() {
 
-    private val sharedPreferences = context.getSharedPreferences(name, mode)
+    @JvmOverloads
+    constructor(
+        context: Context,
+        name: String,
+        mode: Int,
+        codec: BackupCodec? = null,
+        verbose: Boolean = defaultVerbose(),
+        filter: (String, Any?) -> Boolean = { _, _ -> true },
+    ) : this(context, name, context.getSharedPreferences(name, mode), codec, verbose, filter)
 
     /** Name of the intermediate SharedPreferences. */
-    private val intermediateName: String
+    @VisibleForTesting
+    internal val intermediateName: String
         get() = "_br_$name"
 
+    @Suppress("DEPRECATION")
     private val intermediateSharedPreferences: SharedPreferences
         get() {
             // use MODE_MULTI_PROCESS to ensure a reload
@@ -82,12 +95,15 @@
         sharedPreferences.registerOnSharedPreferenceChangeListener(sharedPreferencesListener)
     }
 
+    override fun defaultCodec() = codec ?: super.defaultCodec()
+
     override val backupFile: File
         // use a different file to avoid multi-thread file write
         get() = context.getSharedPreferencesFile(intermediateName)
 
     override fun prepareBackup(file: File) {
-        val editor = intermediateSharedPreferences.merge(sharedPreferences.all, "Backup")
+        val editor =
+            mergeSharedPreferences(intermediateSharedPreferences, sharedPreferences.all, "Backup")
         // commit to ensure data is write to disk synchronously
         if (!editor.commit()) {
             Log.w(LOG_TAG, "[$name] fail to commit")
@@ -104,8 +120,8 @@
         // observers consistently once restore finished.
         sharedPreferences.unregisterOnSharedPreferenceChangeListener(sharedPreferencesListener)
         val restored = intermediateSharedPreferences
-        val editor = sharedPreferences.merge(restored.all, "Restore")
-        editor.apply() // apply to avoid blocking
+        val editor = mergeSharedPreferences(sharedPreferences, restored.all, "Restore")
+        editor.commit() // commit to avoid race condition
         sharedPreferences.registerOnSharedPreferenceChangeListener(sharedPreferencesListener)
         // clear the intermediate SharedPreferences
         restored.delete(intermediateName)
@@ -115,7 +131,7 @@
         if (deleteSharedPreferences(name)) {
             Log.i(LOG_TAG, "SharedPreferences $name deleted")
         } else {
-            edit().clear().apply()
+            edit().clear().commit() // commit to avoid potential race condition
         }
     }
 
@@ -126,11 +142,13 @@
             false
         }
 
-    private fun SharedPreferences.merge(
+    @VisibleForTesting
+    internal open fun mergeSharedPreferences(
+        sharedPreferences: SharedPreferences,
         entries: Map<String, Any?>,
-        operation: String
+        operation: String,
     ): SharedPreferences.Editor {
-        val editor = edit()
+        val editor = sharedPreferences.edit()
         for ((key, value) in entries) {
             if (!filter.invoke(key, value)) {
                 if (verbose) Log.v(LOG_TAG, "[$name] $operation skips $key=$value")
@@ -184,7 +202,7 @@
     companion object {
         private fun Context.getSharedPreferencesFilePath(name: String): String {
             val file = getSharedPreferencesFile(name)
-            return file.relativeTo(ContextCompat.getDataDir(this)!!).toString()
+            return file.relativeTo(dataDirCompat).toString()
         }
 
         /** Returns the absolute path of shared preferences file. */
diff --git a/packages/SettingsLib/DataStore/tests/Android.bp b/packages/SettingsLib/DataStore/tests/Android.bp
index 8770dfa..5d000eb 100644
--- a/packages/SettingsLib/DataStore/tests/Android.bp
+++ b/packages/SettingsLib/DataStore/tests/Android.bp
@@ -9,11 +9,16 @@
 
 android_robolectric_test {
     name: "SettingsLibDataStoreTest",
-    srcs: ["src/**/*"],
+    srcs: [
+        ":SettingsLibDataStore-srcs", // b/240432457
+        "src/**/*",
+    ],
     static_libs: [
-        "SettingsLibDataStore",
+        "androidx.collection_collection-ktx",
+        "androidx.core_core-ktx",
         "androidx.test.ext.junit",
         "guava",
+        "kotlin-test",
         "mockito-robolectric-prebuilt", // mockito deps order matters!
         "mockito-kotlin2",
     ],
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupCodecTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupCodecTest.kt
new file mode 100644
index 0000000..867831b
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupCodecTest.kt
@@ -0,0 +1,75 @@
+/*
+ * 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.settingslib.datastore
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import com.google.common.truth.Truth.assertWithMessage
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
+import kotlin.random.Random
+import kotlin.test.assertFailsWith
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Tests of [BackupCodec]. */
+@RunWith(AndroidJUnit4::class)
+class BackupCodecTest {
+    @Test
+    fun name() {
+        val names = mutableSetOf<String>()
+        for (codec in allCodecs()) {
+            assertThat(names).doesNotContain(codec.name)
+            names.add(codec.name)
+        }
+    }
+
+    @Test
+    fun fromId() {
+        for (codec in allCodecs()) {
+            assertThat(BackupCodec.fromId(codec.id)).isInstanceOf(codec::class.java)
+        }
+    }
+
+    @Test
+    fun fromId_unknownId() {
+        assertFailsWith(IllegalArgumentException::class) { BackupCodec.fromId(-1) }
+    }
+
+    @Test
+    fun encode_decode() {
+        val random = Random.Default
+        fun test(codec: BackupCodec, size: Int) {
+            val data = random.nextBytes(size)
+
+            // encode
+            val outputStream = ByteArrayOutputStream()
+            codec.encode(outputStream).use { it.write(data) }
+
+            // decode
+            val inputStream = ByteArrayInputStream(outputStream.toByteArray())
+            val result = codec.decode(inputStream).use { it.readBytes() }
+
+            assertWithMessage("$size bytes: $data").that(result).isEqualTo(data)
+        }
+
+        for (codec in allCodecs()) {
+            test(codec, 0)
+            repeat(10) { test(codec, random.nextInt(1, 1024)) }
+        }
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreContextTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreContextTest.kt
new file mode 100644
index 0000000..911665a
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreContextTest.kt
@@ -0,0 +1,44 @@
+/*
+ * 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.settingslib.datastore
+
+import android.app.backup.BackupDataOutput
+import android.os.Build
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+
+/** Tests of [BackupContext] and [RestoreContext]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreContextTest {
+    @Test
+    fun backupContext_quota() {
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.O) return
+        val data = mock<BackupDataOutput> { on { quota } doReturn 10L }
+        assertThat(BackupContext(data).quota).isEqualTo(10)
+    }
+
+    @Test
+    fun backupContext_transportFlags() {
+        if (Build.VERSION.SDK_INT < Build.VERSION_CODES.P) return
+        val data = mock<BackupDataOutput> { on { transportFlags } doReturn 5 }
+        assertThat(BackupContext(data).transportFlags).isEqualTo(5)
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileArchiverTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileArchiverTest.kt
new file mode 100644
index 0000000..6cce453
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileArchiverTest.kt
@@ -0,0 +1,275 @@
+/*
+ * 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.settingslib.datastore
+
+import android.app.Application
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import java.io.ByteArrayInputStream
+import java.io.ByteArrayOutputStream
+import java.io.File
+import java.io.IOException
+import java.io.InputStream
+import kotlin.random.Random
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TemporaryFolder
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.verify
+
+/** Tests of [BackupRestoreFileArchiver]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreFileArchiverTest {
+    private val random = Random.Default
+    private val application: Application = getApplicationContext()
+    @get:Rule val temporaryFolder = TemporaryFolder(application.dataDirCompat)
+
+    @Test
+    fun createBackupRestoreEntities() {
+        val fileStorages = mutableListOf<BackupRestoreFileStorage>()
+        for (count in 0 until 3) {
+            val fileArchiver = BackupRestoreFileArchiver(application, fileStorages, "")
+            fileArchiver.createBackupRestoreEntities().apply {
+                assertThat(this).hasSize(fileStorages.size)
+                for (index in 0 until count) {
+                    assertThat(get(index).key).isEqualTo(fileStorages[index].storageFilePath)
+                }
+            }
+            fileStorages.add(FileStorage("storage", "path$count"))
+        }
+    }
+
+    @Test
+    fun wrapBackupOutputStream() {
+        val fileArchiver = BackupRestoreFileArchiver(application, listOf(), "")
+        val outputStream = ByteArrayOutputStream()
+        assertThat(fileArchiver.wrapBackupOutputStream(BackupZipCodec.BEST_SPEED, outputStream))
+            .isSameInstanceAs(outputStream)
+    }
+
+    @Test
+    fun wrapRestoreInputStream() {
+        val fileArchiver = BackupRestoreFileArchiver(application, listOf(), "")
+        val inputStream = ByteArrayInputStream(byteArrayOf())
+        assertThat(fileArchiver.wrapRestoreInputStream(BackupZipCodec.BEST_SPEED, inputStream))
+            .isSameInstanceAs(inputStream)
+    }
+
+    @Test
+    fun restoreEntity_disabled() {
+        val file = temporaryFolder.newFile()
+        val key = file.name
+        val fileStorage = FileStorage("fs", key, restoreEnabled = false)
+
+        BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver").apply {
+            restoreEntity(newBackupDataInputStream(key, byteArrayOf()))
+            assertThat(entityStates.asMap()).isEmpty()
+        }
+    }
+
+    @Test
+    fun restoreEntity_raiseIOException() {
+        val key = "key"
+        val fileStorage = FileStorage("fs", key)
+        BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver").apply {
+            restoreEntity(newBackupDataInputStream(key, byteArrayOf(), IOException()))
+            assertThat(entityStates.asMap()).isEmpty()
+        }
+    }
+
+    @Test
+    fun restoreEntity_onRestoreFinished_raiseException() {
+        val key = "key"
+        val fileStorage = FileStorage("fs", key, restoreException = IllegalStateException())
+        BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver").apply {
+            val data = random.nextBytes(random.nextInt(10))
+            val outputStream = ByteArrayOutputStream()
+            fileStorage.wrapBackupOutputStream(fileStorage.defaultCodec(), outputStream).use {
+                it.write(data)
+            }
+            val payload = outputStream.toByteArray()
+            restoreEntity(newBackupDataInputStream(key, payload))
+            assertThat(entityStates.asMap()).isEmpty()
+        }
+    }
+
+    @Test
+    fun restoreEntity_forwardCompatibility() {
+        val key = "key"
+        val fileStorage = FileStorage("fs", key)
+        for (codec in allCodecs()) {
+            BackupRestoreFileArchiver(application, listOf(), "archiver").apply {
+                val data = random.nextBytes(random.nextInt(MAX_DATA_SIZE))
+                val outputStream = ByteArrayOutputStream()
+                fileStorage.wrapBackupOutputStream(codec, outputStream).use { it.write(data) }
+                val payload = outputStream.toByteArray()
+
+                restoreEntity(newBackupDataInputStream(key, payload))
+
+                assertThat(entityStates.asMap()).apply {
+                    hasSize(1)
+                    containsKey(key)
+                }
+                assertThat(fileStorage.restoreFile.readBytes()).isEqualTo(data)
+            }
+        }
+    }
+
+    @Test
+    fun restoreEntity() {
+        val folder = File(application.dataDirCompat, "backup")
+        val file = File(folder, "file")
+        val key = "${folder.name}${File.separator}${file.name}"
+        fun test(codec: BackupCodec, size: Int) {
+            val fileStorage = FileStorage("fs", key, if (size % 2 == 0) codec else null)
+            val data = random.nextBytes(size)
+            val outputStream = ByteArrayOutputStream()
+            fileStorage.wrapBackupOutputStream(codec, outputStream).use { it.write(data) }
+            val payload = outputStream.toByteArray()
+
+            val fileArchiver =
+                BackupRestoreFileArchiver(application, listOf(fileStorage), "archiver")
+            fileArchiver.restoreEntity(newBackupDataInputStream(key, payload))
+
+            assertThat(fileArchiver.entityStates.asMap()).apply {
+                hasSize(1)
+                containsKey(key)
+            }
+            assertThat(file.readBytes()).isEqualTo(data)
+        }
+
+        for (codec in allCodecs()) {
+            for (size in 0 until 100) test(codec, size)
+            repeat(10) { test(codec, random.nextInt(100, MAX_DATA_SIZE)) }
+        }
+    }
+
+    @Test
+    fun onRestoreFinished() {
+        val fileStorage = mock<BackupRestoreFileStorage>()
+        val fileArchiver = BackupRestoreFileArchiver(application, listOf(fileStorage), "")
+
+        fileArchiver.onRestoreFinished()
+
+        verify(fileStorage).onRestoreFinished()
+    }
+
+    @Test
+    fun toBackupRestoreEntity_backup_disabled() {
+        val context = BackupContext(mock())
+        val fileStorage =
+            mock<BackupRestoreFileStorage> { on { enableBackup(context) } doReturn false }
+
+        assertThat(fileStorage.toBackupRestoreEntity().backup(context, ByteArrayOutputStream()))
+            .isEqualTo(EntityBackupResult.INTACT)
+
+        verify(fileStorage, never()).prepareBackup(any())
+    }
+
+    @Test
+    fun toBackupRestoreEntity_backup_fileNotExist() {
+        val context = BackupContext(mock())
+        val file = File("NotExist")
+        val fileStorage =
+            mock<BackupRestoreFileStorage> {
+                on { enableBackup(context) } doReturn true
+                on { backupFile } doReturn file
+            }
+
+        assertThat(fileStorage.toBackupRestoreEntity().backup(context, ByteArrayOutputStream()))
+            .isEqualTo(EntityBackupResult.DELETE)
+
+        verify(fileStorage).prepareBackup(file)
+        verify(fileStorage, never()).defaultCodec()
+    }
+
+    @Test
+    fun toBackupRestoreEntity_backup() {
+        val context = BackupContext(mock())
+        val file = temporaryFolder.newFile()
+
+        fun test(codec: BackupCodec, size: Int) {
+            val data = random.nextBytes(size)
+            file.outputStream().use { it.write(data) }
+
+            val outputStream = ByteArrayOutputStream()
+            val fileStorage =
+                mock<BackupRestoreFileStorage> {
+                    on { enableBackup(context) } doReturn true
+                    on { backupFile } doReturn file
+                    on { defaultCodec() } doReturn codec
+                    on { wrapBackupOutputStream(any(), any()) }.thenCallRealMethod()
+                    on { wrapRestoreInputStream(any(), any()) }.thenCallRealMethod()
+                    on { prepareBackup(any()) }.thenCallRealMethod()
+                    on { onBackupFinished(any()) }.thenCallRealMethod()
+                }
+
+            assertThat(fileStorage.toBackupRestoreEntity().backup(context, outputStream))
+                .isEqualTo(EntityBackupResult.UPDATE)
+
+            verify(fileStorage).prepareBackup(file)
+            verify(fileStorage).onBackupFinished(file)
+
+            val decodedData =
+                fileStorage
+                    .wrapRestoreInputStream(codec, ByteArrayInputStream(outputStream.toByteArray()))
+                    .readBytes()
+            assertThat(decodedData).isEqualTo(data)
+        }
+
+        for (codec in allCodecs()) {
+            // test small data to ensure correctness
+            for (size in 0 until 100) test(codec, size)
+            repeat(10) { test(codec, random.nextInt(100, MAX_DATA_SIZE)) }
+        }
+    }
+
+    @Test
+    fun toBackupRestoreEntity_restore() {
+        val restoreContext = RestoreContext("storage")
+        val inputStream =
+            object : InputStream() {
+                override fun read() = throw IllegalStateException()
+
+                override fun read(b: ByteArray, off: Int, len: Int) = throw IllegalStateException()
+            }
+        FileStorage("storage", "path").toBackupRestoreEntity().restore(restoreContext, inputStream)
+    }
+
+    private open class FileStorage(
+        override val name: String,
+        filePath: String,
+        private val codec: BackupCodec? = null,
+        private val restoreEnabled: Boolean? = null,
+        private val restoreException: Exception? = null,
+    ) : BackupRestoreFileStorage(getApplicationContext(), filePath) {
+
+        override fun defaultCodec() = codec ?: super.defaultCodec()
+
+        override fun enableRestore() = restoreEnabled ?: super.enableRestore()
+
+        override fun onRestoreFinished(file: File) {
+            super.onRestoreFinished(file)
+            if (restoreException != null) throw restoreException
+        }
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileStorageTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileStorageTest.kt
new file mode 100644
index 0000000..422273d
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreFileStorageTest.kt
@@ -0,0 +1,105 @@
+/*
+ * 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.settingslib.datastore
+
+import android.app.Application
+import android.os.Build
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import java.io.File
+import kotlin.test.assertFailsWith
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/** Tests of [BackupRestoreFileStorage]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreFileStorageTest {
+    private val application: Application = getApplicationContext()
+
+    @Test
+    fun dataDirCompat() {
+        val expected =
+            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.N) {
+                application.dataDir
+            } else {
+                File(application.applicationInfo.dataDir)
+            }
+        assertThat(application.dataDirCompat).isEqualTo(expected)
+    }
+
+    @Test
+    fun backupFile() {
+        assertThat(FileStorage("path").backupFile.toString())
+            .startsWith(application.dataDirCompat.toString())
+    }
+
+    @Test
+    fun restoreFile() {
+        FileStorage("path").apply { assertThat(restoreFile).isEqualTo(backupFile) }
+    }
+
+    @Test
+    fun checkFilePaths() {
+        FileStorage("path").checkFilePaths()
+    }
+
+    @Test
+    fun checkFilePaths_emptyFilePath() {
+        assertFailsWith(IllegalArgumentException::class) { FileStorage("").checkFilePaths() }
+    }
+
+    @Test
+    fun checkFilePaths_absoluteFilePath() {
+        assertFailsWith(IllegalArgumentException::class) {
+            FileStorage("${File.separatorChar}file").checkFilePaths()
+        }
+    }
+
+    @Test
+    fun checkFilePaths_backupFile() {
+        assertFailsWith(IllegalArgumentException::class) {
+            FileStorage("path", fileForBackup = File("path")).checkFilePaths()
+        }
+    }
+
+    @Test
+    fun checkFilePaths_restoreFile() {
+        assertFailsWith(IllegalArgumentException::class) {
+            FileStorage("path", fileForRestore = File("path")).checkFilePaths()
+        }
+    }
+
+    @Test
+    fun createBackupRestoreEntities() {
+        assertThat(FileStorage("path").createBackupRestoreEntities()).isEmpty()
+    }
+
+    private class FileStorage(
+        filePath: String,
+        val fileForBackup: File? = null,
+        val fileForRestore: File? = null,
+    ) : BackupRestoreFileStorage(getApplicationContext(), filePath) {
+        override val name = "storage"
+
+        override val backupFile: File
+            get() = fileForBackup ?: super.backupFile
+
+        override val restoreFile: File
+            get() = fileForRestore ?: super.restoreFile
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt
new file mode 100644
index 0000000..d8f5028
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageManagerTest.kt
@@ -0,0 +1,237 @@
+/*
+ * 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.settingslib.datastore
+
+import android.app.Application
+import android.app.backup.BackupAgentHelper
+import android.app.backup.BackupHelper
+import android.app.backup.BackupManager
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors.directExecutor
+import kotlin.test.assertFailsWith
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.reset
+import org.mockito.kotlin.verify
+import org.robolectric.Shadows
+import org.robolectric.shadows.ShadowBackupManager
+
+/** Tests of [BackupRestoreStorageManager]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreStorageManagerTest {
+    private val application: Application = getApplicationContext()
+    private val manager = BackupRestoreStorageManager.getInstance(application)
+    private val fileStorage = FileStorage("fileStorage")
+    private val keyedStorage = KeyedStorage("keyedStorage")
+
+    private val storage1 = mock<ObservableBackupRestoreStorage> { on { name } doReturn "1" }
+    private val storage2 = mock<ObservableBackupRestoreStorage> { on { name } doReturn "1" }
+
+    @After
+    fun tearDown() {
+        manager.removeAll()
+        ShadowBackupManager.reset()
+    }
+
+    @Test
+    fun getInstance() {
+        assertThat(BackupRestoreStorageManager.getInstance(application)).isSameInstanceAs(manager)
+    }
+
+    @Test
+    fun addBackupAgentHelpers() {
+        val fs = FileStorage("fs")
+        manager.add(keyedStorage, fileStorage, storage1, fs)
+        val backupAgentHelper = DummyBackupAgentHelper()
+        manager.addBackupAgentHelpers(backupAgentHelper)
+        backupAgentHelper.backupHelpers.apply {
+            assertThat(size).isEqualTo(3)
+            assertThat(remove(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+            assertThat(remove(storage1.name)).isSameInstanceAs(storage1)
+            val fileArchiver = entries.first().value as BackupRestoreFileArchiver
+            assertThat(fileArchiver.fileStorages.toSet()).containsExactly(fs, fileStorage)
+        }
+    }
+
+    @Test
+    fun addBackupAgentHelpers_withoutFileStorage() {
+        manager.add(keyedStorage, storage1)
+        val backupAgentHelper = DummyBackupAgentHelper()
+        manager.addBackupAgentHelpers(backupAgentHelper)
+        backupAgentHelper.backupHelpers.apply {
+            assertThat(size).isEqualTo(3)
+            assertThat(remove(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+            assertThat(remove(storage1.name)).isSameInstanceAs(storage1)
+            val fileArchiver = entries.first().value as BackupRestoreFileArchiver
+            assertThat(fileArchiver.fileStorages).isEmpty()
+        }
+    }
+
+    @Test
+    fun add() {
+        manager.add(keyedStorage, fileStorage, storage1)
+        assertThat(manager.storageWrappers).apply {
+            hasSize(3)
+            containsKey(keyedStorage.name)
+            containsKey(fileStorage.name)
+            containsKey(storage1.name)
+        }
+    }
+
+    @Test
+    fun add_identicalName() {
+        manager.add(storage1)
+        assertFailsWith(IllegalStateException::class) { manager.add(storage1) }
+        assertFailsWith(IllegalStateException::class) { manager.add(storage2) }
+    }
+
+    @Test
+    fun add_nonObservable() {
+        assertFailsWith(IllegalArgumentException::class) {
+            manager.add(mock<BackupRestoreStorage>())
+        }
+    }
+
+    @Test
+    fun removeAll() {
+        add()
+        manager.removeAll()
+        assertThat(manager.storageWrappers).isEmpty()
+    }
+
+    @Test
+    fun remove() {
+        manager.add(keyedStorage, fileStorage)
+        assertThat(manager.remove(storage1.name)).isNull()
+        assertThat(manager.remove(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+        assertThat(manager.remove(fileStorage.name)).isSameInstanceAs(fileStorage)
+    }
+
+    @Test
+    fun get() {
+        manager.add(keyedStorage, fileStorage)
+        assertThat(manager.get(storage1.name)).isNull()
+        assertThat(manager.get(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+        assertThat(manager.get(fileStorage.name)).isSameInstanceAs(fileStorage)
+    }
+
+    @Test
+    fun getOrThrow() {
+        manager.add(keyedStorage, fileStorage)
+        assertFailsWith(NullPointerException::class) { manager.getOrThrow(storage1.name) }
+        assertThat(manager.getOrThrow(keyedStorage.name)).isSameInstanceAs(keyedStorage)
+        assertThat(manager.getOrThrow(fileStorage.name)).isSameInstanceAs(fileStorage)
+    }
+
+    @Test
+    fun notifyRestoreFinished() {
+        manager.add(keyedStorage, fileStorage)
+        val keyedObserver = mock<KeyedObserver<String>>()
+        val anyKeyObserver = mock<KeyedObserver<String?>>()
+        val observer = mock<Observer>()
+        val executor = directExecutor()
+        keyedStorage.addObserver("key", keyedObserver, executor)
+        keyedStorage.addObserver(anyKeyObserver, executor)
+        fileStorage.addObserver(observer, executor)
+
+        manager.onRestoreFinished()
+
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.RESTORE)
+        verify(anyKeyObserver).onKeyChanged(null, ChangeReason.RESTORE)
+        verify(observer).onChanged(ChangeReason.RESTORE)
+        if (isRobolectric()) {
+            Shadows.shadowOf(BackupManager(application)).apply {
+                assertThat(isDataChanged).isFalse()
+                assertThat(dataChangedCount).isEqualTo(0)
+            }
+        }
+    }
+
+    @Test
+    fun notifyBackupManager() {
+        manager.add(keyedStorage, fileStorage)
+        val keyedObserver = mock<KeyedObserver<String>>()
+        val anyKeyObserver = mock<KeyedObserver<String?>>()
+        val observer = mock<Observer>()
+        val executor = directExecutor()
+        keyedStorage.addObserver("key", keyedObserver, executor)
+        keyedStorage.addObserver(anyKeyObserver, executor)
+        fileStorage.addObserver(observer, executor)
+
+        val backupManager =
+            if (isRobolectric()) Shadows.shadowOf(BackupManager(application)) else null
+        backupManager?.apply {
+            assertThat(isDataChanged).isFalse()
+            assertThat(dataChangedCount).isEqualTo(0)
+        }
+
+        fileStorage.notifyChange(ChangeReason.UPDATE)
+        verify(observer).onChanged(ChangeReason.UPDATE)
+        verify(keyedObserver, never()).onKeyChanged(any(), any())
+        verify(anyKeyObserver, never()).onKeyChanged(any(), any())
+        reset(observer)
+        backupManager?.apply {
+            assertThat(isDataChanged).isTrue()
+            assertThat(dataChangedCount).isEqualTo(1)
+        }
+
+        keyedStorage.notifyChange("key", ChangeReason.DELETE)
+        verify(observer, never()).onChanged(any())
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.DELETE)
+        verify(anyKeyObserver).onKeyChanged("key", ChangeReason.DELETE)
+        backupManager?.apply {
+            assertThat(isDataChanged).isTrue()
+            assertThat(dataChangedCount).isEqualTo(2)
+        }
+        reset(keyedObserver)
+
+        // backup manager is not notified for restore event
+        fileStorage.notifyChange(ChangeReason.RESTORE)
+        keyedStorage.notifyChange("key", ChangeReason.RESTORE)
+        verify(observer).onChanged(ChangeReason.RESTORE)
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.RESTORE)
+        verify(anyKeyObserver).onKeyChanged("key", ChangeReason.RESTORE)
+        backupManager?.apply {
+            assertThat(isDataChanged).isTrue()
+            assertThat(dataChangedCount).isEqualTo(2)
+        }
+    }
+
+    private class KeyedStorage(override val name: String) :
+        BackupRestoreStorage(), KeyedObservable<String> by KeyedDataObservable() {
+
+        override fun createBackupRestoreEntities(): List<BackupRestoreEntity> = listOf()
+    }
+
+    private class FileStorage(override val name: String) :
+        BackupRestoreFileStorage(getApplicationContext(), "file"), Observable by DataObservable()
+
+    private class DummyBackupAgentHelper : BackupAgentHelper() {
+        val backupHelpers = mutableMapOf<String, BackupHelper>()
+
+        override fun addHelper(keyPrefix: String, helper: BackupHelper) {
+            backupHelpers[keyPrefix] = helper
+        }
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageTest.kt
new file mode 100644
index 0000000..99998ff
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/BackupRestoreStorageTest.kt
@@ -0,0 +1,414 @@
+/*
+ * 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.settingslib.datastore
+
+import android.app.backup.BackupAgentHelper
+import android.app.backup.BackupDataInput
+import android.app.backup.BackupDataInputStream
+import android.app.backup.BackupDataOutput
+import android.os.ParcelFileDescriptor
+import android.os.ParcelFileDescriptor.MODE_APPEND
+import android.os.ParcelFileDescriptor.MODE_READ_ONLY
+import android.os.ParcelFileDescriptor.MODE_WRITE_ONLY
+import androidx.collection.MutableScatterMap
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.google.common.truth.Truth.assertThat
+import java.io.DataOutputStream
+import java.io.File
+import java.io.FileDescriptor
+import java.io.FileOutputStream
+import java.io.InputStream
+import java.io.OutputStream
+import kotlin.random.Random
+import org.junit.Rule
+import org.junit.Test
+import org.junit.rules.TemporaryFolder
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.argumentCaptor
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.doThrow
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
+import org.mockito.kotlin.reset
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.verify
+
+/** Tests of [BackupRestoreStorage]. */
+@RunWith(AndroidJUnit4::class)
+class BackupRestoreStorageTest {
+    @get:Rule val temporaryFolder = TemporaryFolder()
+
+    private val entity1 = Entity("key1", "value1".toByteArray())
+    private val entity1NoOpCodec = Entity("key1", "value1".toByteArray(), BackupNoOpCodec())
+    private val entity2 = Entity("key2", "value2".toByteArray(), BackupZipCodec.BEST_SPEED)
+
+    @Test
+    fun performBackup_disabled() {
+        val storage = spy(TestStorage().apply { enabled = false })
+        val unused = performBackup { data, newState -> storage.performBackup(null, data, newState) }
+        verify(storage, never()).createBackupRestoreEntities()
+        assertThat(storage.entities).isNull()
+        assertThat(storage.entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun performBackup_enabled() {
+        val storage = spy(TestStorage())
+        val unused = performBackup { data, newState -> storage.performBackup(null, data, newState) }
+        verify(storage).createBackupRestoreEntities()
+        assertThat(storage.entities).isNull()
+        assertThat(storage.entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun performBackup_entityBackupWithException() {
+        val entity =
+            mock<BackupRestoreEntity> {
+                on { key } doReturn ""
+                on { backup(any(), any()) } doThrow IllegalStateException()
+            }
+        val storage = TestStorage(entity, entity1)
+
+        val (_, stateFile) =
+            performBackup { data, newState -> storage.performBackup(null, data, newState) }
+
+        assertThat(storage.readEntityStates(stateFile)).apply {
+            hasSize(1)
+            containsKey(entity1.key)
+        }
+    }
+
+    @Test
+    fun performBackup_update_unchanged() {
+        performBackupTest({}) { entityStates, newEntityStates ->
+            assertThat(entityStates).isEqualTo(newEntityStates)
+        }
+    }
+
+    @Test
+    fun performBackup_intact() {
+        performBackupTest({ entity1.backupResult = EntityBackupResult.INTACT }) {
+            entityStates,
+            newEntityStates ->
+            assertThat(entityStates).isEqualTo(newEntityStates)
+        }
+    }
+
+    @Test
+    fun performBackup_delete() {
+        performBackupTest({ entity1.backupResult = EntityBackupResult.DELETE }) { _, newEntityStates
+            ->
+            assertThat(newEntityStates.size).isEqualTo(1)
+            assertThat(newEntityStates).containsKey(entity2.key)
+        }
+    }
+
+    private fun performBackupTest(
+        update: () -> Unit,
+        verification: (Map<String, Long>, Map<String, Long>) -> Unit,
+    ) {
+        val storage = TestStorage(entity1, entity2)
+        val (_, stateFile) =
+            performBackup { data, newState -> storage.performBackup(null, data, newState) }
+
+        val entityStates = storage.readEntityStates(stateFile)
+        assertThat(entityStates).apply {
+            hasSize(2)
+            containsKey(entity1.key)
+            containsKey(entity2.key)
+        }
+
+        update.invoke()
+        val (_, newStateFile) =
+            performBackup { data, newState ->
+                stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+                    storage.performBackup(it, data, newState)
+                }
+            }
+        verification.invoke(entityStates, storage.readEntityStates(newStateFile))
+    }
+
+    @Test
+    fun restoreEntity_disabled() {
+        val storage = spy(TestStorage().apply { enabled = false })
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.restoreEntity(it.toBackupDataInputStream())
+        }
+        verify(storage, never()).createBackupRestoreEntities()
+        assertThat(storage.entities).isNull()
+        assertThat(storage.entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun restoreEntity_entityNotFound() {
+        val storage = TestStorage()
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            val backupDataInputStream = it.toBackupDataInputStream()
+            backupDataInputStream.setKey("")
+            storage.restoreEntity(backupDataInputStream)
+        }
+    }
+
+    @Test
+    fun restoreEntity_exception() {
+        val storage = TestStorage(entity1)
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            val backupDataInputStream = it.toBackupDataInputStream()
+            backupDataInputStream.setKey(entity1.key)
+            storage.restoreEntity(backupDataInputStream)
+        }
+    }
+
+    @Test
+    fun restoreEntity_codecChanged() {
+        assertThat(entity1.codec()).isNotEqualTo(entity1NoOpCodec.codec())
+        backupAndRestore(entity1) { _, data ->
+            TestStorage(entity1NoOpCodec).apply { restoreEntity(data) }
+        }
+        assertThat(entity1.data).isEqualTo(entity1NoOpCodec.restoredData)
+    }
+
+    @Test
+    fun restoreEntity() {
+        val random = Random.Default
+        fun test(codec: BackupCodec, size: Int) {
+            val entity = Entity("key", random.nextBytes(size), codec)
+            backupAndRestore(entity)
+            entity.verifyRestoredData()
+        }
+        for (codec in allCodecs()) {
+            // test small data to ensure correctness
+            for (size in 0 until 100) test(codec, size)
+            repeat(10) { test(codec, random.nextInt(100, MAX_DATA_SIZE)) }
+        }
+    }
+
+    @Test
+    fun readEntityStates_eof_exception() {
+        val storage = TestStorage()
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates.put("", 0) // add an item to verify that exiting elements are clear
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun readEntityStates_other_exception() {
+        val storage = TestStorage()
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates.put("", 0) // add an item to verify that exiting elements are clear
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).apply {
+            close() // cause exception when read state file
+            storage.readEntityStates(this, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun readEntityStates_unknownVersion() {
+        val storage = TestStorage()
+        val stateFile = temporaryFolder.newFile()
+        stateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+            DataOutputStream(FileOutputStream(it.fileDescriptor))
+                .writeByte(BackupRestoreStorage.STATE_VERSION + 1)
+        }
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates.put("", 0) // add an item to verify that exiting elements are clear
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(0)
+    }
+
+    @Test
+    fun writeNewStateDescription() {
+        val storage = spy(TestStorage())
+        // use read only mode to trigger exception when write state file
+        temporaryFolder.newFile().toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.writeNewStateDescription(it)
+        }
+        verify(storage).onRestoreFinished()
+    }
+
+    @Test
+    fun backupAndRestore() {
+        val storage = spy(TestStorage(entity1, entity2))
+        val backupAgentHelper = BackupAgentHelper()
+        backupAgentHelper.addHelper(storage.name, storage)
+
+        // backup
+        val (dataFile, stateFile) =
+            performBackup { data, newState -> backupAgentHelper.onBackup(null, data, newState) }
+        storage.verifyFieldsArePurged()
+
+        // verify state
+        val entityStates = MutableScatterMap<String, Long>()
+        entityStates[""] = 1
+        storage.readEntityStates(null, entityStates)
+        assertThat(entityStates.size).isEqualTo(0)
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.asMap()).apply {
+            hasSize(2)
+            containsKey(entity1.key)
+            containsKey(entity2.key)
+        }
+        reset(storage)
+
+        // restore
+        val newStateFile = temporaryFolder.newFile()
+        dataFile.toParcelFileDescriptor(MODE_READ_ONLY).use { dataPfd ->
+            newStateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+                backupAgentHelper.onRestore(dataPfd.toBackupDataInput(), 0, it)
+            }
+        }
+        verify(storage).onRestoreFinished()
+        storage.verifyFieldsArePurged()
+
+        // ShadowBackupDataOutput does not write data to file, so restore is bypassed
+        if (!isRobolectric()) {
+            entity1.verifyRestoredData()
+            entity2.verifyRestoredData()
+            assertThat(entityStates.asMap()).isEqualTo(storage.readEntityStates(newStateFile))
+        }
+    }
+
+    private fun backupAndRestore(
+        entity: BackupRestoreEntity,
+        restoreEntity: (TestStorage, BackupDataInputStream) -> TestStorage = { storage, data ->
+            storage.restoreEntity(data)
+            storage
+        },
+    ) {
+        val storage = TestStorage(entity)
+        val entityKey = argumentCaptor<String>()
+        val entitySize = argumentCaptor<Int>()
+        val entityData = argumentCaptor<ByteArray>()
+        val data = mock<BackupDataOutput>()
+
+        val stateFile = temporaryFolder.newFile()
+        stateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+            storage.performBackup(null, data, it)
+        }
+        val entityStates = MutableScatterMap<String, Long>()
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use {
+            storage.readEntityStates(it, entityStates)
+        }
+        assertThat(entityStates.size).isEqualTo(1)
+
+        verify(data).writeEntityHeader(entityKey.capture(), entitySize.capture())
+        verify(data).writeEntityData(entityData.capture(), entitySize.capture())
+        assertThat(entityKey.allValues).isEqualTo(listOf(entity.key))
+        assertThat(entityData.allValues).hasSize(1)
+        val payload = entityData.firstValue
+        assertThat(entitySize.allValues).isEqualTo(listOf(payload.size, payload.size))
+
+        val dataFile = temporaryFolder.newFile()
+        dataFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+            FileOutputStream(it.fileDescriptor).write(payload)
+        }
+
+        newBackupDataInputStream(entity.key, payload).apply {
+            restoreEntity.invoke(storage, this).also {
+                assertThat(it.entityStates).isEqualTo(entityStates)
+            }
+        }
+    }
+
+    fun performBackup(backup: (BackupDataOutput, ParcelFileDescriptor) -> Unit): Pair<File, File> {
+        val dataFile = temporaryFolder.newFile()
+        val stateFile = temporaryFolder.newFile()
+        dataFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use { dataPfd ->
+            stateFile.toParcelFileDescriptor(MODE_WRITE_ONLY or MODE_APPEND).use {
+                backup.invoke(dataPfd.toBackupDataOutput(), it)
+            }
+        }
+        return dataFile to stateFile
+    }
+
+    private fun BackupRestoreStorage.verifyFieldsArePurged() {
+        assertThat(entities).isNull()
+        assertThat(entityStates.size).isEqualTo(0)
+        assertThat(entityStates.capacity).isEqualTo(0)
+    }
+
+    private fun BackupRestoreStorage.readEntityStates(stateFile: File): Map<String, Long> {
+        val entityStates = MutableScatterMap<String, Long>()
+        stateFile.toParcelFileDescriptor(MODE_READ_ONLY).use { readEntityStates(it, entityStates) }
+        return entityStates.asMap()
+    }
+
+    private fun File.toParcelFileDescriptor(mode: Int) = ParcelFileDescriptor.open(this, mode)
+
+    private fun ParcelFileDescriptor.toBackupDataOutput() = fileDescriptor.toBackupDataOutput()
+
+    private fun ParcelFileDescriptor.toBackupDataInputStream(): BackupDataInputStream =
+        BackupDataInputStream::class.java.newInstance(toBackupDataInput())
+
+    private fun ParcelFileDescriptor.toBackupDataInput() = fileDescriptor.toBackupDataInput()
+
+    private fun FileDescriptor.toBackupDataOutput(): BackupDataOutput =
+        BackupDataOutput::class.java.newInstance(this)
+
+    private fun FileDescriptor.toBackupDataInput(): BackupDataInput =
+        BackupDataInput::class.java.newInstance(this)
+}
+
+private open class TestStorage(vararg val backupRestoreEntities: BackupRestoreEntity) :
+    ObservableBackupRestoreStorage() {
+    var enabled: Boolean? = null
+
+    override val name
+        get() = "TestBackup"
+
+    override fun createBackupRestoreEntities() = backupRestoreEntities.toList()
+
+    override fun enableBackup(backupContext: BackupContext) =
+        enabled ?: super.enableBackup(backupContext)
+
+    override fun enableRestore() = enabled ?: super.enableRestore()
+}
+
+private class Entity(
+    override val key: String,
+    val data: ByteArray,
+    private val codec: BackupCodec? = null,
+) : BackupRestoreEntity {
+    var restoredData: ByteArray? = null
+    var backupResult = EntityBackupResult.UPDATE
+
+    override fun codec() = codec ?: super.codec()
+
+    override fun backup(
+        backupContext: BackupContext,
+        outputStream: OutputStream,
+    ): EntityBackupResult {
+        outputStream.write(data)
+        return backupResult
+    }
+
+    override fun restore(restoreContext: RestoreContext, inputStream: InputStream) {
+        restoredData = inputStream.readBytes()
+        inputStream.close()
+    }
+
+    fun verifyRestoredData() = assertThat(restoredData).isEqualTo(data)
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt
index b52586c..8638b2f 100644
--- a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/KeyedObserverTest.kt
@@ -16,76 +16,58 @@
 
 package com.android.settingslib.datastore
 
+import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.google.common.truth.Truth.assertThat
-import com.google.common.util.concurrent.MoreExecutors.directExecutor
+import com.google.common.util.concurrent.MoreExecutors
 import java.util.concurrent.Executor
 import java.util.concurrent.atomic.AtomicInteger
 import org.junit.Assert
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.junit.MockitoJUnit
-import org.mockito.junit.MockitoRule
-import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
 import org.mockito.kotlin.reset
 import org.mockito.kotlin.verify
-import org.robolectric.RobolectricTestRunner
 
-@RunWith(RobolectricTestRunner::class)
+@RunWith(AndroidJUnit4::class)
 class KeyedObserverTest {
-    @get:Rule
-    val mockitoRule: MockitoRule = MockitoJUnit.rule()
+    private val observer1 = mock<KeyedObserver<Any?>>()
+    private val observer2 = mock<KeyedObserver<Any?>>()
+    private val keyedObserver1 = mock<KeyedObserver<Any>>()
+    private val keyedObserver2 = mock<KeyedObserver<Any>>()
 
-    @Mock
-    private lateinit var observer1: KeyedObserver<Any?>
+    private val key1 = Object()
+    private val key2 = Object()
 
-    @Mock
-    private lateinit var observer2: KeyedObserver<Any?>
-
-    @Mock
-    private lateinit var keyedObserver1: KeyedObserver<Any>
-
-    @Mock
-    private lateinit var keyedObserver2: KeyedObserver<Any>
-
-    @Mock
-    private lateinit var key1: Any
-
-    @Mock
-    private lateinit var key2: Any
-
-    @Mock
-    private lateinit var executor: Executor
-
+    private val executor1: Executor = MoreExecutors.directExecutor()
+    private val executor2: Executor = MoreExecutors.newDirectExecutorService()
     private val keyedObservable = KeyedDataObservable<Any>()
 
     @Test
     fun addObserver_sameExecutor() {
-        keyedObservable.addObserver(observer1, executor)
-        keyedObservable.addObserver(observer1, executor)
+        keyedObservable.addObserver(observer1, executor1)
+        keyedObservable.addObserver(observer1, executor1)
     }
 
     @Test
     fun addObserver_keyedObserver_sameExecutor() {
-        keyedObservable.addObserver(key1, keyedObserver1, executor)
-        keyedObservable.addObserver(key1, keyedObserver1, executor)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
     }
 
     @Test
     fun addObserver_differentExecutor() {
-        keyedObservable.addObserver(observer1, executor)
+        keyedObservable.addObserver(observer1, executor1)
         Assert.assertThrows(IllegalStateException::class.java) {
-            keyedObservable.addObserver(observer1, directExecutor())
+            keyedObservable.addObserver(observer1, executor2)
         }
     }
 
     @Test
     fun addObserver_keyedObserver_differentExecutor() {
-        keyedObservable.addObserver(key1, keyedObserver1, executor)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
         Assert.assertThrows(IllegalStateException::class.java) {
-            keyedObservable.addObserver(key1, keyedObserver1, directExecutor())
+            keyedObservable.addObserver(key1, keyedObserver1, executor2)
         }
     }
 
@@ -93,7 +75,7 @@
     fun addObserver_weaklyReferenced() {
         val counter = AtomicInteger()
         var observer: KeyedObserver<Any?>? = KeyedObserver { _, _ -> counter.incrementAndGet() }
-        keyedObservable.addObserver(observer!!, directExecutor())
+        keyedObservable.addObserver(observer!!, executor1)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         assertThat(counter.get()).isEqualTo(1)
@@ -111,7 +93,7 @@
     fun addObserver_keyedObserver_weaklyReferenced() {
         val counter = AtomicInteger()
         var keyObserver: KeyedObserver<Any>? = KeyedObserver { _, _ -> counter.incrementAndGet() }
-        keyedObservable.addObserver(key1, keyObserver!!, directExecutor())
+        keyedObservable.addObserver(key1, keyObserver!!, executor1)
 
         keyedObservable.notifyChange(key1, ChangeReason.UPDATE)
         assertThat(counter.get()).isEqualTo(1)
@@ -127,45 +109,43 @@
 
     @Test
     fun addObserver_notifyObservers_removeObserver() {
-        keyedObservable.addObserver(observer1, directExecutor())
-        keyedObservable.addObserver(observer2, executor)
+        keyedObservable.addObserver(observer1, executor1)
+        keyedObservable.addObserver(observer2, executor2)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         verify(observer1).onKeyChanged(null, ChangeReason.UPDATE)
-        verify(observer2, never()).onKeyChanged(any(), any())
-        verify(executor).execute(any())
+        verify(observer2).onKeyChanged(null, ChangeReason.UPDATE)
 
-        reset(observer1, executor)
+        reset(observer1, observer2)
         keyedObservable.removeObserver(observer2)
 
         keyedObservable.notifyChange(ChangeReason.DELETE)
         verify(observer1).onKeyChanged(null, ChangeReason.DELETE)
-        verify(executor, never()).execute(any())
+        verify(observer2, never()).onKeyChanged(null, ChangeReason.DELETE)
     }
 
     @Test
     fun addObserver_keyedObserver_notifyObservers_removeObserver() {
-        keyedObservable.addObserver(key1, keyedObserver1, directExecutor())
-        keyedObservable.addObserver(key2, keyedObserver2, executor)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
+        keyedObservable.addObserver(key2, keyedObserver2, executor2)
 
         keyedObservable.notifyChange(key1, ChangeReason.UPDATE)
         verify(keyedObserver1).onKeyChanged(key1, ChangeReason.UPDATE)
-        verify(keyedObserver2, never()).onKeyChanged(any(), any())
-        verify(executor, never()).execute(any())
+        verify(keyedObserver2, never()).onKeyChanged(key2, ChangeReason.UPDATE)
 
-        reset(keyedObserver1, executor)
-        keyedObservable.removeObserver(key2, keyedObserver2)
+        reset(keyedObserver1, keyedObserver2)
+        keyedObservable.removeObserver(key1, keyedObserver1)
 
         keyedObservable.notifyChange(key1, ChangeReason.DELETE)
-        verify(keyedObserver1).onKeyChanged(key1, ChangeReason.DELETE)
-        verify(executor, never()).execute(any())
+        verify(keyedObserver1, never()).onKeyChanged(key1, ChangeReason.DELETE)
+        verify(keyedObserver2, never()).onKeyChanged(key2, ChangeReason.DELETE)
     }
 
     @Test
     fun notifyChange_addMoreTypeObservers_checkOnKeyChanged() {
-        keyedObservable.addObserver(observer1, directExecutor())
-        keyedObservable.addObserver(key1, keyedObserver1, directExecutor())
-        keyedObservable.addObserver(key2, keyedObserver2, directExecutor())
+        keyedObservable.addObserver(observer1, executor1)
+        keyedObservable.addObserver(key1, keyedObserver1, executor1)
+        keyedObservable.addObserver(key2, keyedObserver2, executor1)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         verify(observer1).onKeyChanged(null, ChangeReason.UPDATE)
@@ -191,10 +171,10 @@
     fun notifyChange_addObserverWithinCallback() {
         // ConcurrentModificationException is raised if it is not implemented correctly
         val observer: KeyedObserver<Any?> = KeyedObserver { _, _ ->
-            keyedObservable.addObserver(observer1, executor)
+            keyedObservable.addObserver(observer1, executor1)
         }
 
-        keyedObservable.addObserver(observer, directExecutor())
+        keyedObservable.addObserver(observer, executor1)
 
         keyedObservable.notifyChange(ChangeReason.UPDATE)
         keyedObservable.removeObserver(observer)
@@ -204,12 +184,12 @@
     fun notifyChange_KeyedObserver_addObserverWithinCallback() {
         // ConcurrentModificationException is raised if it is not implemented correctly
         val keyObserver: KeyedObserver<Any?> = KeyedObserver { _, _ ->
-            keyedObservable.addObserver(key1, keyedObserver1, executor)
+            keyedObservable.addObserver(key1, keyedObserver1, executor1)
         }
 
-        keyedObservable.addObserver(key1, keyObserver, directExecutor())
+        keyedObservable.addObserver(key1, keyObserver, executor1)
 
         keyedObservable.notifyChange(key1, ChangeReason.UPDATE)
         keyedObservable.removeObserver(key1, keyObserver)
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
index f065829..173c2b1 100644
--- a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/ObserverTest.kt
@@ -22,40 +22,33 @@
 import java.util.concurrent.Executor
 import java.util.concurrent.atomic.AtomicInteger
 import org.junit.Assert.assertThrows
-import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.junit.MockitoJUnit
-import org.mockito.junit.MockitoRule
-import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.never
 import org.mockito.kotlin.reset
 import org.mockito.kotlin.verify
 
 @RunWith(AndroidJUnit4::class)
 class ObserverTest {
-    @get:Rule val mockitoRule: MockitoRule = MockitoJUnit.rule()
+    private val observer1 = mock<Observer>()
+    private val observer2 = mock<Observer>()
 
-    @Mock private lateinit var observer1: Observer
-
-    @Mock private lateinit var observer2: Observer
-
-    @Mock private lateinit var executor: Executor
-
+    private val executor1: Executor = MoreExecutors.directExecutor()
+    private val executor2: Executor = MoreExecutors.newDirectExecutorService()
     private val observable = DataObservable()
 
     @Test
     fun addObserver_sameExecutor() {
-        observable.addObserver(observer1, executor)
-        observable.addObserver(observer1, executor)
+        observable.addObserver(observer1, executor1)
+        observable.addObserver(observer1, executor1)
     }
 
     @Test
     fun addObserver_differentExecutor() {
-        observable.addObserver(observer1, executor)
+        observable.addObserver(observer1, executor1)
         assertThrows(IllegalStateException::class.java) {
-            observable.addObserver(observer1, MoreExecutors.directExecutor())
+            observable.addObserver(observer1, executor2)
         }
     }
 
@@ -63,7 +56,7 @@
     fun addObserver_weaklyReferenced() {
         val counter = AtomicInteger()
         var observer: Observer? = Observer { counter.incrementAndGet() }
-        observable.addObserver(observer!!, MoreExecutors.directExecutor())
+        observable.addObserver(observer!!, executor1)
 
         observable.notifyChange(ChangeReason.UPDATE)
         assertThat(counter.get()).isEqualTo(1)
@@ -79,31 +72,27 @@
 
     @Test
     fun addObserver_notifyObservers_removeObserver() {
-        observable.addObserver(observer1, MoreExecutors.directExecutor())
-        observable.addObserver(observer2, executor)
+        observable.addObserver(observer1, executor1)
+        observable.addObserver(observer2, executor2)
 
         observable.notifyChange(ChangeReason.DELETE)
 
         verify(observer1).onChanged(ChangeReason.DELETE)
-        verify(observer2, never()).onChanged(any())
-        verify(executor).execute(any())
+        verify(observer2).onChanged(ChangeReason.DELETE)
 
-        reset(observer1, executor)
+        reset(observer1, observer2)
         observable.removeObserver(observer2)
 
         observable.notifyChange(ChangeReason.UPDATE)
         verify(observer1).onChanged(ChangeReason.UPDATE)
-        verify(executor, never()).execute(any())
+        verify(observer2, never()).onChanged(ChangeReason.UPDATE)
     }
 
     @Test
     fun notifyChange_addObserverWithinCallback() {
         // ConcurrentModificationException is raised if it is not implemented correctly
-        val observer = Observer { observable.addObserver(observer1, executor) }
-        observable.addObserver(
-            observer,
-            MoreExecutors.directExecutor()
-        )
+        val observer = Observer { observable.addObserver(observer1, executor1) }
+        observable.addObserver(observer, executor1)
         observable.notifyChange(ChangeReason.UPDATE)
         observable.removeObserver(observer)
     }
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/SharedPreferencesStorageTest.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/SharedPreferencesStorageTest.kt
new file mode 100644
index 0000000..fec7d75
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/SharedPreferencesStorageTest.kt
@@ -0,0 +1,175 @@
+/*
+ * 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.settingslib.datastore
+
+import android.app.Application
+import android.content.Context
+import android.content.SharedPreferences
+import android.os.Build
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.platform.app.InstrumentationRegistry
+import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors
+import java.io.ByteArrayOutputStream
+import java.io.File
+import java.util.concurrent.Executor
+import kotlin.random.Random
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.verify
+
+/** Tests of [SharedPreferencesStorage]. */
+@RunWith(AndroidJUnit4::class)
+class SharedPreferencesStorageTest {
+    private val random = Random.Default
+    private val application: Application = ApplicationProvider.getApplicationContext()
+    private val map =
+        mapOf(
+            "boolean" to true,
+            "float" to random.nextFloat(),
+            "int" to random.nextInt(),
+            "long" to random.nextLong(),
+            "string" to "string",
+            "set" to setOf("string"),
+        )
+
+    @After
+    fun tearDown() {
+        application.getSharedPreferences(NAME, MODE).edit().clear().applySync()
+    }
+
+    @Test
+    fun constructors() {
+        val storage1 = SharedPreferencesStorage(application, NAME, MODE)
+        val storage2 =
+            SharedPreferencesStorage(
+                application,
+                NAME,
+                application.getSharedPreferences(NAME, MODE),
+            )
+        assertThat(storage1.sharedPreferences).isSameInstanceAs(storage2.sharedPreferences)
+    }
+
+    @Test
+    fun observer() {
+        val observer = mock<KeyedObserver<Any?>>()
+        val keyedObserver = mock<KeyedObserver<Any>>()
+        val storage = SharedPreferencesStorage(application, NAME, MODE)
+        val executor: Executor = MoreExecutors.directExecutor()
+        storage.addObserver(observer, executor)
+        storage.addObserver("key", keyedObserver, executor)
+
+        storage.sharedPreferences.edit().putString("key", "string").applySync()
+        verify(observer).onKeyChanged("key", ChangeReason.UPDATE)
+        verify(keyedObserver).onKeyChanged("key", ChangeReason.UPDATE)
+
+        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
+            storage.sharedPreferences.edit().clear().applySync()
+            verify(observer).onKeyChanged(null, ChangeReason.DELETE)
+            verify(keyedObserver).onKeyChanged("key", ChangeReason.DELETE)
+        }
+    }
+
+    @Test
+    fun prepareBackup_commitFailed() {
+        val editor = mock<SharedPreferences.Editor> { on { commit() } doReturn false }
+        val storage =
+            spy(SharedPreferencesStorage(application, NAME, MODE)) {
+                onGeneric { mergeSharedPreferences(any(), any(), any()) } doReturn editor
+            }
+        storage.prepareBackup(File(""))
+    }
+
+    @Test
+    fun backupAndRestore() {
+        fun test(codec: BackupCodec) {
+            val storage = SharedPreferencesStorage(application, NAME, MODE, codec)
+            storage.mergeSharedPreferences(storage.sharedPreferences, map, "op").commit()
+            assertThat(storage.sharedPreferences.all).isEqualTo(map)
+
+            val outputStream = ByteArrayOutputStream()
+            assertThat(storage.toBackupRestoreEntity().backup(BackupContext(mock()), outputStream))
+                .isEqualTo(EntityBackupResult.UPDATE)
+            val payload = outputStream.toByteArray()
+
+            storage.sharedPreferences.edit().clear().commit()
+            assertThat(storage.sharedPreferences.all).isEmpty()
+
+            BackupRestoreFileArchiver(application, listOf(storage), "archiver")
+                .restoreEntity(newBackupDataInputStream(storage.storageFilePath, payload))
+            assertThat(storage.sharedPreferences.all).isEqualTo(map)
+        }
+
+        for (codec in allCodecs()) test(codec)
+    }
+
+    @Test
+    fun mergeSharedPreferences_filter() {
+        val storage =
+            SharedPreferencesStorage(application, NAME, MODE) { key, value ->
+                key == "float" || value is String
+            }
+        storage.mergeSharedPreferences(storage.sharedPreferences, map, "op").apply()
+        assertThat(storage.sharedPreferences.all)
+            .containsExactly("float", map["float"], "string", map["string"])
+    }
+
+    @Test
+    fun mergeSharedPreferences_invalidSet() {
+        val storage = SharedPreferencesStorage(application, NAME, MODE, verbose = true)
+        storage
+            .mergeSharedPreferences(
+                storage.sharedPreferences,
+                mapOf<String, Any>("set" to setOf(Any())),
+                "op"
+            )
+            .apply()
+        assertThat(storage.sharedPreferences.all).isEmpty()
+    }
+
+    @Test
+    fun mergeSharedPreferences_unknownType() {
+        val storage = SharedPreferencesStorage(application, NAME, MODE)
+        storage
+            .mergeSharedPreferences(storage.sharedPreferences, map + ("key" to Any()), "op")
+            .apply()
+        assertThat(storage.sharedPreferences.all).isEqualTo(map)
+    }
+
+    @Test
+    fun mergeSharedPreferences() {
+        val storage = SharedPreferencesStorage(application, NAME, MODE, verbose = true)
+        storage.mergeSharedPreferences(storage.sharedPreferences, map, "op").apply()
+        assertThat(storage.sharedPreferences.all).isEqualTo(map)
+    }
+
+    private fun SharedPreferences.Editor.applySync() {
+        apply()
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+    }
+
+    companion object {
+        private const val NAME = "pref"
+        private const val MODE = Context.MODE_PRIVATE
+    }
+}
diff --git a/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/TestUtils.kt b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/TestUtils.kt
new file mode 100644
index 0000000..823d222
--- /dev/null
+++ b/packages/SettingsLib/DataStore/tests/src/com/android/settingslib/datastore/TestUtils.kt
@@ -0,0 +1,83 @@
+/*
+ * 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.settingslib.datastore
+
+import android.app.backup.BackupDataInput
+import android.app.backup.BackupDataInputStream
+import android.os.Build
+import java.io.ByteArrayInputStream
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.mock
+
+internal const val MAX_DATA_SIZE = 1 shl 12
+
+internal fun allCodecs() =
+    arrayOf<BackupCodec>(
+        BackupNoOpCodec(),
+    ) + zipCodecs()
+
+internal fun zipCodecs() =
+    arrayOf<BackupCodec>(
+        BackupZipCodec.DEFAULT_COMPRESSION,
+        BackupZipCodec.BEST_COMPRESSION,
+        BackupZipCodec.BEST_SPEED,
+    )
+
+internal fun <T : Any> Class<T>.newInstance(arg: Any, type: Class<*> = arg.javaClass): T =
+    getDeclaredConstructor(type).apply { isAccessible = true }.newInstance(arg)
+
+internal fun newBackupDataInputStream(
+    key: String,
+    data: ByteArray,
+    e: Exception? = null,
+): BackupDataInputStream {
+    // ShadowBackupDataOutput does not write data to file, so mock for reading data
+    val inputStream = ByteArrayInputStream(data)
+    val backupDataInput =
+        mock<BackupDataInput> {
+            on { readEntityData(any(), any(), any()) } doAnswer
+                {
+                    if (e != null) throw e
+                    val buf = it.arguments[0] as ByteArray
+                    val offset = it.arguments[1] as Int
+                    val size = it.arguments[2] as Int
+                    inputStream.read(buf, offset, size)
+                }
+        }
+    return BackupDataInputStream::class
+        .java
+        .newInstance(backupDataInput, BackupDataInput::class.java)
+        .apply {
+            setKey(key)
+            setDataSize(data.size)
+        }
+}
+
+internal fun BackupDataInputStream.setKey(value: Any) {
+    val field = javaClass.getDeclaredField("key")
+    field.isAccessible = true
+    field.set(this, value)
+}
+
+internal fun BackupDataInputStream.setDataSize(dataSize: Int) {
+    val field = javaClass.getDeclaredField("dataSize")
+    field.isAccessible = true
+    field.setInt(this, dataSize)
+}
+
+internal fun isRobolectric() = Build.FINGERPRINT.contains("robolectric")