Merge "Add configurations to support for backup and restore of MediaProvider" into main

12 files changed, 729 insertions, 203 deletions

diff --git a/AndroidManifest.xml b/AndroidManifest.xml
index a18581111..e4249ce1c 100644
--- a/AndroidManifest.xml
+++ b/AndroidManifest.xml
@@ -84,11 +84,17 @@
             android:name="com.android.providers.media.MediaApplication"
             android:label="@string/picker_app_label"
             android:icon="@mipmap/picker_app_icon"
-            android:allowBackup="false"
             android:supportsRtl="true"
             android:forceQueryable="true"
             android:usesCleartextTraffic="true"
-            android:crossProfile="true">
+            android:crossProfile="true"
+            android:allowBackup="true"
+            android:fullBackupOnly="true"
+            android:backupInForeground="true"
+            android:killAfterRestore="false"
+            android:restoreAnyVersion="true"
+            android:backupAgent="com.android.providers.media.backupandrestore.MediaBackupAgent"
+            android:dataExtractionRules="@xml/data_extraction_rules">
         <provider
                 android:name="com.android.providers.media.MediaProvider"
                 android:authorities="media"
diff --git a/res/xml/data_extraction_rules.xml b/res/xml/data_extraction_rules.xml
new file mode 100644
index 000000000..77d698b96
--- /dev/null
+++ b/res/xml/data_extraction_rules.xml
@@ -0,0 +1,35 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+  ~ Copyright (C) 2025 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.
+  -->
+
+<data-extraction-rules>
+    <!-- Disable cloud backup -->
+    <cloud-backup>
+        <exclude domain="root" path="."/>
+        <exclude domain="file" path="."/>
+        <exclude domain="database" path="."/>
+        <exclude domain="sharedpref" path="."/>
+        <exclude domain="external" path="."/>
+        <exclude domain="device_root" path="."/>
+        <exclude domain="device_file" path="."/>
+        <exclude domain="device_database" path="."/>
+        <exclude domain="device_sharedpref" path="."/>
+    </cloud-backup>
+    <!-- Only transfer files/backup folder for device transfer -->
+    <device-transfer>
+        <include domain="file" path="backup" />
+    </device-transfer>
+</data-extraction-rules>
diff --git a/src/com/android/providers/media/MediaProvider.java b/src/com/android/providers/media/MediaProvider.java
index 3985e79fa..0beb07b8d 100644
--- a/src/com/android/providers/media/MediaProvider.java
+++ b/src/com/android/providers/media/MediaProvider.java
@@ -300,6 +300,7 @@ import com.android.modules.utils.BackgroundThread;
 import com.android.modules.utils.build.SdkLevel;
 import com.android.providers.media.DatabaseHelper.OnFilesChangeListener;
 import com.android.providers.media.DatabaseHelper.OnLegacyMigrationListener;
+import com.android.providers.media.backupandrestore.BackupAndRestoreUtils;
 import com.android.providers.media.backupandrestore.BackupExecutor;
 import com.android.providers.media.dao.FileRow;
 import com.android.providers.media.flags.Flags;
@@ -1794,6 +1795,7 @@ public class MediaProvider extends ContentProvider {
                 return null;
             });
         }
+        BackupAndRestoreUtils.doCleanUpAfterRestoreIfRequired(getContext());
 
         // Delete any stale thumbnails
         final int staleThumbnails = mExternalDatabase.runWithTransaction((db) -> {
@@ -7376,6 +7378,13 @@ public class MediaProvider extends ContentProvider {
         return null;
     }
 
+    /**
+     * Triggers backup for MediaProvider.
+     */
+    public void triggerBackup() {
+        mExternalPrimaryBackupExecutor.doBackup(null);
+    }
+
     @Nullable
     private Bundle getResultForWaitForIdle() {
         // TODO(b/195009139): Remove after overriding wait for idle in test to sync picker
diff --git a/src/com/android/providers/media/backupandrestore/BackupAndRestoreUtils.java b/src/com/android/providers/media/backupandrestore/BackupAndRestoreUtils.java
index 4ad086784..c162c8c06 100644
--- a/src/com/android/providers/media/backupandrestore/BackupAndRestoreUtils.java
+++ b/src/com/android/providers/media/backupandrestore/BackupAndRestoreUtils.java
@@ -16,10 +16,21 @@
 
 package com.android.providers.media.backupandrestore;
 
+import static com.android.providers.media.flags.Flags.enableBackupAndRestore;
+
+import android.annotation.NonNull;
+import android.content.Context;
+import android.content.pm.PackageManager;
 import android.provider.MediaStore;
+import android.util.Log;
+
+import com.android.modules.utils.build.SdkLevel;
+import com.android.providers.media.util.FileUtils;
 
 import com.google.common.collect.HashBiMap;
 
+import java.io.File;
+
 /**
  * Class containing common constants and methods for backup and restore.
  */
@@ -56,6 +67,11 @@ public final class BackupAndRestoreUtils {
     static final String RESTORE_COMPLETED = "RESTORE_COMPLETED";
 
     /**
+     * TAG to be used for logging purposes
+     */
+    static final String TAG = BackupAndRestoreUtils.class.getSimpleName();
+
+    /**
      * Array of columns backed up for restore in the future.
      */
     static final String[] BACKUP_COLUMNS = new String[]{
@@ -156,4 +172,117 @@ public final class BackupAndRestoreUtils {
         // Adding number gap to allow addition of new values
         sIdToColumnBiMap.put("80", MediaStore.MediaColumns.XMP);
     }
+
+    /**
+     * Checks whether backup and restore operations are supported and enabled on the current device.
+     *
+     * <p>This method verifies that the required backup and restore flag is enabled, SDK version is
+     * S+ and ensures the device hardware is suitable for these operations. Backup and restore are
+     * supported only on mobile phones and tablets, excluding devices like automotive systems, TVs,
+     * PCs, and smartwatches.</p>
+     *
+     * @param context the application {@link Context}, used to access system resources.
+     * @return {@code true} if backup and restore is enabled and supported on the device,
+     *         {@code false} otherwise.
+     */
+    static boolean isBackupAndRestoreSupported(Context context) {
+        if (!enableBackupAndRestore() || !SdkLevel.isAtLeastS()) {
+            return false;
+        }
+
+        if (context == null || context.getPackageManager() == null) {
+            return false;
+        }
+
+        final PackageManager pm = context.getPackageManager();
+        return !pm.hasSystemFeature(PackageManager.FEATURE_AUTOMOTIVE)
+                && !pm.hasSystemFeature(PackageManager.FEATURE_TELEVISION)
+                && !pm.hasSystemFeature(PackageManager.FEATURE_PC)
+                && !pm.hasSystemFeature(PackageManager.FEATURE_WATCH);
+    }
+
+    /**
+     * Deletes the restore directory and unsets shared preference.
+     *
+     * <p>
+     * This method is triggered during idle maintenance after a media scan.
+     * During the scan, restored values are read and metadata is updated.
+     * Once the scan is complete, the restore directory is no longer needed and is deleted.
+     * The shared preference is unset to indicate that no recent restoration has occurred.
+     * </p>
+     *
+     * @param context The context to check shared preference and delete the restore directory
+     */
+    public static void doCleanUpAfterRestoreIfRequired(Context context) {
+        if (isBackupAndRestoreSupported(context) && isRestoringFromRecentBackup(context)) {
+            try {
+                deleteRestoreDirectory(context);
+            } catch (Exception e) {
+                Log.e(TAG, "Failed to delete restore directory", e);
+            }
+            disableRestoreFromRecentBackup(context);
+        }
+    }
+
+    /**
+     * Indicates that the values should be read from the recent backup. Sets shared preference's
+     * value to true.
+     *
+     * @param context The context used to access shared preferences.
+     */
+    static void enableRestoreFromRecentBackup(@NonNull Context context) {
+        context.getSharedPreferences(SHARED_PREFERENCE_NAME,
+                Context.MODE_PRIVATE).edit().putBoolean(RESTORE_COMPLETED, true).apply();
+    }
+
+    /**
+     * Indicates that values shouldn't be read from backup. Sets shared preference value to false.
+     *
+     * @param context The context used to access shared preferences.
+     */
+    static void disableRestoreFromRecentBackup(@NonNull Context context) {
+        context.getSharedPreferences(SHARED_PREFERENCE_NAME,
+                Context.MODE_PRIVATE).edit().putBoolean(RESTORE_COMPLETED, false).apply();
+    }
+
+    /**
+     * Checks if the shared preference is set, indicating a recent restore operation.
+     *
+     * @param context The application context used to access shared preferences.
+     * @return {@code true} if a restore operation was recently completed, {@code false} otherwise.
+     */
+    static boolean isRestoringFromRecentBackup(@NonNull Context context) {
+        return context.getSharedPreferences(SHARED_PREFERENCE_NAME,
+                Context.MODE_PRIVATE).getBoolean(RESTORE_COMPLETED, false);
+    }
+
+    /**
+     * Deletes the backup directory if it exists.
+     *
+     * @param context The application context used to locate and delete the backup directory.
+     */
+    static void deleteBackupDirectory(@NonNull Context context) {
+        File filesDir = context.getFilesDir();
+        File backupDir = new File(filesDir, BACKUP_DIRECTORY_NAME);
+
+        if (backupDir.exists() && backupDir.isDirectory()) {
+            FileUtils.deleteContents(backupDir);
+            backupDir.delete();
+        }
+    }
+
+    /**
+     * Deletes the restore directory if it exists.
+     *
+     * @param context The application context used to locate and delete the restore directory.
+     */
+    static void deleteRestoreDirectory(@NonNull Context context) {
+        File filesDir = context.getFilesDir();
+        File restoreDir = new File(filesDir, RESTORE_DIRECTORY_NAME);
+
+        if (restoreDir.exists() && restoreDir.isDirectory()) {
+            FileUtils.deleteContents(restoreDir);
+            restoreDir.delete();
+        }
+    }
 }
diff --git a/src/com/android/providers/media/backupandrestore/BackupExecutor.java b/src/com/android/providers/media/backupandrestore/BackupExecutor.java
index 9985477e6..118ddd88f 100644
--- a/src/com/android/providers/media/backupandrestore/BackupExecutor.java
+++ b/src/com/android/providers/media/backupandrestore/BackupExecutor.java
@@ -22,19 +22,17 @@ import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.BACKUP_DIRECTORY_NAME;
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.FIELD_SEPARATOR;
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.KEY_VALUE_SEPARATOR;
-import static com.android.providers.media.flags.Flags.enableBackupAndRestore;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.isBackupAndRestoreSupported;
 import static com.android.providers.media.util.Logging.TAG;
 
 import android.annotation.SuppressLint;
 import android.content.Context;
-import android.content.pm.PackageManager;
 import android.database.Cursor;
 import android.os.CancellationSignal;
 import android.provider.MediaStore.Files.FileColumns;
 import android.provider.MediaStore.MediaColumns;
 import android.util.Log;
 
-import com.android.modules.utils.build.SdkLevel;
 import com.android.providers.media.DatabaseHelper;
 import com.android.providers.media.leveldb.LevelDBEntry;
 import com.android.providers.media.leveldb.LevelDBInstance;
@@ -111,34 +109,6 @@ public final class BackupExecutor {
         updateLastBackedUpGenerationNumber(lastGenerationNumber);
     }
 
-    /**
-     * Checks whether backup and restore operations are supported and enabled on the current device.
-     *
-     * <p>This method verifies that the required backup and restore flag is enabled, SDK version is
-     * S+ and ensures the device hardware is suitable for these operations. Backup and restore are
-     * supported only on mobile phones and tablets, excluding devices like automotive systems, TVs,
-     * PCs, and smartwatches.</p>
-     *
-     * @param context the application {@link Context}, used to access system resources.
-     * @return {@code true} if backup and restore is enabled and supported on the device,
-     *         {@code false} otherwise.
-     */
-    public static boolean isBackupAndRestoreSupported(Context context) {
-        if (!enableBackupAndRestore() || !SdkLevel.isAtLeastS()) {
-            return false;
-        }
-
-        if (context == null || context.getPackageManager() == null) {
-            return false;
-        }
-
-        final PackageManager pm = context.getPackageManager();
-        return !pm.hasSystemFeature(PackageManager.FEATURE_AUTOMOTIVE)
-                && !pm.hasSystemFeature(PackageManager.FEATURE_TELEVISION)
-                && !pm.hasSystemFeature(PackageManager.FEATURE_PC)
-                && !pm.hasSystemFeature(PackageManager.FEATURE_WATCH);
-    }
-
     private long clearBackupIfNeededAndReturnLastBackedUpNumber(long currentDbGenerationNumber,
             long lastBackedUpGenerationNumber) {
         if (currentDbGenerationNumber < lastBackedUpGenerationNumber) {
diff --git a/src/com/android/providers/media/backupandrestore/MediaBackupAgent.java b/src/com/android/providers/media/backupandrestore/MediaBackupAgent.java
new file mode 100644
index 000000000..6fa98220e
--- /dev/null
+++ b/src/com/android/providers/media/backupandrestore/MediaBackupAgent.java
@@ -0,0 +1,164 @@
+/*
+ * Copyright (C) 2025 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.providers.media.backupandrestore;
+
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.BACKUP_DIRECTORY_NAME;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.RESTORE_DIRECTORY_NAME;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.deleteBackupDirectory;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.deleteRestoreDirectory;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.enableRestoreFromRecentBackup;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.isBackupAndRestoreSupported;
+
+import android.annotation.NonNull;
+import android.app.backup.BackupAgent;
+import android.app.backup.BackupDataInput;
+import android.app.backup.BackupDataOutput;
+import android.app.backup.FullBackupDataOutput;
+import android.content.ContentProviderClient;
+import android.content.Context;
+import android.os.ParcelFileDescriptor;
+import android.provider.MediaStore;
+import android.util.Log;
+
+import com.android.providers.media.MediaProvider;
+
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
+import java.util.stream.Stream;
+
+/**
+ * This custom BackupAgent is used for backing up and restoring MediaProvider's metadata.
+ * <p>
+ * It implements {@link BackupAgent#onFullBackup} and {@link BackupAgent#onRestoreFinished}
+ * to handle pre-processing tasks before the backup is initiated and post-processing tasks after
+ * the restore is completed, respectively.
+ * </p>
+ */
+public final class MediaBackupAgent extends BackupAgent {
+
+    private static final String TAG = MediaBackupAgent.class.getSimpleName();
+
+    @Override
+    public void onBackup(ParcelFileDescriptor oldState, BackupDataOutput data,
+            ParcelFileDescriptor newState) throws IOException {
+    }
+
+    @Override
+    public void onRestore(BackupDataInput data, int appVersionCode, ParcelFileDescriptor newState)
+            throws IOException {
+    }
+
+    /**
+     * {@inheritDoc}
+     *
+     * <p>
+     *     Checks if media provider's backup and restore is supported for the device. If supported,
+     *     triggers a backup
+     * </p>
+     */
+    @Override
+    public void onFullBackup(FullBackupDataOutput data) throws IOException {
+        if ((data.getTransportFlags() & FLAG_DEVICE_TO_DEVICE_TRANSFER) == 0) {
+            Log.v(TAG, "Skip cloud backup for media provider");
+            return;
+        }
+
+        Context context = getApplicationContext();
+        if (isBackupAndRestoreSupported(context)) {
+            try (ContentProviderClient cpc = context.getContentResolver()
+                    .acquireContentProviderClient(MediaStore.AUTHORITY)) {
+                final MediaProvider provider = ((MediaProvider) cpc.getLocalContentProvider());
+                provider.triggerBackup();
+            } catch (Exception e) {
+                Log.e(TAG, "Failed to trigger backup", e);
+            }
+        }
+
+        super.onFullBackup(data);
+    }
+
+    /**
+     * This method is called on the target device as a part of restore process after files transfer
+     * is completed and before the first media scan is triggered by restore apk. Checks if media
+     * provider's backup and restore is supported for the device. If supported,
+     *  1. Copies over files from backup directory to restore directory & deletes backup directory
+     *  2. Sets shared preference to true
+     */
+    @Override
+    public void onRestoreFinished() {
+        super.onRestoreFinished();
+
+        Context context = getApplicationContext();
+        if (isBackupAndRestoreSupported(context)) {
+            // The backed-up data from the source device will be read from the restore directory,
+            // while the device will create its own backup directory.
+            copyContentsFromBackupToRestoreDirectory(context);
+
+            // Delete the copied over backup directory
+            deleteBackupDirectory(context);
+
+            // Indicates restore is completed and metadata can be read from restore directory
+            enableRestoreFromRecentBackup(context);
+        }
+    }
+
+    /**
+     * Copies the contents of the backup directory to the restore directory.
+     *
+     * @param context The application context, used to retrieve the files directory.
+     */
+    private static void copyContentsFromBackupToRestoreDirectory(@NonNull Context context) {
+        deleteRestoreDirectory(context);
+
+        File filesDir = context.getFilesDir();
+        File backupDir = new File(filesDir, BACKUP_DIRECTORY_NAME);
+        File restoreDir = new File(filesDir, RESTORE_DIRECTORY_NAME);
+
+        try {
+            if (backupDir.exists() && backupDir.isDirectory()) {
+                copyDirectory(backupDir.toPath(), restoreDir.toPath());
+            } else {
+                Log.e(TAG, "Backup directory does not exist.");
+            }
+        } catch (Exception ex) {
+            Log.e(TAG, "Failed to copy backup directory to restore directory", ex);
+        }
+    }
+
+    private static void copyDirectory(Path source, Path target) throws IOException {
+        try (Stream<Path> paths = Files.walk(source)) {
+            paths.forEach(path -> {
+                try {
+                    Path destination = target.resolve(source.relativize(path));
+                    if (Files.isDirectory(path)) {
+                        if (!Files.exists(destination)) {
+                            Files.createDirectories(destination);
+                        }
+                    } else {
+                        Files.copy(path, destination, StandardCopyOption.REPLACE_EXISTING);
+                    }
+                } catch (IOException e) {
+                    Log.e(TAG, "Failed to copy contents of backup "
+                            + "directory to restore directory", e);
+                }
+            });
+        }
+    }
+}
diff --git a/src/com/android/providers/media/backupandrestore/RestoreExecutor.java b/src/com/android/providers/media/backupandrestore/RestoreExecutor.java
index e6674d6eb..05db447ff 100644
--- a/src/com/android/providers/media/backupandrestore/RestoreExecutor.java
+++ b/src/com/android/providers/media/backupandrestore/RestoreExecutor.java
@@ -20,13 +20,12 @@ import static android.provider.MediaStore.VOLUME_EXTERNAL_PRIMARY;
 
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.FIELD_SEPARATOR;
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.KEY_VALUE_SEPARATOR;
-import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.RESTORE_COMPLETED;
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.RESTORE_DIRECTORY_NAME;
-import static com.android.providers.media.backupandrestore.BackupExecutor.isBackupAndRestoreSupported;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.isBackupAndRestoreSupported;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.isRestoringFromRecentBackup;
 
 import android.content.ContentValues;
 import android.content.Context;
-import android.content.SharedPreferences;
 
 import com.android.providers.media.leveldb.LevelDBInstance;
 import com.android.providers.media.leveldb.LevelDBManager;
@@ -84,14 +83,6 @@ public final class RestoreExecutor {
         return Optional.of(contentValues);
     }
 
-    private static boolean isRestoringFromRecentBackup(Context context) {
-        // Shared preference with key "RESTORE_COMPLETED" should be set to true for recovery to
-        // take place.
-        SharedPreferences sharedPreferences = context.getSharedPreferences(
-                BackupAndRestoreUtils.SHARED_PREFERENCE_NAME, Context.MODE_PRIVATE);
-        return sharedPreferences.getBoolean(RESTORE_COMPLETED, false);
-    }
-
     private Map<String, String> deSerialiseValueString(String valueString) {
         String[] values = valueString.split(FIELD_SEPARATOR);
         Map<String, String> map = new HashMap<>();
diff --git a/src/com/android/providers/media/scan/ModernMediaScanner.java b/src/com/android/providers/media/scan/ModernMediaScanner.java
index d7f146f26..8563cd8a2 100644
--- a/src/com/android/providers/media/scan/ModernMediaScanner.java
+++ b/src/com/android/providers/media/scan/ModernMediaScanner.java
@@ -47,6 +47,7 @@ import static android.media.MediaMetadataRetriever.METADATA_KEY_WRITER;
 import static android.media.MediaMetadataRetriever.METADATA_KEY_YEAR;
 import static android.provider.MediaStore.AUTHORITY;
 import static android.provider.MediaStore.UNKNOWN_STRING;
+import static android.provider.MediaStore.VOLUME_EXTERNAL;
 import static android.text.format.DateUtils.HOUR_IN_MILLIS;
 import static android.text.format.DateUtils.MINUTE_IN_MILLIS;
 
@@ -1212,21 +1213,20 @@ public class ModernMediaScanner implements MediaScanner {
         }
 
         // Recovery is performed on first scan of file in target device
-        if (existingId == -1) {
-            try {
-                if (restoreExecutor != null) {
-                    Optional<ContentValues> restoredDataOptional = restoreExecutor
-                            .getMetadataForFileIfBackedUp(file.getAbsolutePath(), mContext);
-                    if (restoredDataOptional.isPresent()) {
-                        ContentValues valuesRestored = restoredDataOptional.get();
-                        if (isRestoredMetadataOfActualFile(valuesRestored, attrs)) {
-                            return restoreDataFromBackup(valuesRestored, file, attrs, mimeType);
-                        }
+        try {
+            if (restoreExecutor != null) {
+                Optional<ContentValues> restoredDataOptional = restoreExecutor
+                        .getMetadataForFileIfBackedUp(file.getAbsolutePath(), mContext);
+                if (restoredDataOptional.isPresent()) {
+                    ContentValues valuesRestored = restoredDataOptional.get();
+                    if (isRestoredMetadataOfActualFile(valuesRestored, attrs)) {
+                        return restoreDataFromBackup(valuesRestored, file, attrs, mimeType,
+                                existingId);
                     }
                 }
-            } catch (Exception e) {
-                Log.e(TAG, "Error while attempting to restore metadata from backup", e);
             }
+        } catch (Exception e) {
+            Log.e(TAG, "Error while attempting to restore metadata from backup", e);
         }
 
         switch (mediaType) {
@@ -1259,8 +1259,9 @@ public class ModernMediaScanner implements MediaScanner {
     }
 
     private ContentProviderOperation.Builder restoreDataFromBackup(
-            ContentValues restoredValues, File file, BasicFileAttributes attrs, String mimeType) {
-        final ContentProviderOperation.Builder op = newUpsert(MediaStore.VOLUME_EXTERNAL, -1);
+            ContentValues restoredValues, File file, BasicFileAttributes attrs, String mimeType,
+            long existingId) {
+        final ContentProviderOperation.Builder op = newUpsert(VOLUME_EXTERNAL, existingId);
         withGenericValues(op, file, attrs, mimeType, /* mediaType */ null);
         op.withValues(restoredValues);
         return op;
diff --git a/tests/src/com/android/providers/media/backupandrestore/BackupAndRestoreTestUtils.java b/tests/src/com/android/providers/media/backupandrestore/BackupAndRestoreTestUtils.java
new file mode 100644
index 000000000..228be1b40
--- /dev/null
+++ b/tests/src/com/android/providers/media/backupandrestore/BackupAndRestoreTestUtils.java
@@ -0,0 +1,175 @@
+/*
+ * Copyright (C) 2025 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.providers.media.backupandrestore;
+
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.FIELD_SEPARATOR;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.KEY_VALUE_SEPARATOR;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.RESTORE_COMPLETED;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.SHARED_PREFERENCE_NAME;
+
+import android.content.Context;
+import android.provider.MediaStore;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class BackupAndRestoreTestUtils {
+
+    /**
+     * Map used to store column name for given key id.
+     */
+    private static Map<String, String> sColumnIdToKeyMap;
+
+    /**
+     * Map used to store key id for given column name.
+     */
+    private static Map<String, String> sColumnNameToIdMap;
+
+    static void createKeyToColumnNameMap() {
+        sColumnIdToKeyMap = new HashMap<>();
+        sColumnIdToKeyMap.put("0", MediaStore.Files.FileColumns.IS_FAVORITE);
+        sColumnIdToKeyMap.put("1", MediaStore.Files.FileColumns.MEDIA_TYPE);
+        sColumnIdToKeyMap.put("2", MediaStore.Files.FileColumns.MIME_TYPE);
+        sColumnIdToKeyMap.put("3", MediaStore.Files.FileColumns._USER_ID);
+        sColumnIdToKeyMap.put("4", MediaStore.Files.FileColumns.SIZE);
+        sColumnIdToKeyMap.put("5", MediaStore.MediaColumns.DATE_TAKEN);
+        sColumnIdToKeyMap.put("6", MediaStore.MediaColumns.CD_TRACK_NUMBER);
+        sColumnIdToKeyMap.put("7", MediaStore.MediaColumns.ALBUM);
+        sColumnIdToKeyMap.put("8", MediaStore.MediaColumns.ARTIST);
+        sColumnIdToKeyMap.put("9", MediaStore.MediaColumns.AUTHOR);
+        sColumnIdToKeyMap.put("10", MediaStore.MediaColumns.COMPOSER);
+        sColumnIdToKeyMap.put("11", MediaStore.MediaColumns.GENRE);
+        sColumnIdToKeyMap.put("12", MediaStore.MediaColumns.TITLE);
+        sColumnIdToKeyMap.put("13", MediaStore.MediaColumns.YEAR);
+        sColumnIdToKeyMap.put("14", MediaStore.MediaColumns.DURATION);
+        sColumnIdToKeyMap.put("15", MediaStore.MediaColumns.NUM_TRACKS);
+        sColumnIdToKeyMap.put("16", MediaStore.MediaColumns.WRITER);
+        sColumnIdToKeyMap.put("17", MediaStore.MediaColumns.ALBUM_ARTIST);
+        sColumnIdToKeyMap.put("18", MediaStore.MediaColumns.DISC_NUMBER);
+        sColumnIdToKeyMap.put("19", MediaStore.MediaColumns.COMPILATION);
+        sColumnIdToKeyMap.put("20", MediaStore.MediaColumns.BITRATE);
+        sColumnIdToKeyMap.put("21", MediaStore.MediaColumns.CAPTURE_FRAMERATE);
+        sColumnIdToKeyMap.put("22", MediaStore.Audio.AudioColumns.TRACK);
+        sColumnIdToKeyMap.put("23", MediaStore.MediaColumns.DOCUMENT_ID);
+        sColumnIdToKeyMap.put("24", MediaStore.MediaColumns.INSTANCE_ID);
+        sColumnIdToKeyMap.put("25", MediaStore.MediaColumns.ORIGINAL_DOCUMENT_ID);
+        sColumnIdToKeyMap.put("26", MediaStore.MediaColumns.RESOLUTION);
+        sColumnIdToKeyMap.put("27", MediaStore.MediaColumns.ORIENTATION);
+        sColumnIdToKeyMap.put("28", MediaStore.Video.VideoColumns.COLOR_STANDARD);
+        sColumnIdToKeyMap.put("29", MediaStore.Video.VideoColumns.COLOR_TRANSFER);
+        sColumnIdToKeyMap.put("30", MediaStore.Video.VideoColumns.COLOR_RANGE);
+        sColumnIdToKeyMap.put("31", MediaStore.Files.FileColumns._VIDEO_CODEC_TYPE);
+        sColumnIdToKeyMap.put("32", MediaStore.MediaColumns.WIDTH);
+        sColumnIdToKeyMap.put("33", MediaStore.MediaColumns.HEIGHT);
+        sColumnIdToKeyMap.put("34", MediaStore.Images.ImageColumns.DESCRIPTION);
+        sColumnIdToKeyMap.put("35", MediaStore.Images.ImageColumns.EXPOSURE_TIME);
+        sColumnIdToKeyMap.put("36", MediaStore.Images.ImageColumns.F_NUMBER);
+        sColumnIdToKeyMap.put("37", MediaStore.Images.ImageColumns.ISO);
+        sColumnIdToKeyMap.put("38", MediaStore.Images.ImageColumns.SCENE_CAPTURE_TYPE);
+        sColumnIdToKeyMap.put("39", MediaStore.Files.FileColumns._SPECIAL_FORMAT);
+        sColumnIdToKeyMap.put("40", MediaStore.Files.FileColumns.OWNER_PACKAGE_NAME);
+        // Adding number gap to allow addition of new values
+        sColumnIdToKeyMap.put("80", MediaStore.MediaColumns.XMP);
+    }
+
+    static void createColumnNameToKeyMap() {
+        sColumnNameToIdMap = new HashMap<>();
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.IS_FAVORITE, "0");
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.MEDIA_TYPE, "1");
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.MIME_TYPE, "2");
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns._USER_ID, "3");
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.SIZE, "4");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.DATE_TAKEN, "5");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.CD_TRACK_NUMBER, "6");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.ALBUM, "7");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.ARTIST, "8");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.AUTHOR, "9");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.COMPOSER, "10");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.GENRE, "11");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.TITLE, "12");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.YEAR, "13");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.DURATION, "14");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.NUM_TRACKS, "15");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.WRITER, "16");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.ALBUM_ARTIST, "17");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.DISC_NUMBER, "18");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.COMPILATION, "19");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.BITRATE, "20");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.CAPTURE_FRAMERATE, "21");
+        sColumnNameToIdMap.put(MediaStore.Audio.AudioColumns.TRACK, "22");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.DOCUMENT_ID, "23");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.INSTANCE_ID, "24");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.ORIGINAL_DOCUMENT_ID, "25");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.RESOLUTION, "26");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.ORIENTATION, "27");
+        sColumnNameToIdMap.put(MediaStore.Video.VideoColumns.COLOR_STANDARD, "28");
+        sColumnNameToIdMap.put(MediaStore.Video.VideoColumns.COLOR_TRANSFER, "29");
+        sColumnNameToIdMap.put(MediaStore.Video.VideoColumns.COLOR_RANGE, "30");
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns._VIDEO_CODEC_TYPE, "31");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.WIDTH, "32");
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.HEIGHT, "33");
+        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.DESCRIPTION, "34");
+        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.EXPOSURE_TIME, "35");
+        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.F_NUMBER, "36");
+        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.ISO, "37");
+        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.SCENE_CAPTURE_TYPE, "38");
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns._SPECIAL_FORMAT, "39");
+        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.OWNER_PACKAGE_NAME, "40");
+        // Adding number gap to allow addition of new values
+        sColumnNameToIdMap.put(MediaStore.MediaColumns.XMP, "80");
+    }
+
+    static Map<String, String> deSerialiseValueString(String valueString) {
+        if (sColumnIdToKeyMap == null) {
+            createKeyToColumnNameMap();
+        }
+
+        String[] values = valueString.split(":::");
+        Map<String, String> map = new HashMap<>();
+        for (String value : values) {
+            if (value == null || value.isEmpty()) {
+                continue;
+            }
+
+            String[] keyValue = value.split("=", 2);
+            map.put(sColumnIdToKeyMap.get(keyValue[0]), keyValue[1]);
+        }
+
+        return map;
+    }
+
+    static String createSerialisedValue(Map<String, String> entries) {
+        if (sColumnNameToIdMap == null) {
+            createColumnNameToKeyMap();
+        }
+
+        StringBuilder sb = new StringBuilder();
+        for (String backupColumn : sColumnNameToIdMap.keySet()) {
+            if (entries.containsKey(backupColumn)) {
+                sb.append(sColumnNameToIdMap.get(backupColumn)).append(KEY_VALUE_SEPARATOR).append(
+                        entries.get(backupColumn));
+                sb.append(FIELD_SEPARATOR);
+            }
+        }
+        return sb.toString();
+    }
+
+    static boolean getSharedPreferenceValue(Context context) {
+        return context.getSharedPreferences(SHARED_PREFERENCE_NAME,
+                Context.MODE_PRIVATE).getBoolean(RESTORE_COMPLETED, false);
+    }
+}
diff --git a/tests/src/com/android/providers/media/backupandrestore/BackupExecutorTest.java b/tests/src/com/android/providers/media/backupandrestore/BackupExecutorTest.java
index d464dffc3..70c7eb373 100644
--- a/tests/src/com/android/providers/media/backupandrestore/BackupExecutorTest.java
+++ b/tests/src/com/android/providers/media/backupandrestore/BackupExecutorTest.java
@@ -16,6 +16,7 @@
 
 package com.android.providers.media.backupandrestore;
 
+import static com.android.providers.media.backupandrestore.BackupAndRestoreTestUtils.deSerialiseValueString;
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.BACKUP_COLUMNS;
 import static com.android.providers.media.scan.MediaScanner.REASON_UNKNOWN;
 import static com.android.providers.media.scan.MediaScannerTest.stage;
@@ -51,7 +52,6 @@ import com.android.providers.media.util.FileUtils;
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -75,11 +75,6 @@ public final class BackupExecutorTest {
     @Rule
     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
 
-    /**
-     * Map used to store key id for given column and vice versa.
-     */
-    private static Map<String, String> sColumnIdToKeyMap;
-
     private Set<File> mStagedFiles = new HashSet<>();
 
     private Context mIsolatedContext;
@@ -90,11 +85,6 @@ public final class BackupExecutorTest {
 
     private File mDownloadsDir;
 
-    @BeforeClass
-    public static void setupBeforeClass() {
-        createColumnToKeyMap();
-    }
-
     private String mLevelDbPath;
 
     @Before
@@ -216,71 +206,9 @@ public final class BackupExecutorTest {
         }
     }
 
-    static Map<String, String> deSerialiseValueString(String valueString) {
-        String[] values = valueString.split(":::");
-        Map<String, String> map = new HashMap<>();
-        for (String value : values) {
-            if (value == null || value.isEmpty()) {
-                continue;
-            }
-
-            String[] keyValue = value.split("=", 2);
-            map.put(sColumnIdToKeyMap.get(keyValue[0]), keyValue[1]);
-        }
-
-        return map;
-    }
-
     private void stageNewFile(int resId, File file) throws IOException {
         file.createNewFile();
         mStagedFiles.add(file);
         stage(resId, file);
     }
-
-    static void createColumnToKeyMap() {
-        sColumnIdToKeyMap = new HashMap<>();
-        sColumnIdToKeyMap.put("0", MediaStore.Files.FileColumns.IS_FAVORITE);
-        sColumnIdToKeyMap.put("1", MediaStore.Files.FileColumns.MEDIA_TYPE);
-        sColumnIdToKeyMap.put("2", MediaStore.Files.FileColumns.MIME_TYPE);
-        sColumnIdToKeyMap.put("3", MediaStore.Files.FileColumns._USER_ID);
-        sColumnIdToKeyMap.put("4", MediaStore.Files.FileColumns.SIZE);
-        sColumnIdToKeyMap.put("5", MediaStore.MediaColumns.DATE_TAKEN);
-        sColumnIdToKeyMap.put("6", MediaStore.MediaColumns.CD_TRACK_NUMBER);
-        sColumnIdToKeyMap.put("7", MediaStore.MediaColumns.ALBUM);
-        sColumnIdToKeyMap.put("8", MediaStore.MediaColumns.ARTIST);
-        sColumnIdToKeyMap.put("9", MediaStore.MediaColumns.AUTHOR);
-        sColumnIdToKeyMap.put("10", MediaStore.MediaColumns.COMPOSER);
-        sColumnIdToKeyMap.put("11", MediaStore.MediaColumns.GENRE);
-        sColumnIdToKeyMap.put("12", MediaStore.MediaColumns.TITLE);
-        sColumnIdToKeyMap.put("13", MediaStore.MediaColumns.YEAR);
-        sColumnIdToKeyMap.put("14", MediaStore.MediaColumns.DURATION);
-        sColumnIdToKeyMap.put("15", MediaStore.MediaColumns.NUM_TRACKS);
-        sColumnIdToKeyMap.put("16", MediaStore.MediaColumns.WRITER);
-        sColumnIdToKeyMap.put("17", MediaStore.MediaColumns.ALBUM_ARTIST);
-        sColumnIdToKeyMap.put("18", MediaStore.MediaColumns.DISC_NUMBER);
-        sColumnIdToKeyMap.put("19", MediaStore.MediaColumns.COMPILATION);
-        sColumnIdToKeyMap.put("20", MediaStore.MediaColumns.BITRATE);
-        sColumnIdToKeyMap.put("21", MediaStore.MediaColumns.CAPTURE_FRAMERATE);
-        sColumnIdToKeyMap.put("22", MediaStore.Audio.AudioColumns.TRACK);
-        sColumnIdToKeyMap.put("23", MediaStore.MediaColumns.DOCUMENT_ID);
-        sColumnIdToKeyMap.put("24", MediaStore.MediaColumns.INSTANCE_ID);
-        sColumnIdToKeyMap.put("25", MediaStore.MediaColumns.ORIGINAL_DOCUMENT_ID);
-        sColumnIdToKeyMap.put("26", MediaStore.MediaColumns.RESOLUTION);
-        sColumnIdToKeyMap.put("27", MediaStore.MediaColumns.ORIENTATION);
-        sColumnIdToKeyMap.put("28", MediaStore.Video.VideoColumns.COLOR_STANDARD);
-        sColumnIdToKeyMap.put("29", MediaStore.Video.VideoColumns.COLOR_TRANSFER);
-        sColumnIdToKeyMap.put("30", MediaStore.Video.VideoColumns.COLOR_RANGE);
-        sColumnIdToKeyMap.put("31", MediaStore.Files.FileColumns._VIDEO_CODEC_TYPE);
-        sColumnIdToKeyMap.put("32", MediaStore.MediaColumns.WIDTH);
-        sColumnIdToKeyMap.put("33", MediaStore.MediaColumns.HEIGHT);
-        sColumnIdToKeyMap.put("34", MediaStore.Images.ImageColumns.DESCRIPTION);
-        sColumnIdToKeyMap.put("35", MediaStore.Images.ImageColumns.EXPOSURE_TIME);
-        sColumnIdToKeyMap.put("36", MediaStore.Images.ImageColumns.F_NUMBER);
-        sColumnIdToKeyMap.put("37", MediaStore.Images.ImageColumns.ISO);
-        sColumnIdToKeyMap.put("38", MediaStore.Images.ImageColumns.SCENE_CAPTURE_TYPE);
-        sColumnIdToKeyMap.put("39", MediaStore.Files.FileColumns._SPECIAL_FORMAT);
-        sColumnIdToKeyMap.put("40", MediaStore.Files.FileColumns.OWNER_PACKAGE_NAME);
-        // Adding number gap to allow addition of new values
-        sColumnIdToKeyMap.put("80", MediaStore.MediaColumns.XMP);
-    }
 }
diff --git a/tests/src/com/android/providers/media/backupandrestore/MediaBackupAgentTest.java b/tests/src/com/android/providers/media/backupandrestore/MediaBackupAgentTest.java
new file mode 100644
index 000000000..5f892f126
--- /dev/null
+++ b/tests/src/com/android/providers/media/backupandrestore/MediaBackupAgentTest.java
@@ -0,0 +1,189 @@
+/*
+ * Copyright (C) 2025 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.providers.media.backupandrestore;
+
+import static com.android.providers.media.backupandrestore.BackupAndRestoreTestUtils.deSerialiseValueString;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreTestUtils.getSharedPreferenceValue;
+import static com.android.providers.media.scan.MediaScanner.REASON_UNKNOWN;
+import static com.android.providers.media.scan.MediaScannerTest.stage;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
+
+import android.Manifest;
+import android.content.ContentResolver;
+import android.content.ContentValues;
+import android.content.Context;
+import android.database.Cursor;
+import android.net.Uri;
+import android.os.Build;
+import android.os.Environment;
+import android.os.SystemClock;
+import android.platform.test.annotations.EnableFlags;
+import android.platform.test.flag.junit.SetFlagsRule;
+import android.provider.MediaStore;
+
+import androidx.test.InstrumentationRegistry;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SdkSuppress;
+
+import com.android.providers.media.IsolatedContext;
+import com.android.providers.media.R;
+import com.android.providers.media.TestConfigStore;
+import com.android.providers.media.flags.Flags;
+import com.android.providers.media.leveldb.LevelDBInstance;
+import com.android.providers.media.leveldb.LevelDBManager;
+import com.android.providers.media.leveldb.LevelDBResult;
+import com.android.providers.media.scan.ModernMediaScanner;
+import com.android.providers.media.util.FileUtils;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.Map;
+
+@RunWith(AndroidJUnit4.class)
+@EnableFlags(Flags.FLAG_ENABLE_BACKUP_AND_RESTORE)
+@SdkSuppress(minSdkVersion = Build.VERSION_CODES.S)
+public class MediaBackupAgentTest {
+    @Rule
+    public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
+
+    private Context mIsolatedContext;
+
+    private ContentResolver mIsolatedResolver;
+
+    private ModernMediaScanner mModern;
+
+    private File mDownloadsDir;
+
+    private File mRestoreDir;
+    private File mBackupDir;
+
+    private String mLevelDbPath;
+    private MediaBackupAgent mMediaBackupAgent;
+
+    @Before
+    public void setUp() {
+        final Context context = InstrumentationRegistry.getTargetContext();
+        androidx.test.platform.app.InstrumentationRegistry.getInstrumentation().getUiAutomation()
+                .adoptShellPermissionIdentity(Manifest.permission.LOG_COMPAT_CHANGE,
+                        Manifest.permission.READ_COMPAT_CHANGE_CONFIG,
+                        Manifest.permission.DUMP,
+                        Manifest.permission.READ_DEVICE_CONFIG);
+
+        mIsolatedContext = new IsolatedContext(context, "modern", /*asFuseThread*/ false);
+        mIsolatedResolver = mIsolatedContext.getContentResolver();
+        mModern = new ModernMediaScanner(mIsolatedContext, new TestConfigStore());
+        mRestoreDir = new File(mIsolatedContext.getFilesDir(), "restore");
+        mBackupDir = new File(mIsolatedContext.getFilesDir(), "backup");
+        mDownloadsDir = new File(Environment.getExternalStorageDirectory(),
+                Environment.DIRECTORY_DOWNLOADS);
+        mLevelDbPath =
+                mIsolatedContext.getFilesDir().getAbsolutePath() + "/backup/external_primary/";
+        FileUtils.deleteContents(mDownloadsDir);
+
+        mMediaBackupAgent = new MediaBackupAgent();
+        mMediaBackupAgent.attach(mIsolatedContext);
+    }
+
+    @Test
+    public void testCompleteFlow() throws Exception {
+        //create new test file & stage it
+        File file = new File(mDownloadsDir, "testImage_"
+                + SystemClock.elapsedRealtimeNanos() + ".jpg");
+        file.createNewFile();
+        stage(R.raw.test_image, file);
+
+        try {
+            String path = file.getAbsolutePath();
+            // scan directory to have entry in files table
+            mModern.scanDirectory(mDownloadsDir, REASON_UNKNOWN);
+
+            // set is_favorite value to 1. We will check this value later after restoration.
+            updateFavoritesValue(path, 1);
+
+            // run idle maintenance, this will save file's metadata in leveldb with is_favorite = 1
+            MediaStore.runIdleMaintenance(mIsolatedResolver);
+            assertTrue(mBackupDir.exists());
+
+            assertLevelDbExistsAndHasLatestValues(path);
+
+            // run the backup agent. This will copy over backup directory to restore directory and
+            // set shared preference.
+            mMediaBackupAgent.onRestoreFinished();
+            assertTrue(getSharedPreferenceValue(mIsolatedContext));
+            assertTrue(mRestoreDir.exists());
+            assertFalse(mBackupDir.exists());
+
+            //delete existing external db database having old values
+            mIsolatedContext.deleteDatabase("external.db");
+
+            // run media scan, this will populate db and read value from backup
+            mModern.scanDirectory(mDownloadsDir, REASON_UNKNOWN);
+            assertEquals(1, queryFavoritesValue(path));
+
+            // on idle maintenance, clean up is called. It should delete restore directory and set
+            // shared preference to false
+            MediaStore.runIdleMaintenance(mIsolatedResolver);
+            assertFalse(getSharedPreferenceValue(mIsolatedContext));
+            assertFalse(mRestoreDir.exists());
+        } finally {
+            file.delete();
+        }
+    }
+
+    private void assertLevelDbExistsAndHasLatestValues(String path) {
+        // check that entry is created in level db for the file
+        LevelDBInstance levelDBInstance = LevelDBManager.getInstance(mLevelDbPath);
+        assertNotNull(levelDBInstance);
+
+        // check that entry created in level db has latest value(is_favorite = 1)
+        LevelDBResult levelDBResult = levelDBInstance.query(path);
+        assertNotNull(levelDBResult);
+        Map<String, String> actualResultMap = deSerialiseValueString(levelDBResult.getValue());
+        assertEquals(1,
+                Integer.parseInt(actualResultMap.get(MediaStore.MediaColumns.IS_FAVORITE)));
+    }
+
+    private void updateFavoritesValue(String path, int value) {
+        Uri uri = MediaStore.Files.getContentUri(MediaStore.VOLUME_EXTERNAL);
+        String selection = MediaStore.Files.FileColumns.DATA + " LIKE ?";
+        String[] selectionArgs = new String[]{path};
+
+        ContentValues values = new ContentValues();
+        values.put(MediaStore.Files.FileColumns.IS_FAVORITE, value);
+        values.put(MediaStore.MediaColumns.DATE_MODIFIED, System.currentTimeMillis());
+
+        mIsolatedResolver.update(uri, values, selection, selectionArgs);
+    }
+
+    private int queryFavoritesValue(String path) {
+        Uri uri = MediaStore.Files.getContentUri(MediaStore.VOLUME_EXTERNAL);
+        String selection = MediaStore.Files.FileColumns.DATA + " LIKE ?";
+        String[] selectionArgs = new String[]{path};
+
+        Cursor cursor = mIsolatedResolver.query(uri, null, selection, selectionArgs, null);
+        cursor.moveToFirst();
+        return cursor.getInt(cursor.getColumnIndex(MediaStore.Files.FileColumns.IS_FAVORITE));
+    }
+}
diff --git a/tests/src/com/android/providers/media/backupandrestore/RestoreExecutorTest.java b/tests/src/com/android/providers/media/backupandrestore/RestoreExecutorTest.java
index ebffecc13..9f0b41a21 100644
--- a/tests/src/com/android/providers/media/backupandrestore/RestoreExecutorTest.java
+++ b/tests/src/com/android/providers/media/backupandrestore/RestoreExecutorTest.java
@@ -16,8 +16,7 @@
 
 package com.android.providers.media.backupandrestore;
 
-import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.FIELD_SEPARATOR;
-import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.KEY_VALUE_SEPARATOR;
+import static com.android.providers.media.backupandrestore.BackupAndRestoreTestUtils.createSerialisedValue;
 import static com.android.providers.media.backupandrestore.BackupAndRestoreUtils.RESTORE_COMPLETED;
 import static com.android.providers.media.scan.MediaScanner.REASON_UNKNOWN;
 import static com.android.providers.media.scan.MediaScannerTest.stage;
@@ -53,7 +52,6 @@ import com.android.providers.media.util.FileUtils;
 
 import org.junit.After;
 import org.junit.Before;
-import org.junit.BeforeClass;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -73,11 +71,6 @@ public final class RestoreExecutorTest {
     @Rule
     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
 
-    /**
-     * Map used to store key id for given column and vice versa.
-     */
-    private static Map<String, String> sColumnNameToIdMap;
-
     private Context mIsolatedContext;
 
     private ContentResolver mIsolatedResolver;
@@ -86,11 +79,6 @@ public final class RestoreExecutorTest {
 
     private File mDownloadsDir;
 
-    @BeforeClass
-    public static void setupBeforeClass() {
-        createColumnToKeyMap();
-    }
-
     @Before
     public void setUp() {
         final Context context = InstrumentationRegistry.getTargetContext();
@@ -347,63 +335,4 @@ public final class RestoreExecutorTest {
         file.createNewFile();
         stage(resId, file);
     }
-
-    private String createSerialisedValue(Map<String, String> entries) {
-        StringBuilder sb = new StringBuilder();
-        for (String backupColumn : sColumnNameToIdMap.keySet()) {
-            if (entries.containsKey(backupColumn)) {
-                sb.append(sColumnNameToIdMap.get(backupColumn)).append(KEY_VALUE_SEPARATOR).append(
-                        entries.get(backupColumn));
-                sb.append(FIELD_SEPARATOR);
-            }
-        }
-        return sb.toString();
-    }
-
-    private static void createColumnToKeyMap() {
-        sColumnNameToIdMap = new HashMap<>();
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.IS_FAVORITE, "0");
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.MEDIA_TYPE, "1");
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.MIME_TYPE, "2");
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns._USER_ID, "3");
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.SIZE, "4");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.DATE_TAKEN, "5");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.CD_TRACK_NUMBER, "6");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.ALBUM, "7");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.ARTIST, "8");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.AUTHOR, "9");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.COMPOSER, "10");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.GENRE, "11");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.TITLE, "12");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.YEAR, "13");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.DURATION, "14");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.NUM_TRACKS, "15");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.WRITER, "16");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.ALBUM_ARTIST, "17");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.DISC_NUMBER, "18");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.COMPILATION, "19");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.BITRATE, "20");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.CAPTURE_FRAMERATE, "21");
-        sColumnNameToIdMap.put(MediaStore.Audio.AudioColumns.TRACK, "22");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.DOCUMENT_ID, "23");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.INSTANCE_ID, "24");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.ORIGINAL_DOCUMENT_ID, "25");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.RESOLUTION, "26");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.ORIENTATION, "27");
-        sColumnNameToIdMap.put(MediaStore.Video.VideoColumns.COLOR_STANDARD, "28");
-        sColumnNameToIdMap.put(MediaStore.Video.VideoColumns.COLOR_TRANSFER, "29");
-        sColumnNameToIdMap.put(MediaStore.Video.VideoColumns.COLOR_RANGE, "30");
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns._VIDEO_CODEC_TYPE, "31");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.WIDTH, "32");
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.HEIGHT, "33");
-        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.DESCRIPTION, "34");
-        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.EXPOSURE_TIME, "35");
-        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.F_NUMBER, "36");
-        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.ISO, "37");
-        sColumnNameToIdMap.put(MediaStore.Images.ImageColumns.SCENE_CAPTURE_TYPE, "38");
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns._SPECIAL_FORMAT, "39");
-        sColumnNameToIdMap.put(MediaStore.Files.FileColumns.OWNER_PACKAGE_NAME, "40");
-        // Adding number gap to allow addition of new values
-        sColumnNameToIdMap.put(MediaStore.MediaColumns.XMP, "80");
-    }
 }