Add tests for DE_sys and CE apex data rollback.

This adds tests in StagedRollbackTest for apex data directory rollback.

The test performs the following steps:

- Push files into the apex data directory
- Update the APEX with rollback enabled
- Remove files from data directory
- Add new files to data directory
- Roll back the update
- Verify old files are restored

Bug: 141148175
Test: atest StagedRollbackTest#testRollbackApexDataDirectories_DeSys
Test: atest StagedRollbackTest#testRollbackApexDataDirectories_Ce
Change-Id: Ie1e7ed626c0b654c19b23bfcb008ac4d717aafe2

2 files changed, 134 insertions, 18 deletions

diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
index 80491cd98bc9..0fdcbc52bffa 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -584,6 +584,25 @@ public class StagedRollbackTest {
         assertThat(InstallUtils.getInstalledVersion(TestApp.A)).isEqualTo(1);
     }
 
+    @Test
+    public void testRollbackApexDataDirectories_Phase1() throws Exception {
+        int sessionId = Install.single(TEST_APEX_WITH_APK_V2).setStaged().setEnableRollback()
+                .commit();
+        InstallUtils.waitForSessionReady(sessionId);
+    }
+
+    @Test
+    public void testRollbackApexDataDirectories_Phase2() throws Exception {
+        RollbackInfo available = RollbackUtils.getAvailableRollback(APK_IN_APEX_TESTAPEX_NAME);
+
+        RollbackUtils.rollback(available.getRollbackId(), TEST_APEX_WITH_APK_V2);
+        RollbackInfo committed = RollbackUtils.getCommittedRollbackById(available.getRollbackId());
+
+        // Note: The app is not rolled back until after the rollback is staged
+        // and the device has been rebooted.
+        InstallUtils.waitForSessionReady(committed.getCommittedSessionId());
+    }
+
     private static void runShellCommand(String cmd) {
         ParcelFileDescriptor pfd = InstrumentationRegistry.getInstrumentation().getUiAutomation()
                 .executeShellCommand(cmd);
diff --git a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java b/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
index 672cbb084dae..ce2f7c50e581 100644
--- a/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
+++ b/tests/RollbackTest/StagedRollbackTest/src/com/android/tests/rollback/host/StagedRollbackTest.java
@@ -16,6 +16,8 @@
 
 package com.android.tests.rollback.host;
 
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 import static org.testng.Assert.assertThrows;
@@ -53,6 +55,17 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
 
     private static final String APK_IN_APEX_TESTAPEX_NAME = "com.android.apex.apkrollback.test";
 
+    private static final String TEST_SUBDIR = "/subdir/";
+
+    private static final String TEST_FILENAME_1 = "test_file.txt";
+    private static final String TEST_STRING_1 = "hello this is a test";
+    private static final String TEST_FILENAME_2 = "another_file.txt";
+    private static final String TEST_STRING_2 = "this is a different file";
+    private static final String TEST_FILENAME_3 = "also.xyz";
+    private static final String TEST_STRING_3 = "also\n a\n test\n string";
+    private static final String TEST_FILENAME_4 = "one_more.test";
+    private static final String TEST_STRING_4 = "once more unto the test";
+
     @Before
     public void setUp() throws Exception {
         if (!getDevice().isAdbRoot()) {
@@ -77,6 +90,10 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
         getDevice().executeShellCommand(
                 "rm -f /system/apex/" + APK_IN_APEX_TESTAPEX_NAME + "*.apex "
                         + "/data/apex/active/" + APK_IN_APEX_TESTAPEX_NAME + "*.apex");
+        getDevice().executeShellCommand(
+                "rm -rf " + apexDataDirDeSys(APK_IN_APEX_TESTAPEX_NAME) + "*");
+        getDevice().executeShellCommand(
+                "rm -rf " + apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + "*");
         getDevice().reboot();
     }
 
@@ -245,15 +262,7 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
     @Test
     public void testRollbackApexWithApk() throws Exception {
         getDevice().uninstallPackage("com.android.cts.install.lib.testapp.A");
-        CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(getBuild());
-        final String fileName = APK_IN_APEX_TESTAPEX_NAME + "_v1.apex";
-        final File apex = buildHelper.getTestFile(fileName);
-        if (!getDevice().isAdbRoot()) {
-            getDevice().enableAdbRoot();
-        }
-        getDevice().remountSystemWritable();
-        assertTrue(getDevice().pushFile(apex, "/system/apex/" + fileName));
-        getDevice().reboot();
+        pushTestApex();
         runPhase("testRollbackApexWithApk_Phase1");
         getDevice().reboot();
         runPhase("testRollbackApexWithApk_Phase2");
@@ -267,15 +276,7 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
     @Test
     public void testRollbackApexWithApkCrashing() throws Exception {
         getDevice().uninstallPackage("com.android.cts.install.lib.testapp.A");
-        CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(getBuild());
-        final String fileName = APK_IN_APEX_TESTAPEX_NAME + "_v1.apex";
-        final File apex = buildHelper.getTestFile(fileName);
-        if (!getDevice().isAdbRoot()) {
-            getDevice().enableAdbRoot();
-        }
-        getDevice().remountSystemWritable();
-        assertTrue(getDevice().pushFile(apex, "/system/apex/" + fileName));
-        getDevice().reboot();
+        pushTestApex();
 
         // Install an apex with apk that crashes
         runPhase("testRollbackApexWithApkCrashing_Phase1");
@@ -289,6 +290,102 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
         runPhase("testRollbackApexWithApkCrashing_Phase3");
     }
 
+    /**
+     * Tests that data in DE_sys apex data directory is restored when apex is rolled back.
+     */
+    @Test
+    public void testRollbackApexDataDirectories_DeSys() throws Exception {
+        pushTestApex();
+
+        // Push files to apex data directory
+        String oldFilePath1 = apexDataDirDeSys(APK_IN_APEX_TESTAPEX_NAME) + "/" + TEST_FILENAME_1;
+        String oldFilePath2 =
+                apexDataDirDeSys(APK_IN_APEX_TESTAPEX_NAME) + TEST_SUBDIR + TEST_FILENAME_2;
+        assertTrue(getDevice().pushString(TEST_STRING_1, oldFilePath1));
+        assertTrue(getDevice().pushString(TEST_STRING_2, oldFilePath2));
+
+        // Install new version of the APEX with rollback enabled
+        runPhase("testRollbackApexDataDirectories_Phase1");
+        getDevice().reboot();
+
+        // Replace files in data directory
+        getDevice().deleteFile(oldFilePath1);
+        getDevice().deleteFile(oldFilePath2);
+        String newFilePath3 = apexDataDirDeSys(APK_IN_APEX_TESTAPEX_NAME) + "/" + TEST_FILENAME_3;
+        String newFilePath4 =
+                apexDataDirDeSys(APK_IN_APEX_TESTAPEX_NAME) + TEST_SUBDIR + TEST_FILENAME_4;
+        assertTrue(getDevice().pushString(TEST_STRING_3, newFilePath3));
+        assertTrue(getDevice().pushString(TEST_STRING_4, newFilePath4));
+
+        // Roll back the APEX
+        runPhase("testRollbackApexDataDirectories_Phase2");
+        getDevice().reboot();
+
+        // Verify that old files have been restored and new files are gone
+        assertEquals(TEST_STRING_1, getDevice().pullFileContents(oldFilePath1));
+        assertEquals(TEST_STRING_2, getDevice().pullFileContents(oldFilePath2));
+        assertNull(getDevice().pullFile(newFilePath3));
+        assertNull(getDevice().pullFile(newFilePath4));
+    }
+
+    /**
+     * Tests that data in CE apex data directory is restored when apex is rolled back.
+     */
+    @Test
+    public void testRollbackApexDataDirectories_Ce() throws Exception {
+        pushTestApex();
+
+        // Push files to apex data directory
+        String oldFilePath1 = apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + "/" + TEST_FILENAME_1;
+        String oldFilePath2 =
+                apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + TEST_SUBDIR + TEST_FILENAME_2;
+        assertTrue(getDevice().pushString(TEST_STRING_1, oldFilePath1));
+        assertTrue(getDevice().pushString(TEST_STRING_2, oldFilePath2));
+
+        // Install new version of the APEX with rollback enabled
+        runPhase("testRollbackApexDataDirectories_Phase1");
+        getDevice().reboot();
+
+        // Replace files in data directory
+        getDevice().deleteFile(oldFilePath1);
+        getDevice().deleteFile(oldFilePath2);
+        String newFilePath3 = apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + "/" + TEST_FILENAME_3;
+        String newFilePath4 =
+                apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + TEST_SUBDIR + TEST_FILENAME_4;
+        assertTrue(getDevice().pushString(TEST_STRING_3, newFilePath3));
+        assertTrue(getDevice().pushString(TEST_STRING_4, newFilePath4));
+
+        // Roll back the APEX
+        runPhase("testRollbackApexDataDirectories_Phase2");
+        getDevice().reboot();
+
+        // Verify that old files have been restored and new files are gone
+        assertEquals(TEST_STRING_1, getDevice().pullFileContents(oldFilePath1));
+        assertEquals(TEST_STRING_2, getDevice().pullFileContents(oldFilePath2));
+        assertNull(getDevice().pullFile(newFilePath3));
+        assertNull(getDevice().pullFile(newFilePath4));
+    }
+
+    private void pushTestApex() throws Exception {
+        CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(getBuild());
+        final String fileName = APK_IN_APEX_TESTAPEX_NAME + "_v1.apex";
+        final File apex = buildHelper.getTestFile(fileName);
+        if (!getDevice().isAdbRoot()) {
+            getDevice().enableAdbRoot();
+        }
+        getDevice().remountSystemWritable();
+        assertTrue(getDevice().pushFile(apex, "/system/apex/" + fileName));
+        getDevice().reboot();
+    }
+
+    private static String apexDataDirDeSys(String apexName) {
+        return String.format("/data/misc/apexdata/%s", apexName);
+    }
+
+    private static String apexDataDirCe(String apexName, int userId) {
+        return String.format("/data/misc_ce/%d/apexdata/%s", userId, apexName);
+    }
+
     private void crashProcess(String processName, int numberOfCrashes) throws Exception {
         String pid = "";
         String lastPid = "invalid";