2 files changed, 21 insertions, 0 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 a616c61b34f8..6c9ffe2a7fac 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -477,4 +477,12 @@ public class StagedRollbackTest {
         StorageManager sm = (StorageManager) context.getSystemService(Context.STORAGE_SERVICE);
         assertThat(sm.isCheckpointSupported()).isTrue();
     }
+
+    @Test
+    public void hasMainlineModule() throws Exception {
+        String pkgName = getModuleMetadataPackageName();
+        boolean existed =  InstrumentationRegistry.getInstrumentation().getContext()
+                .getPackageManager().getModuleInfo(pkgName, 0) != null;
+        assertThat(existed).isTrue();
+    }
 }
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 282f012dbf6f..78775be84828 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
@@ -243,6 +243,7 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
      */
     @Test
     public void testRollbackWhitelistedApp() throws Exception {
+        assumeTrue(hasMainlineModule());
         runPhase("testRollbackWhitelistedApp_Phase1");
         getDevice().reboot();
         runPhase("testRollbackWhitelistedApp_Phase2");
@@ -460,4 +461,16 @@ public class StagedRollbackTest extends BaseHostJUnit4Test {
             return false;
         }
     }
+
+    /**
+     * True if this build has mainline modules installed.
+     */
+    private boolean hasMainlineModule() throws Exception {
+        try {
+            runPhase("hasMainlineModule");
+            return true;
+        } catch (AssertionError ignore) {
+            return false;
+        }
+    }
 }