diff --git a/tests/RollbackTest/Android.bp b/tests/RollbackTest/Android.bp
index 6e1cef4..9f6ce4e 100644
--- a/tests/RollbackTest/Android.bp
+++ b/tests/RollbackTest/Android.bp
@@ -31,7 +31,8 @@
     test_config: "RollbackTest.xml",
     java_resources: [
         ":com.android.apex.apkrollback.test_v2",
-        ":com.android.apex.apkrollback.test_v2Crashing"
+        ":com.android.apex.apkrollback.test_v2Crashing",
+        ":test.rebootless_apex_v2",
     ],
 }
 
@@ -47,7 +48,10 @@
     ],
     test_suites: ["general-tests"],
     test_config: "StagedRollbackTest.xml",
-    data: [":com.android.apex.apkrollback.test_v1"],
+    data: [
+        ":com.android.apex.apkrollback.test_v1",
+        ":test.rebootless_apex_v1",
+    ],
 }
 
 java_test_host {
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 539090b..4512876 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -241,6 +241,33 @@
     }
 
     @Test
+    public void testRollbackRebootlessApex() throws Exception {
+        final String packageName = "test.apex.rebootless";
+        assertThat(InstallUtils.getInstalledVersion(packageName)).isEqualTo(1);
+
+        // install
+        TestApp apex1 = new TestApp("TestRebootlessApexV1", packageName, 1,
+                /* isApex= */ true, "test.rebootless_apex_v1.apex");
+        TestApp apex2 = new TestApp("TestRebootlessApexV2", packageName, 2,
+                /* isApex= */ true, "test.rebootless_apex_v2.apex");
+        Install.single(apex2).setEnableRollback(PackageManager.ROLLBACK_DATA_POLICY_RETAIN)
+                .commit();
+
+        // verify rollback
+        assertThat(InstallUtils.getInstalledVersion(packageName)).isEqualTo(2);
+        RollbackManager rm = RollbackUtils.getRollbackManager();
+        RollbackInfo rollback = getUniqueRollbackInfoForPackage(
+                rm.getAvailableRollbacks(), packageName);
+        assertThat(rollback).isNotNull();
+        assertThat(rollback).packagesContainsExactly(Rollback.from(apex2).to(apex1));
+        assertThat(rollback).isNotStaged();
+
+        // rollback
+        RollbackUtils.rollback(rollback.getRollbackId());
+        assertThat(InstallUtils.getInstalledVersion(packageName)).isEqualTo(1);
+    }
+
+    @Test
     public void hasMainlineModule() throws Exception {
         String pkgName = getModuleMetadataPackageName();
         boolean existed =  InstrumentationRegistry.getInstrumentation().getContext()
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 293f159..3576a78 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
@@ -96,7 +96,9 @@
         deleteFiles("/system/apex/" + APK_IN_APEX_TESTAPEX_NAME + "*.apex",
                 "/data/apex/active/" + APK_IN_APEX_TESTAPEX_NAME + "*.apex",
                 apexDataDirDeSys(APK_IN_APEX_TESTAPEX_NAME) + "*",
-                apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + "*");
+                apexDataDirCe(APK_IN_APEX_TESTAPEX_NAME, 0) + "*",
+                "/system/apex/test.rebootless_apex_v*.apex",
+                "/data/apex/active/test.apex.rebootless*.apex");
     }
 
     /**
@@ -160,7 +162,7 @@
      */
     @Test
     public void testRollbackApexWithApkCrashing() throws Exception {
-        pushTestApex();
+        pushTestApex(APK_IN_APEX_TESTAPEX_NAME + "_v1.apex");
 
         // Install an apex with apk that crashes
         runPhase("testRollbackApexWithApkCrashing_Phase1_Install");
@@ -181,6 +183,15 @@
     }
 
     /**
+     * Tests rollback is supported correctly for rebootless apex
+     */
+    @Test
+    public void testRollbackRebootlessApex() throws Exception {
+        pushTestApex("test.rebootless_apex_v1.apex");
+        runPhase("testRollbackRebootlessApex");
+    }
+
+    /**
      * Tests that packages are monitored across multiple reboots.
      */
     @Test
@@ -204,9 +215,8 @@
         runPhase("testWatchdogMonitorsAcrossReboots_Phase3_VerifyRollback");
     }
 
-    private void pushTestApex() throws Exception {
+    private void pushTestApex(String fileName) throws Exception {
         CompatibilityBuildHelper buildHelper = new CompatibilityBuildHelper(getBuild());
-        final String fileName = APK_IN_APEX_TESTAPEX_NAME + "_v1.apex";
         final File apex = buildHelper.getTestFile(fileName);
         try {
             getDevice().enableAdbRoot();