diff --git a/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivity.kt b/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivity.kt
index e3a60781de7..c40ec96eaea 100644
--- a/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivity.kt
+++ b/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivity.kt
@@ -39,7 +39,6 @@ import androidx.lifecycle.viewmodel.MutableCreationExtras
 import com.android.settings.R
 import com.android.settings.Utils
 import com.android.settings.biometrics.BiometricEnrollBase
-import com.android.settings.biometrics2.data.repository.FingerprintRepository
 import com.android.settings.biometrics2.factory.BiometricsViewModelFactory
 import com.android.settings.biometrics2.factory.BiometricsViewModelFactory.CHALLENGE_GENERATOR_KEY
 import com.android.settings.biometrics2.factory.BiometricsViewModelFactory.ENROLLMENT_REQUEST_KEY
@@ -92,12 +91,6 @@ open class FingerprintEnrollmentActivity : FragmentActivity() {
     /** Internal activity for FingerprintSettings */
     class InternalActivity : FingerprintEnrollmentActivity()
 
-    /**
-     * This flag is used for addBackStack(), we do not save it in ViewModel because it is just used
-     * during FragmentManager calls
-     */
-    private var isFirstFragmentAdded = false
-
     private val viewModelProvider: ViewModelProvider by lazy {
         ViewModelProvider(this)
     }
@@ -168,7 +161,7 @@ open class FingerprintEnrollmentActivity : FragmentActivity() {
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
-        viewModel.setSavedInstanceState(savedInstanceState)
+        viewModel.onRestoreInstanceState(savedInstanceState)
         autoCredentialViewModel.setCredentialModel(savedInstanceState, intent)
 
         // Theme
@@ -189,10 +182,9 @@ open class FingerprintEnrollmentActivity : FragmentActivity() {
         }
         if (fragment == null) {
             checkCredential()
-            val request: EnrollmentRequest = viewModel.getRequest()
-            if (request.isSkipFindSensor) {
+            if (viewModel.request.isSkipFindSensor) {
                 startEnrollingFragment()
-            } else if (request.isSkipIntro) {
+            } else if (viewModel.request.isSkipIntro) {
                 startFindSensorFragment()
             } else {
                 startIntroFragment()
@@ -229,12 +221,12 @@ open class FingerprintEnrollmentActivity : FragmentActivity() {
     }
 
     private fun startFragment(fragmentClass: Class<out Fragment>, tag: String) {
-        if (!isFirstFragmentAdded) {
+        if (!viewModel.isFirstFragmentAdded) {
             supportFragmentManager.beginTransaction()
                 .setReorderingAllowed(true)
                 .replace(R.id.fragment_container_view, fragmentClass, null, tag)
                 .commit()
-            isFirstFragmentAdded = true
+            viewModel.setIsFirstFragmentAdded()
         } else {
             supportFragmentManager.beginTransaction()
                 .setReorderingAllowed(true)
diff --git a/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModel.java b/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModel.java
index 33799f39f68..4c702fa6f2e 100644
--- a/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModel.java
+++ b/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModel.java
@@ -53,12 +53,17 @@ public class FingerprintEnrollmentViewModel extends AndroidViewModel {
     @VisibleForTesting
     static final String SAVED_STATE_IS_NEW_FINGERPRINT_ADDED = "is_new_fingerprint_added";
 
+    @VisibleForTesting
+    static final String SAVED_STATE_IS_FIRST_FRAGMENT_ADDED = "is_first_fragment_added";
+
     @NonNull private final FingerprintRepository mFingerprintRepository;
 
     private final AtomicBoolean mIsWaitingActivityResult = new AtomicBoolean(false);
     private final MutableLiveData<ActivityResult> mSetResultLiveData = new MutableLiveData<>();
     @NonNull private final EnrollmentRequest mRequest;
     private boolean mIsNewFingerprintAdded = false;
+    /** Flag for FragmentManager::addToBackStack() */
+    private boolean mIsFirstFragmentAdded = false;
 
     public FingerprintEnrollmentViewModel(
             @NonNull Application application,
@@ -145,7 +150,7 @@ public class FingerprintEnrollmentViewModel extends AndroidViewModel {
     /**
      * Handle savedInstanceState from activity onCreated()
      */
-    public void setSavedInstanceState(@Nullable Bundle savedInstanceState) {
+    public void onRestoreInstanceState(@Nullable Bundle savedInstanceState) {
         if (savedInstanceState == null) {
             return;
         }
@@ -154,6 +159,8 @@ public class FingerprintEnrollmentViewModel extends AndroidViewModel {
         );
         mIsNewFingerprintAdded = savedInstanceState.getBoolean(
                 SAVED_STATE_IS_NEW_FINGERPRINT_ADDED, false);
+        mIsFirstFragmentAdded = savedInstanceState.getBoolean(
+                SAVED_STATE_IS_FIRST_FRAGMENT_ADDED, false);
     }
 
     /**
@@ -162,6 +169,7 @@ public class FingerprintEnrollmentViewModel extends AndroidViewModel {
     public void onSaveInstanceState(@NonNull Bundle outState) {
         outState.putBoolean(SAVED_STATE_IS_WAITING_ACTIVITY_RESULT, mIsWaitingActivityResult.get());
         outState.putBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED, mIsNewFingerprintAdded);
+        outState.putBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED, mIsFirstFragmentAdded);
     }
 
     /**
@@ -193,6 +201,17 @@ public class FingerprintEnrollmentViewModel extends AndroidViewModel {
         mIsNewFingerprintAdded = true;
     }
 
+    public boolean isFirstFragmentAdded() {
+        return mIsFirstFragmentAdded;
+    }
+
+    /**
+     * set mIsFirstFragmentAdded to true, this flag will be used during adding fragment
+     */
+    public void setIsFirstFragmentAdded() {
+        mIsFirstFragmentAdded = true;
+    }
+
     /**
      * Update FINGERPRINT_SUGGESTION_ACTIVITY into package manager
      */
diff --git a/tests/uitests/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivityTest.kt b/tests/uitests/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivityTest.kt
index 49c5ac94411..33c8d3dc017 100644
--- a/tests/uitests/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivityTest.kt
+++ b/tests/uitests/src/com/android/settings/biometrics2/ui/view/FingerprintEnrollmentActivityTest.kt
@@ -22,13 +22,13 @@ import android.hardware.fingerprint.FingerprintManager
 import android.hardware.fingerprint.FingerprintSensorPropertiesInternal
 import android.hardware.fingerprint.IFingerprintAuthenticatorsRegisteredCallback
 import android.os.UserHandle
-import android.support.test.uiautomator.By
-import android.support.test.uiautomator.UiDevice
-import android.support.test.uiautomator.UiObject2
-import android.support.test.uiautomator.Until
 import android.util.Log
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.platform.app.InstrumentationRegistry
+import androidx.test.uiautomator.By
+import androidx.test.uiautomator.UiDevice
+import androidx.test.uiautomator.UiObject2
+import androidx.test.uiautomator.Until
 import com.android.internal.widget.LockPatternChecker
 import com.android.internal.widget.LockPatternUtils
 import com.android.internal.widget.LockscreenCredential
@@ -57,6 +57,7 @@ class FingerprintEnrollmentActivityTest {
     private var canAssumeUdfps = false
     private var canAssumeSfps = false
     private var enrollingPageTitle: String = ""
+    private var runAsLandscape = false
 
     private val device: UiDevice by lazy {
         UiDevice.getInstance(InstrumentationRegistry.getInstrumentation())
@@ -108,12 +109,16 @@ class FingerprintEnrollmentActivityTest {
     @After
     @Throws(Exception::class)
     fun tearDown() {
+        runAsLandscape = false
+        setDeviceOrientation()
+
         LockScreenUtil.resetLockscreen(TEST_PIN)
         device.pressHome()
     }
 
     @Test
     fun testIntroChooseLock() {
+        setDeviceOrientation()
         val intent = newActivityIntent(false)
         context.startActivity(intent)
         assertThat(
@@ -124,6 +129,12 @@ class FingerprintEnrollmentActivityTest {
         ).isTrue()
     }
 
+    @Test
+    fun testIntroChooseLock_landscape() {
+        runAsLandscape = true
+        testIntroChooseLock()
+    }
+
     private fun verifyIntroPage() {
         device.waitForIdle()
         run {
@@ -141,7 +152,8 @@ class FingerprintEnrollmentActivityTest {
         // Click more btn at most twice and the introduction should stay in the last page
         var moreBtn: UiObject2? = null
         var i = 0
-        while (i < 2 && device.findObject(By.text("More")).also { moreBtn = it } != null) {
+        val more = if (runAsLandscape) 5 else 2
+        while (i < more && device.findObject(By.text("More")).also { moreBtn = it } != null) {
             moreBtn!!.click()
             device.waitForIdle()
             device.wait(Until.hasObject(By.text("More")), IDLE_TIMEOUT)
@@ -154,6 +166,8 @@ class FingerprintEnrollmentActivityTest {
     @Test
     fun testIntroWithGkPwHandle_withUdfps_clickStart() {
         Assume.assumeTrue(canAssumeUdfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(false)
 
@@ -178,9 +192,17 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_withUdfps_clickStart_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_withUdfps_clickStart()
+    }
+
     @Test
     fun testIntroWithGkPwHandle_withUdfps_clickLottie() {
         Assume.assumeTrue(canAssumeUdfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(false)
 
@@ -203,9 +225,17 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_withUdfps_clickLottie_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_withUdfps_clickLottie()
+    }
+
     @Test
     fun testIntroWithGkPwHandle_withSfps() {
         Assume.assumeTrue(canAssumeSfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(false)
 
@@ -225,9 +255,17 @@ class FingerprintEnrollmentActivityTest {
         // We don't have view which can be clicked to run to next page, stop at here.
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_withSfps_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_withSfps()
+    }
+
     @Test
     fun testIntroWithGkPwHandle_withRfps() {
         Assume.assumeFalse(canAssumeUdfps || canAssumeSfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(false)
 
@@ -252,8 +290,15 @@ class FingerprintEnrollmentActivityTest {
         }
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_withRfps_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_withRfps()
+    }
+
     @Test
     fun testIntroWithGkPwHandle_clickNoThanksInIntroPage() {
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(false)
 
@@ -268,8 +313,15 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.findObject(By.text("No thanks"))).isNull()
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_clickNoThanksInIntroPage_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_clickNoThanksInIntroPage()
+    }
+
     @Test
     fun testIntroWithGkPwHandle_clickSkipInFindSensor() {
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(false)
 
@@ -291,8 +343,15 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.findObject(By.text(DO_IT_LATER))).isNull()
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_clickSkipInFindSensor_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_clickSkipInFindSensor()
+    }
+
     @Test
     fun testIntroWithGkPwHandle_clickSkipAnywayInFindFpsDialog_whenIsSuw() {
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(true)
 
@@ -322,8 +381,15 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.findObject(By.text(DO_IT_LATER))).isNull()
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_clickSkipAnywayInFindFpsDialog_whenIsSuw_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_clickSkipAnywayInFindFpsDialog_whenIsSuw()
+    }
+
     @Test
     fun testIntroWithGkPwHandle_clickGoBackInFindFpsDialog_whenIsSuw() {
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchIntroWithGkPwHandle(true)
 
@@ -351,8 +417,15 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(DO_IT_LATER)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testIntroWithGkPwHandle_clickGoBackInFindFpsDialog_whenIsSuw_landscape() {
+        runAsLandscape = true
+        testIntroWithGkPwHandle_clickGoBackInFindFpsDialog_whenIsSuw()
+    }
+
     @Test
     fun testIntroCheckPin() {
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         val intent = newActivityIntent(false)
         context.startActivity(intent)
@@ -366,6 +439,7 @@ class FingerprintEnrollmentActivityTest {
 
     @Test
     fun testEnrollingWithGkPwHandle() {
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchEnrollingWithGkPwHandle()
 
@@ -374,9 +448,17 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testEnrollingWithGkPwHandle_landscape() {
+        runAsLandscape = true
+        testEnrollingWithGkPwHandle()
+    }
+
     @Test
     fun testEnrollingIconTouchDialog_withSfps() {
         Assume.assumeTrue(canAssumeSfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchEnrollingWithGkPwHandle()
 
@@ -409,9 +491,17 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testEnrollingIconTouchDialog_withSfps_landscape() {
+        runAsLandscape = true
+        testEnrollingIconTouchDialog_withSfps()
+    }
+
     @Test
     fun testEnrollingIconTouchDialog_withRfps() {
         Assume.assumeFalse(canAssumeUdfps || canAssumeSfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchEnrollingWithGkPwHandle()
 
@@ -443,9 +533,17 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testEnrollingIconTouchDialog_withRfps_landscape() {
+        runAsLandscape = true
+        testEnrollingIconTouchDialog_withRfps()
+    }
+
     @Test
     fun testFindUdfpsWithGkPwHandle_clickStart() {
         Assume.assumeTrue(canAssumeUdfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchFindSensorWithGkPwHandle()
 
@@ -464,9 +562,54 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testFindUdfpsWithGkPwHandle_clickStart_landscape() {
+        runAsLandscape = true
+        testFindUdfpsWithGkPwHandle_clickStart()
+    }
+
+    @Test
+    fun testFindUdfpsLandscapeWithGkPwHandle_clickStartThenBack() {
+        Assume.assumeTrue(canAssumeUdfps)
+
+        setDeviceOrientation()
+        LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
+        launchFindSensorWithGkPwHandle()
+
+        // FindUdfps page (portrait)
+        assertThat(device.wait(Until.hasObject(By.text(DO_IT_LATER)), IDLE_TIMEOUT)).isTrue()
+
+        // rotate device
+        device.setOrientationLandscape()
+        device.waitForIdle()
+
+        // FindUdfps page (landscape)
+        assertThat(device.wait(Until.hasObject(By.text(DO_IT_LATER)), IDLE_TIMEOUT)).isTrue()
+        val lottie = device.findObject(
+            By.res(SETTINGS_PACKAGE_NAME, "illustration_lottie")
+        )
+        assertThat(lottie).isNotNull()
+        assertThat(lottie.isClickable).isTrue()
+        val startBtn = device.findObject(By.text("Start"))
+        assertThat(startBtn.isClickable).isTrue()
+        startBtn.click()
+
+        // Enrolling page
+        assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
+
+        // Press back
+        device.pressBack()
+        device.waitForIdle()
+
+        // FindUdfps page (landscape-again)
+        assertThat(device.wait(Until.hasObject(By.text(DO_IT_LATER)), IDLE_TIMEOUT)).isTrue()
+    }
+
     @Test
     fun testFindUdfpsWithGkPwHandle_clickLottie() {
         Assume.assumeTrue(canAssumeUdfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchFindSensorWithGkPwHandle()
 
@@ -485,9 +628,17 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.hasObject(By.text(enrollingPageTitle)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testFindUdfpsWithGkPwHandle_clickLottie_landscape() {
+        runAsLandscape = true
+        testFindUdfpsWithGkPwHandle_clickLottie()
+    }
+
     @Test
     fun testFindSfpsWithGkPwHandle() {
         Assume.assumeTrue(canAssumeSfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchFindSensorWithGkPwHandle()
 
@@ -501,9 +652,17 @@ class FingerprintEnrollmentActivityTest {
         // We don't have view which can be clicked to run to next page, stop at here.
     }
 
+    @Test
+    fun testFindSfpsWithGkPwHandle_landscape() {
+        runAsLandscape = true
+        testFindSfpsWithGkPwHandle()
+    }
+
     @Test
     fun testFindRfpsWithGkPwHandle() {
         Assume.assumeFalse(canAssumeUdfps || canAssumeSfps)
+
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchFindSensorWithGkPwHandle()
 
@@ -528,8 +687,15 @@ class FingerprintEnrollmentActivityTest {
         }
     }
 
+    @Test
+    fun testFindRfpsWithGkPwHandle_landscape() {
+        runAsLandscape = true
+        testFindRfpsWithGkPwHandle()
+    }
+
     @Test
     fun testFindSensorWithGkPwHandle_clickSkipInFindSensor() {
+        setDeviceOrientation()
         LockScreenUtil.setLockscreen(LockScreenUtil.LockscreenType.PIN, TEST_PIN, true)
         launchFindSensorWithGkPwHandle()
 
@@ -545,6 +711,12 @@ class FingerprintEnrollmentActivityTest {
         assertThat(device.wait(Until.gone(By.text(DO_IT_LATER)), IDLE_TIMEOUT)).isTrue()
     }
 
+    @Test
+    fun testFindSensorWithGkPwHandle_clickSkipInFindSensor_landscape() {
+        runAsLandscape = true
+        testFindSensorWithGkPwHandle_clickSkipInFindSensor()
+    }
+
     private fun launchIntroWithGkPwHandle(isSuw: Boolean) {
         val lockPatternUtils = LockPatternUtils(context)
         val lockscreenCredential = LockscreenCredential.createPin(TEST_PIN)
@@ -610,6 +782,15 @@ class FingerprintEnrollmentActivityTest {
         return intent
     }
 
+    private fun setDeviceOrientation() {
+        if (runAsLandscape) {
+            device.setOrientationLandscape()
+        } else {
+            device.setOrientationPortrait()
+        }
+        device.waitForIdle()
+    }
+
     companion object {
         private const val TAG = "FingerprintEnrollmentActivityTest"
         const val SETTINGS_PACKAGE_NAME = "com.android.settings"
diff --git a/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModelTest.java b/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModelTest.java
index 7850677916c..73fab791e1a 100644
--- a/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModelTest.java
+++ b/tests/unit/src/com/android/settings/biometrics2/ui/viewmodel/FingerprintEnrollmentViewModelTest.java
@@ -18,8 +18,8 @@ package com.android.settings.biometrics2.ui.viewmodel;
 
 import static android.hardware.fingerprint.FingerprintSensorProperties.TYPE_UDFPS_OPTICAL;
 
+import static com.android.settings.biometrics2.ui.viewmodel.FingerprintEnrollmentViewModel.SAVED_STATE_IS_FIRST_FRAGMENT_ADDED;
 import static com.android.settings.biometrics2.ui.viewmodel.FingerprintEnrollmentViewModel.SAVED_STATE_IS_NEW_FINGERPRINT_ADDED;
-import static com.android.settings.biometrics2.ui.viewmodel.FingerprintEnrollmentViewModel.SAVED_STATE_IS_WAITING_ACTIVITY_RESULT;
 import static com.android.settings.biometrics2.utils.EnrollmentRequestUtils.newAllFalseRequest;
 import static com.android.settings.biometrics2.utils.FingerprintRepositoryUtils.newFingerprintRepository;
 import static com.android.settings.biometrics2.utils.FingerprintRepositoryUtils.setupFingerprintEnrolledFingerprints;
@@ -58,13 +58,17 @@ public class FingerprintEnrollmentViewModelTest {
     private FingerprintRepository mFingerprintRepository;
     private FingerprintEnrollmentViewModel mViewModel;
 
+    private FingerprintEnrollmentViewModel newViewModelInstance() {
+        return new FingerprintEnrollmentViewModel(mApplication, mFingerprintRepository,
+                newAllFalseRequest(mApplication));
+    }
+
     @Before
     public void setUp() {
         mApplication = ApplicationProvider.getApplicationContext();
         mFingerprintRepository = newFingerprintRepository(mFingerprintManager, TYPE_UDFPS_OPTICAL,
                 5);
-        mViewModel = new FingerprintEnrollmentViewModel(mApplication, mFingerprintRepository,
-                newAllFalseRequest(mApplication));
+        mViewModel = newViewModelInstance();
     }
 
     @Test
@@ -73,63 +77,137 @@ public class FingerprintEnrollmentViewModelTest {
     }
 
     @Test
-    public void testSetSavedInstanceState() {
-        // setSavedInstanceState() as false
-        final Bundle bundle = new Bundle();
-        final Bundle outBundle = new Bundle();
-
-        // Set SAVED_STATE_IS_WAITING_ACTIVITY_RESULT to true
-        bundle.putBoolean(SAVED_STATE_IS_WAITING_ACTIVITY_RESULT, false);
-        mViewModel.setSavedInstanceState(bundle);
+    public void testIsWaitingActivityResult() {
+        // Default false
         assertThat(mViewModel.isWaitingActivityResult().get()).isFalse();
 
-        // Set SAVED_STATE_IS_WAITING_ACTIVITY_RESULT to true
-        bundle.clear();
-        bundle.putBoolean(SAVED_STATE_IS_WAITING_ACTIVITY_RESULT, true);
-        mViewModel.setSavedInstanceState(bundle);
-        assertThat(mViewModel.isWaitingActivityResult().get()).isTrue();
+        // false if null bundle
+        mViewModel = newViewModelInstance();
+        mViewModel.onRestoreInstanceState(null);
+        assertThat(mViewModel.isWaitingActivityResult().get()).isFalse();
 
-        // Set SAVED_STATE_IS_NEW_FINGERPRINT_ADDED to false
-        bundle.clear();
-        bundle.putBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED, false);
-        mViewModel.setSavedInstanceState(bundle);
-        outBundle.clear();
-        mViewModel.onSaveInstanceState(outBundle);
-        assertThat(outBundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isFalse();
+        // false if empty bundle
+        mViewModel.onRestoreInstanceState(new Bundle());
+        assertThat(mViewModel.isWaitingActivityResult().get()).isFalse();
 
-        // Set SAVED_STATE_IS_NEW_FINGERPRINT_ADDED to true
-        bundle.clear();
-        bundle.putBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED, true);
-        mViewModel.setSavedInstanceState(bundle);
-        outBundle.clear();
-        mViewModel.onSaveInstanceState(outBundle);
-        assertThat(outBundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        // False value can be saved during onSaveInstanceState() and restore after
+        // onSaveInstanceState()
+        final Bundle falseSavedInstance = new Bundle();
+        mViewModel.onSaveInstanceState(falseSavedInstance);
+        final FingerprintEnrollmentViewModel falseViewModel = newViewModelInstance();
+        falseViewModel.onRestoreInstanceState(falseSavedInstance);
+        assertThat(falseViewModel.isWaitingActivityResult().get()).isFalse();
+
+        // True value can be saved during onSaveInstanceState() and restore after
+        // onSaveInstanceState()
+        final Bundle trueSavedInstance = new Bundle();
+        mViewModel.isWaitingActivityResult().set(true);
+        mViewModel.onSaveInstanceState(trueSavedInstance);
+        final FingerprintEnrollmentViewModel trueViewModel = newViewModelInstance();
+        trueViewModel.onRestoreInstanceState(trueSavedInstance);
+        assertThat(trueViewModel.isWaitingActivityResult().get()).isTrue();
     }
 
     @Test
-    public void testOnSaveInstanceState() {
-        // Test isWaitingActivityResult false
-        mViewModel.isWaitingActivityResult().set(false);
-        final Bundle bundle = new Bundle();
-        mViewModel.onSaveInstanceState(bundle);
-        assertThat(bundle.getBoolean(SAVED_STATE_IS_WAITING_ACTIVITY_RESULT)).isFalse();
+    public void testIsNewFingerprintAdded() {
+        // Default false
+        final Bundle outBundle = new Bundle();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isFalse();
 
-        // Test isWaitingActivityResult true
-        mViewModel.isWaitingActivityResult().set(true);
-        bundle.clear();
-        mViewModel.onSaveInstanceState(bundle);
-        assertThat(bundle.getBoolean(SAVED_STATE_IS_WAITING_ACTIVITY_RESULT)).isTrue();
+        // false if null bundle
+        mViewModel = newViewModelInstance();
+        mViewModel.onRestoreInstanceState(null);
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isFalse();
 
-        // Test isNewFingerprintAdded default false
-        bundle.clear();
-        mViewModel.onSaveInstanceState(bundle);
-        assertThat(bundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isFalse();
+        // false if empty bundle
+        mViewModel = newViewModelInstance();
+        mViewModel.onRestoreInstanceState(new Bundle());
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isFalse();
 
-        // Test isNewFingerprintAdded true
-        mViewModel.setIsNewFingerprintAdded();
-        bundle.clear();
-        mViewModel.onSaveInstanceState(bundle);
-        assertThat(bundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        // False value can be saved during onSaveInstanceState() and restore after
+        // onSaveInstanceState()
+        final Bundle falseSavedInstance = new Bundle();
+        falseSavedInstance.putBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED, false);
+        mViewModel.onRestoreInstanceState(falseSavedInstance);
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isFalse();
+
+        // True value can be saved during onSaveInstanceState() and restore after
+        // onSaveInstanceState()
+        final Bundle trueSavedInstance = new Bundle();
+        trueSavedInstance.putBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED, true);
+        mViewModel.onRestoreInstanceState(trueSavedInstance);
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+
+        // setIsFirstFragmentAdded() can be saved during onSaveInstanceState()
+        mViewModel.setIsFirstFragmentAdded();
+        mViewModel.onSaveInstanceState(trueSavedInstance);
+        assertThat(trueSavedInstance.containsKey(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+        assertThat(trueSavedInstance.getBoolean(SAVED_STATE_IS_NEW_FINGERPRINT_ADDED)).isTrue();
+    }
+
+    @Test
+    public void testIsFirstFragmentAdded() {
+        // Default false
+        final Bundle outBundle = new Bundle();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isFalse();
+
+        // false if null bundle
+        mViewModel = newViewModelInstance();
+        mViewModel.onRestoreInstanceState(null);
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isFalse();
+
+        // false if empty bundle
+        mViewModel = newViewModelInstance();
+        mViewModel.onRestoreInstanceState(new Bundle());
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isFalse();
+
+        // False value can be saved during onSaveInstanceState() and restore after
+        // onSaveInstanceState()
+        final Bundle falseSavedInstance = new Bundle();
+        falseSavedInstance.putBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED, false);
+        mViewModel.onRestoreInstanceState(falseSavedInstance);
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isFalse();
+
+        // True value can be saved during onSaveInstanceState() and restore after
+        // onSaveInstanceState()
+        final Bundle trueSavedInstance = new Bundle();
+        trueSavedInstance.putBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED, true);
+        mViewModel.onRestoreInstanceState(trueSavedInstance);
+        outBundle.clear();
+        mViewModel.onSaveInstanceState(outBundle);
+        assertThat(outBundle.containsKey(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
+        assertThat(outBundle.getBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
+
+        // setIsFirstFragmentAdded() can be saved during onSaveInstanceState()
+        mViewModel.setIsFirstFragmentAdded();
+        mViewModel.onSaveInstanceState(trueSavedInstance);
+        assertThat(trueSavedInstance.containsKey(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
+        assertThat(trueSavedInstance.getBoolean(SAVED_STATE_IS_FIRST_FRAGMENT_ADDED)).isTrue();
     }
 
     @Test