Revert^4 "Removing all restored backups except one so we don't have old backups"

To prevent the error I introduce the method restorePreviousState().

This reverts commit 1234090589be2f8b99284737be23113c19fd0239.

Reason for revert: Fixing presubmit breaks

Bug: 325285743
Flag:  ACONFIG narrow_grid_restore disabled
Test: BackupAndRestoreDBSelectionTest
Change-Id: I63666c29eb91dada4243af019ab0538c40bc0f78
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 42d4d50..2e0f676 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -72,6 +72,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.stream.Collectors;
 
 public class InvariantDeviceProfile {
@@ -578,6 +579,45 @@
     }
 
     /**
+     * Returns the GridOption associated to the given file name or null if the fileName is not
+     * supported.
+     * Ej, launcher.db -> "normal grid", launcher_4_by_4.db -> "practical grid"
+     */
+    public GridOption getGridOptionFromFileName(Context context, String fileName) {
+        return parseAllGridOptions(context).stream()
+                .filter(gridOption -> Objects.equals(gridOption.dbFile, fileName))
+                .findFirst()
+                .orElse(null);
+    }
+
+    /**
+     * Returns the name of the given size on the current device or empty string if the size is not
+     * supported. Ej. 4x4 -> normal, 5x4 -> practical, etc.
+     * (Note: the name of the grid can be different for the same grid size depending of
+     * the values of the InvariantDeviceProfile)
+     *
+     */
+    public String getGridNameFromSize(Context context, Point size) {
+        return parseAllGridOptions(context).stream()
+                .filter(gridOption -> gridOption.numColumns == size.x
+                        && gridOption.numRows == size.y)
+                .map(gridOption -> gridOption.name)
+                .findFirst()
+                .orElse("");
+    }
+
+    /**
+     * Returns the grid option for the given gridName on the current device (Note: the gridOption
+     * be different for the same gridName depending on the values of the InvariantDeviceProfile).
+     */
+    public GridOption getGridOptionFromName(Context context, String gridName) {
+        return parseAllGridOptions(context).stream()
+                .filter(gridOption -> Objects.equals(gridOption.name, gridName))
+                .findFirst()
+                .orElse(null);
+    }
+
+    /**
      * @return all the grid options that can be shown on the device
      */
     public List<GridOption> parseAllGridOptions(Context context) {
diff --git a/src/com/android/launcher3/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 22bc13b..f2b7d18 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -50,8 +50,10 @@
 import androidx.annotation.VisibleForTesting;
 import androidx.annotation.WorkerThread;
 
+import com.android.launcher3.Flags;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherFiles;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
@@ -121,7 +123,48 @@
         // executed again.
         LauncherPrefs.get(context).removeSync(RESTORE_DEVICE);
 
-        idp.reinitializeAfterRestore(context);
+        if (Flags.narrowGridRestore()) {
+            String oldPhoneFileName = idp.dbFile;
+            removeOldDBs(context, oldPhoneFileName);
+            trySettingPreviousGidAsCurrent(context, idp, oldPhoneFileName);
+        } else {
+            idp.reinitializeAfterRestore(context);
+        }
+    }
+
+    /**
+     * Try setting the gird used in the previous phone to the new one. If the current device doesn't
+     * support the previous grid option it will not be set.
+     */
+    private static void trySettingPreviousGidAsCurrent(Context context, InvariantDeviceProfile idp,
+            String oldPhoneDbFileName) {
+        InvariantDeviceProfile.GridOption gridOption = idp.getGridOptionFromFileName(context,
+                oldPhoneDbFileName);
+        if (gridOption != null) {
+            /*
+             * We do this because in some cases different devices have different names for grid
+             * options, in one device the grid option "normal" can be 4x4 while in other it
+             * could be "practical". Calling this changes the current device grid to the same
+             * we had in the other phone, in the case the current phone doesn't support the grid
+             * option we use the default and migrate the db to the default. Migration occurs on
+             * {@code GridSizeMigrationUtil#migrateGridIfNeeded}
+             */
+            idp.setCurrentGrid(context, gridOption.name);
+        }
+    }
+
+    /**
+     * Only keep the last database used on the previous device.
+     */
+    private static void removeOldDBs(Context context, String oldPhoneDbFileName) {
+        // At this point idp.dbFile contains the name of the dbFile from the previous phone
+        LauncherFiles.GRID_DB_FILES.stream()
+                .filter(dbName -> !dbName.equals(oldPhoneDbFileName))
+                .forEach(dbName -> {
+                    if (context.getDatabasePath(dbName).delete()) {
+                        FileLog.d(TAG, "Removed old grid db file: " + dbName);
+                    }
+                });
     }
 
     private static boolean performRestore(Context context, ModelDbController controller) {
diff --git a/tests/assets/databases/BackupAndRestore/launcher.db b/tests/assets/databases/BackupAndRestore/launcher.db
new file mode 100644
index 0000000..126d166
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher.db
Binary files differ
diff --git a/tests/assets/databases/BackupAndRestore/launcher_3_by_3.db b/tests/assets/databases/BackupAndRestore/launcher_3_by_3.db
new file mode 100644
index 0000000..6d8cd73
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher_3_by_3.db
Binary files differ
diff --git a/tests/assets/databases/BackupAndRestore/launcher_4_by_4.db b/tests/assets/databases/BackupAndRestore/launcher_4_by_4.db
new file mode 100644
index 0000000..00061dd
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher_4_by_4.db
Binary files differ
diff --git a/tests/assets/databases/BackupAndRestore/launcher_4_by_5.db b/tests/assets/databases/BackupAndRestore/launcher_4_by_5.db
new file mode 100644
index 0000000..e2e65aa
--- /dev/null
+++ b/tests/assets/databases/BackupAndRestore/launcher_4_by_5.db
Binary files differ
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/rule/BackAndRestoreRule.kt b/tests/multivalentTests/src/com/android/launcher3/util/rule/BackAndRestoreRule.kt
new file mode 100644
index 0000000..da96939
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/rule/BackAndRestoreRule.kt
@@ -0,0 +1,119 @@
+/*
+ * 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.launcher3.util.rule
+
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherPrefs
+import java.io.File
+import java.nio.file.Paths
+import kotlin.io.path.pathString
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+
+/**
+ * Removes all launcher's DBs from the device and copies the dbs in
+ * assets/databases/BackupAndRestore to the device. It also set's the needed LauncherPrefs variables
+ * needed to kickstart a backup and restore.
+ */
+class BackAndRestoreRule : TestRule {
+
+    private val phoneContext = getInstrumentation().targetContext
+
+    private fun dbBackUp() = File(phoneContext.dataDir.path, "/databasesBackUp")
+
+    private fun dbDirectory() = File(phoneContext.dataDir.path, "/databases")
+
+    private fun isWorkspaceDatabase(rawFileName: String): Boolean {
+        val fileName = Paths.get(rawFileName).fileName.pathString
+        return fileName.startsWith("launcher") && fileName.endsWith(".db")
+    }
+
+    fun getDatabaseFiles() = dbDirectory().listFiles().filter { isWorkspaceDatabase(it.name) }
+
+    /**
+     * Setting RESTORE_DEVICE would trigger a restore next time the Launcher starts, and we remove
+     * the widgets and apps ids to prevent issues when loading the database.
+     */
+    private fun setRestoreConstants() {
+        LauncherPrefs.get(phoneContext)
+            .put(LauncherPrefs.RESTORE_DEVICE.to(InvariantDeviceProfile.TYPE_MULTI_DISPLAY))
+        LauncherPrefs.get(phoneContext)
+            .remove(LauncherPrefs.OLD_APP_WIDGET_IDS, LauncherPrefs.APP_WIDGET_IDS)
+    }
+
+    private fun uploadDatabase(dbName: String) {
+        val file = File(File(getInstrumentation().targetContext.dataDir, "/databases"), dbName)
+        file.writeBytes(
+            getInstrumentation()
+                .context
+                .assets
+                .open("databases/BackupAndRestore/$dbName")
+                .readBytes()
+        )
+        file.setWritable(true, false)
+    }
+
+    private fun uploadDbs() {
+        uploadDatabase("launcher.db")
+        uploadDatabase("launcher_4_by_4.db")
+        uploadDatabase("launcher_4_by_5.db")
+        uploadDatabase("launcher_3_by_3.db")
+    }
+
+    private fun savePreviousState() {
+        dbBackUp().deleteRecursively()
+        if (!dbDirectory().renameTo(dbBackUp())) {
+            throw Exception("Unable to move databases to backup directory")
+        }
+        dbDirectory().mkdir()
+        if (!dbDirectory().exists()) {
+            throw Exception("Databases directory doesn't exists")
+        }
+    }
+
+    private fun restorePreviousState() {
+        dbDirectory().deleteRecursively()
+        if (!dbBackUp().renameTo(dbDirectory())) {
+            throw Exception("Unable to restore backup directory to databases directory")
+        }
+        dbBackUp().delete()
+    }
+
+    fun before() {
+        savePreviousState()
+        setRestoreConstants()
+        uploadDbs()
+    }
+
+    fun after() {
+        restorePreviousState()
+    }
+
+    override fun apply(base: Statement?, description: Description?): Statement =
+        object : Statement() {
+            override fun evaluate() {
+                before()
+                try {
+                    base?.evaluate()
+                } finally {
+                    after()
+                }
+            }
+        }
+}
diff --git a/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
new file mode 100644
index 0000000..3e36bbb
--- /dev/null
+++ b/tests/src/com/android/launcher3/backuprestore/BackupAndRestoreDBSelectionTest.kt
@@ -0,0 +1,70 @@
+/*
+ * 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.launcher3.backuprestore
+
+import android.platform.test.flag.junit.SetFlagsRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.MediumTest
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.Flags
+import com.android.launcher3.LauncherPrefs
+import com.android.launcher3.model.ModelDbController
+import com.android.launcher3.util.Executors.MODEL_EXECUTOR
+import com.android.launcher3.util.TestUtil
+import com.android.launcher3.util.rule.BackAndRestoreRule
+import com.android.launcher3.util.rule.setFlags
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Makes sure to test {@code RestoreDbTask#removeOldDBs}, we need to remove all the dbs that are not
+ * the last one used when we restore the device.
+ */
+@RunWith(AndroidJUnit4::class)
+@MediumTest
+class BackupAndRestoreDBSelectionTest {
+
+    @JvmField @Rule var backAndRestoreRule = BackAndRestoreRule()
+
+    @JvmField
+    @Rule
+    val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
+
+    @Before
+    fun setUp() {
+        setFlagsRule.setFlags(true, Flags.FLAG_NARROW_GRID_RESTORE)
+    }
+
+    @Test
+    fun oldDatabasesNotPresentAfterRestore() {
+        val dbController = ModelDbController(getInstrumentation().targetContext)
+        dbController.tryMigrateDB(null)
+        TestUtil.runOnExecutorSync(MODEL_EXECUTOR) {
+            assert(backAndRestoreRule.getDatabaseFiles().size == 1) {
+                "There should only be one database after restoring, the last one used. Actual databases ${backAndRestoreRule.getDatabaseFiles()}"
+            }
+            assert(
+                !LauncherPrefs.get(getInstrumentation().targetContext)
+                    .has(LauncherPrefs.RESTORE_DEVICE)
+            ) {
+                "RESTORE_DEVICE shouldn't be present after a backup and restore."
+            }
+        }
+    }
+}