1 /*
2  * Copyright (C) 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.launcher3.util.rule
18 
19 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
20 import com.android.launcher3.InvariantDeviceProfile
21 import com.android.launcher3.LauncherPrefs
22 import java.io.File
23 import java.nio.file.Paths
24 import kotlin.io.path.pathString
25 import org.junit.rules.TestRule
26 import org.junit.runner.Description
27 import org.junit.runners.model.Statement
28 
29 /**
30  * Removes all launcher's DBs from the device and copies the dbs in
31  * assets/databases/BackupAndRestore to the device. It also set's the needed LauncherPrefs variables
32  * needed to kickstart a backup and restore.
33  */
34 class BackAndRestoreRule : TestRule {
35 
36     private val phoneContext = getInstrumentation().targetContext
37 
dbBackUpnull38     private fun dbBackUp() = File(phoneContext.dataDir.path, "/databasesBackUp")
39 
40     private fun dbDirectory() = File(phoneContext.dataDir.path, "/databases")
41 
42     private fun isWorkspaceDatabase(rawFileName: String): Boolean {
43         val fileName = Paths.get(rawFileName).fileName.pathString
44         return fileName.startsWith("launcher") && fileName.endsWith(".db")
45     }
46 
<lambda>null47     fun getDatabaseFiles() = dbDirectory().listFiles().filter { isWorkspaceDatabase(it.name) }
48 
49     /**
50      * Setting RESTORE_DEVICE would trigger a restore next time the Launcher starts, and we remove
51      * the widgets and apps ids to prevent issues when loading the database.
52      */
setRestoreConstantsnull53     private fun setRestoreConstants() {
54         LauncherPrefs.get(phoneContext)
55             .put(LauncherPrefs.RESTORE_DEVICE.to(InvariantDeviceProfile.TYPE_MULTI_DISPLAY))
56         LauncherPrefs.get(phoneContext)
57             .remove(LauncherPrefs.OLD_APP_WIDGET_IDS, LauncherPrefs.APP_WIDGET_IDS)
58     }
59 
uploadDatabasenull60     private fun uploadDatabase(dbName: String) {
61         val file = File(File(getInstrumentation().targetContext.dataDir, "/databases"), dbName)
62         file.writeBytes(
63             getInstrumentation()
64                 .context
65                 .assets
66                 .open("databases/BackupAndRestore/$dbName")
67                 .readBytes()
68         )
69         file.setWritable(true, false)
70     }
71 
uploadDbsnull72     private fun uploadDbs() {
73         uploadDatabase("launcher.db")
74         uploadDatabase("launcher_4_by_4.db")
75         uploadDatabase("launcher_4_by_5.db")
76         uploadDatabase("launcher_3_by_3.db")
77     }
78 
savePreviousStatenull79     private fun savePreviousState() {
80         dbBackUp().deleteRecursively()
81         if (!dbDirectory().renameTo(dbBackUp())) {
82             throw Exception("Unable to move databases to backup directory")
83         }
84         dbDirectory().mkdir()
85         if (!dbDirectory().exists()) {
86             throw Exception("Databases directory doesn't exists")
87         }
88     }
89 
restorePreviousStatenull90     private fun restorePreviousState() {
91         dbDirectory().deleteRecursively()
92         if (!dbBackUp().renameTo(dbDirectory())) {
93             throw Exception("Unable to restore backup directory to databases directory")
94         }
95         dbBackUp().delete()
96     }
97 
beforenull98     fun before() {
99         savePreviousState()
100         setRestoreConstants()
101         uploadDbs()
102     }
103 
afternull104     fun after() {
105         restorePreviousState()
106     }
107 
applynull108     override fun apply(base: Statement?, description: Description?): Statement =
109         object : Statement() {
110             override fun evaluate() {
111                 before()
112                 try {
113                     base?.evaluate()
114                 } finally {
115                     after()
116                 }
117             }
118         }
119 }
120