2 files changed, 49 insertions, 1 deletions
diff --git a/tests/apex/Android.bp b/tests/apex/Android.bp
index 18f1bea75..9dfbdf589 100644
--- a/tests/apex/Android.bp
+++ b/tests/apex/Android.bp
@@ -31,6 +31,8 @@ android_test {
         "androidx.test.rules",
         "androidx.test.ext.junit",
         "androidx.test.ext.truth",
+        "com.android.permission.flags-aconfig-java",
+        "flag-junit",
         "mockito-target-extended-minus-junit4",
     ],
     jni_libs: [
diff --git a/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt b/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt
index 6500b3926..e9c93a33a 100644
--- a/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt
+++ b/tests/apex/java/com/android/role/persistence/RolesPersistenceTest.kt
@@ -20,12 +20,18 @@ import android.content.ApexEnvironment
 import android.content.Context
 import android.os.Process
 import android.os.UserHandle
+import android.platform.test.annotations.RequiresFlagsDisabled
+import android.platform.test.annotations.RequiresFlagsEnabled
+import android.platform.test.flag.junit.DeviceFlagsValueProvider
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.permission.flags.Flags
 import com.google.common.truth.Truth.assertThat
 import java.io.File
 import org.junit.After
+import org.junit.Assume.assumeFalse
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.Parameterized
@@ -49,11 +55,22 @@ class RolesPersistenceTest {
 
     private val persistence = RolesPersistenceImpl {}
     private val defaultRoles = mapOf(ROLE_NAME to setOf(HOLDER_1, HOLDER_2))
+    private val activeUserIds = mapOf(ROLE_NAME to USER_ID)
     private val stateVersionUndefined = RolesState(VERSION_UNDEFINED, PACKAGE_HASH, defaultRoles)
     private val stateVersionFallbackMigrated =
         RolesState(VERSION_FALLBACK_MIGRATED, PACKAGE_HASH, defaultRoles, setOf(ROLE_NAME))
+    private val stateVersionActiveUserIds =
+        RolesState(
+            VERSION_ACTIVE_USER_IDS,
+            PACKAGE_HASH,
+            defaultRoles,
+            setOf(ROLE_NAME),
+            activeUserIds,
+        )
     private val user = Process.myUserHandle()
 
+    @get:Rule val flagsRule = DeviceFlagsValueProvider.createCheckFlagsRule()
+
     @Before
     fun setUp() {
         createMockDataDirectory()
@@ -84,16 +101,41 @@ class RolesPersistenceTest {
         mockitoSession.finishMocking()
     }
 
+    @RequiresFlagsDisabled(Flags.FLAG_CROSS_USER_ROLE_ENABLED)
     @Test
     fun testWriteRead() {
+        assumeFalse(stateVersion == StateVersion.VERSION_ACTIVE_USER_IDS)
         persistence.writeForUser(state, user)
         val persistedState = persistence.readForUser(user)
 
         assertThat(persistedState).isEqualTo(state)
     }
 
+    @RequiresFlagsEnabled(Flags.FLAG_CROSS_USER_ROLE_ENABLED)
+    @Test
+    fun testWriteRead_supportsActiveUser() {
+        persistence.writeForUser(state, user)
+        val persistedState = persistence.readForUser(user)
+
+        assertThat(persistedState).isEqualTo(state)
+    }
+
+    @RequiresFlagsDisabled(Flags.FLAG_CROSS_USER_ROLE_ENABLED)
     @Test
     fun testWriteCorruptReadFromReserveCopy() {
+        assumeFalse(stateVersion == StateVersion.VERSION_ACTIVE_USER_IDS)
+        persistence.writeForUser(state, user)
+        // Corrupt the primary file.
+        RolesPersistenceImpl.getFile(user)
+            .writeText("<roles version=\"-1\"><role name=\"com.foo.bar\"><holder")
+        val persistedState = persistence.readForUser(user)
+
+        assertThat(persistedState).isEqualTo(state)
+    }
+
+    @RequiresFlagsEnabled(Flags.FLAG_CROSS_USER_ROLE_ENABLED)
+    @Test
+    fun testWriteCorruptReadFromReserveCopy_supportsActiveUser() {
         persistence.writeForUser(state, user)
         // Corrupt the primary file.
         RolesPersistenceImpl.getFile(user)
@@ -116,11 +158,13 @@ class RolesPersistenceTest {
         when (stateVersion) {
             StateVersion.VERSION_UNDEFINED -> stateVersionUndefined
             StateVersion.VERSION_FALLBACK_MIGRATED -> stateVersionFallbackMigrated
+            StateVersion.VERSION_ACTIVE_USER_IDS -> stateVersionActiveUserIds
         }
 
     enum class StateVersion {
         VERSION_UNDEFINED,
-        VERSION_FALLBACK_MIGRATED
+        VERSION_FALLBACK_MIGRATED,
+        VERSION_ACTIVE_USER_IDS,
     }
 
     companion object {
@@ -130,10 +174,12 @@ class RolesPersistenceTest {
 
         private const val VERSION_UNDEFINED = -1
         private const val VERSION_FALLBACK_MIGRATED = 1
+        private const val VERSION_ACTIVE_USER_IDS = 2
         private const val APEX_MODULE_NAME = "com.android.permission"
         private const val PACKAGE_HASH = "packagesHash"
         private const val ROLE_NAME = "roleName"
         private const val HOLDER_1 = "holder1"
         private const val HOLDER_2 = "holder2"
+        private const val USER_ID = 10
     }
 }