Support landscape mode for provisioning logic

Test: atest ManagedProvisioningTests
Test: atest ManagedProvisioningRoboTests
Bug: 166531689
Change-Id: I318feca535fa15cf086ac6b88d1ada1a2b54c559
diff --git a/src/com/android/managedprovisioning/provisioning/AbstractProvisioningActivity.java b/src/com/android/managedprovisioning/provisioning/AbstractProvisioningActivity.java
index 41dadbb..8bafe16 100644
--- a/src/com/android/managedprovisioning/provisioning/AbstractProvisioningActivity.java
+++ b/src/com/android/managedprovisioning/provisioning/AbstractProvisioningActivity.java
@@ -50,11 +50,13 @@
 
     static final int STATE_PROVISIONING_INTIIALIZING = 1;
     static final int STATE_PROVISIONING_STARTED = 2;
-    static final int STATE_PROVISIONING_FINALIZED = 3;
+    static final int STATE_PROVISIONING_COMPLETED = 3;
+    static final int STATE_PROVISIONING_FINALIZED = 4;
 
     @Retention(RetentionPolicy.SOURCE)
     @IntDef({STATE_PROVISIONING_INTIIALIZING,
             STATE_PROVISIONING_STARTED,
+            STATE_PROVISIONING_COMPLETED,
             STATE_PROVISIONING_FINALIZED})
     private @interface ProvisioningState {}
 
@@ -65,7 +67,6 @@
     @VisibleForTesting static final String CANCEL_PROVISIONING_DIALOG_RESET
             = "CancelProvisioningDialogReset";
 
-    protected ProvisioningManagerInterface mProvisioningManager;
     protected ProvisioningParams mParams;
     protected @ProvisioningState int mState;
 
diff --git a/src/com/android/managedprovisioning/provisioning/AdminIntegratedFlowPrepareActivity.java b/src/com/android/managedprovisioning/provisioning/AdminIntegratedFlowPrepareActivity.java
index c9663a0..d5136e6 100644
--- a/src/com/android/managedprovisioning/provisioning/AdminIntegratedFlowPrepareActivity.java
+++ b/src/com/android/managedprovisioning/provisioning/AdminIntegratedFlowPrepareActivity.java
@@ -47,6 +47,7 @@
 public class AdminIntegratedFlowPrepareActivity extends AbstractProvisioningActivity {
 
     private RepeatingVectorAnimation mRepeatingVectorAnimation;
+    private AdminIntegratedFlowPrepareManager mProvisioningManager;
 
     public AdminIntegratedFlowPrepareActivity() {
         this(new Utils(), new SettingsFacade(),
diff --git a/src/com/android/managedprovisioning/provisioning/ProvisioningActivity.java b/src/com/android/managedprovisioning/provisioning/ProvisioningActivity.java
index 71409bd..a9c7fe3 100644
--- a/src/com/android/managedprovisioning/provisioning/ProvisioningActivity.java
+++ b/src/com/android/managedprovisioning/provisioning/ProvisioningActivity.java
@@ -21,6 +21,8 @@
 import static com.android.internal.logging.nano.MetricsProto.MetricsEvent.PROVISIONING_PROVISIONING_ACTIVITY_TIME_MS;
 import static com.android.internal.util.Preconditions.checkNotNull;
 
+import static java.util.Objects.requireNonNull;
+
 import android.Manifest.permission;
 import android.annotation.IntDef;
 import android.app.Activity;
@@ -97,7 +99,6 @@
     static final int PROVISIONING_MODE_WORK_PROFILE_ON_FULLY_MANAGED_DEVICE = 3;
     static final int PROVISIONING_MODE_FINANCED_DEVICE = 4;
     static final int PROVISIONING_MODE_WORK_PROFILE_ON_ORG_OWNED_DEVICE = 5;
-
     @IntDef(prefix = { "PROVISIONING_MODE_" }, value = {
         PROVISIONING_MODE_WORK_PROFILE,
         PROVISIONING_MODE_FULLY_MANAGED_DEVICE,
@@ -125,6 +126,7 @@
     private RepeatingVectorAnimation mRepeatingVectorAnimation;
     private UserProvisioningStateHelper mUserProvisioningStateHelper;
     private PolicyComplianceUtils mPolicyComplianceUtils;
+    private ProvisioningManager mProvisioningManager;
 
     public ProvisioningActivity() {
         this(
@@ -153,22 +155,34 @@
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
 
+        // assign this Activity as the view store owner to access saved state and receive updates
+        getProvisioningManager().setViewModelStoreOwner(this);
+
         if (mUserProvisioningStateHelper == null) {
             mUserProvisioningStateHelper = new UserProvisioningStateHelper(this);
         }
+
+        if (mState == STATE_PROVISIONING_FINALIZED) {
+            updateProvisioningFinalizedScreen();
+        }
     }
 
     @Override
-    protected ProvisioningManagerInterface getProvisioningManager() {
+    protected ProvisioningManager getProvisioningManager() {
         if (mProvisioningManager == null) {
             mProvisioningManager = ProvisioningManager.getInstance(this);
         }
         return mProvisioningManager;
     }
 
+    @VisibleForTesting
+    protected void setProvisioningManager(ProvisioningManager provisioningManager) {
+        mProvisioningManager = requireNonNull(provisioningManager);
+    }
+
     @Override
     public void preFinalizationCompleted() {
-        if (mState == STATE_PROVISIONING_FINALIZED) {
+        if (mState == STATE_PROVISIONING_COMPLETED || mState == STATE_PROVISIONING_FINALIZED) {
             return;
         }
 
@@ -182,13 +196,12 @@
 
         ProvisionLogger.logi("ProvisioningActivity pre-finalization completed");
 
-        // TODO: call this for the new flow after new NFC flow has been added
-        // maybeLaunchNfcUserSetupCompleteIntent();
+        // TODO(183094412): Decouple state from AbstractProvisioningActivity
+        mState = STATE_PROVISIONING_COMPLETED;
 
         if (shouldSkipEducationScreens() || mTransitionAnimationHelper.areAllTransitionsShown()) {
             updateProvisioningFinalizedScreen();
         }
-        mState = STATE_PROVISIONING_FINALIZED;
     }
 
     // Enforces DPCs to implement the POLICY_COMPLIANCE handler for NFC and financed device
@@ -218,6 +231,9 @@
         if (shouldSkipEducationScreens() || Utils.isSilentProvisioning(this, mParams)) {
             onNextButtonClicked();
         }
+
+        // TODO(183094412): Decouple state from AbstractProvisioningActivity
+        mState = STATE_PROVISIONING_FINALIZED;
     }
 
     @VisibleForTesting
@@ -306,7 +322,8 @@
 
     @Override
     protected void decideCancelProvisioningDialog() {
-        if (mState == STATE_PROVISIONING_FINALIZED && !mParams.isOrganizationOwnedProvisioning) {
+        if ((mState == STATE_PROVISIONING_COMPLETED || mState == STATE_PROVISIONING_FINALIZED)
+                && !mParams.isOrganizationOwnedProvisioning) {
             return;
         }
 
@@ -336,16 +353,25 @@
         } else {
             endTransitionAnimation();
         }
+        // remove this Activity as the view store owner to avoid memory leaks
+        if (isFinishing()) {
+            getProvisioningManager().clearViewModelStoreOwner();
+        }
     }
 
     @Override
     public void onAllTransitionsShown() {
-        if (mState == STATE_PROVISIONING_FINALIZED) {
+        if (mState == STATE_PROVISIONING_COMPLETED) {
             updateProvisioningFinalizedScreen();
         }
     }
 
     @Override
+    public void onTransitionStart(int screenIndex, AnimatedVectorDrawable currentAnimation) {
+        getProvisioningManager().setCurrentTransitionAnimation(screenIndex);
+    }
+
+    @Override
     protected void initializeUi(ProvisioningParams params) {
         final boolean isPoProvisioning = mUtils.isProfileOwnerAction(params.provisioningAction);
         final int titleResId =
@@ -405,7 +431,9 @@
                         drawable, providerInfo);
         mTransitionAnimationHelper = new TransitionAnimationHelper(provisioningMode,
                 /* adminCanGrantSensorsPermissions= */ !mParams.deviceOwnerPermissionGrantOptOut,
-                animationComponents, this);
+                animationComponents,
+                this,
+                getProvisioningManager().getCurrentTransitionAnimation());
     }
 
     private @ProvisioningMode int getProvisioningMode() {
diff --git a/src/com/android/managedprovisioning/provisioning/ProvisioningManager.java b/src/com/android/managedprovisioning/provisioning/ProvisioningManager.java
index 706b944..9fbf39a 100644
--- a/src/com/android/managedprovisioning/provisioning/ProvisioningManager.java
+++ b/src/com/android/managedprovisioning/provisioning/ProvisioningManager.java
@@ -17,12 +17,16 @@
 package com.android.managedprovisioning.provisioning;
 
 import static com.android.internal.logging.nano.MetricsProto.MetricsEvent.PROVISIONING_TOTAL_TASK_TIME_MS;
-import static com.android.internal.util.Preconditions.checkNotNull;
 import static com.android.managedprovisioning.analytics.ProvisioningAnalyticsTracker.CANCELLED_DURING_PROVISIONING;
 
+import static java.util.Objects.requireNonNull;
+
 import android.content.Context;
-import android.os.Handler;
-import android.os.Looper;
+import android.os.Bundle;
+
+import androidx.appcompat.app.AppCompatActivity;
+import androidx.lifecycle.ViewModelProvider;
+import androidx.lifecycle.ViewModelStoreOwner;
 
 import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
@@ -33,10 +37,16 @@
 import com.android.managedprovisioning.common.ProvisionLogger;
 import com.android.managedprovisioning.common.SettingsFacade;
 import com.android.managedprovisioning.model.ProvisioningParams;
+import com.android.managedprovisioning.provisioning.ProvisioningViewModel.ProvisioningViewModelFactory;
 
 /**
  * Singleton instance that provides communications between the ongoing provisioning process and the
  * UI layer.
+ * <p>{@link #setViewModelStoreOwner(ViewModelStoreOwner)} must be called when the view model store
+ * owner has been created (for example, {@link AppCompatActivity#onCreate(Bundle)}). In addition,
+ * to prevent memory leaks when the view model store owner is destroyed, {@link
+ * #clearViewModelStoreOwner()} must be called (for example, in {@link
+ * AppCompatActivity#onDestroy()} when {@link AppCompatActivity#isFinishing()} is {@code true}).
  */
 public class ProvisioningManager implements ProvisioningControllerCallback,
         ProvisioningManagerInterface {
@@ -48,6 +58,7 @@
     private final ProvisioningAnalyticsTracker mProvisioningAnalyticsTracker;
     private final TimeLogger mTimeLogger;
     private final ProvisioningManagerHelper mHelper;
+    private ProvisioningViewModel mViewModel;
 
     @GuardedBy("this")
     private AbstractProvisioningController mController;
@@ -62,7 +73,6 @@
     private ProvisioningManager(Context context) {
         this(
                 context,
-                new Handler(Looper.getMainLooper()),
                 new ProvisioningControllerFactory(),
                 new ProvisioningAnalyticsTracker(
                         MetricsWriterFactory.getMetricsWriter(context, new SettingsFacade()),
@@ -73,17 +83,37 @@
     @VisibleForTesting
     ProvisioningManager(
             Context context,
-            Handler uiHandler,
             ProvisioningControllerFactory factory,
             ProvisioningAnalyticsTracker analyticsTracker,
             TimeLogger timeLogger) {
-        mContext = checkNotNull(context);
-        mFactory = checkNotNull(factory);
-        mProvisioningAnalyticsTracker = checkNotNull(analyticsTracker);
-        mTimeLogger = checkNotNull(timeLogger);
+        mContext = requireNonNull(context);
+        mFactory = requireNonNull(factory);
+        mProvisioningAnalyticsTracker = requireNonNull(analyticsTracker);
+        mTimeLogger = requireNonNull(timeLogger);
         mHelper = new ProvisioningManagerHelper(context);
     }
 
+    /**
+     * Sets the view model store owner.
+     * <p>Must be called when the view model store owner has been created (for example, {@link
+     * AppCompatActivity#onCreate(Bundle)}).
+     * @see #clearViewModelStoreOwner()
+     */
+    void setViewModelStoreOwner(ViewModelStoreOwner owner) {
+        mViewModel = new ViewModelProvider(owner, new ProvisioningViewModelFactory())
+                .get(ProvisioningViewModel.class);
+    }
+
+    /**
+     * Clears the view model store owner.
+     * <p>Must be called when the view model store owner is getting destroyed (for example, in
+     * {@link AppCompatActivity#onDestroy()} when {@link AppCompatActivity#isFinishing()} is {@code
+     * true}), in order to clean the reference and prevent memory leaks.
+     */
+    void clearViewModelStoreOwner() {
+        mViewModel = null;
+    }
+
     @Override
     public void maybeStartProvisioning(final ProvisioningParams params) {
         synchronized (this) {
@@ -149,6 +179,14 @@
         mHelper.error(titleId, messageId, factoryResetRequired);
     }
 
+    void setCurrentTransitionAnimation(int currentTransitionAnimation) {
+        mViewModel.setCurrentTransitionScreen(currentTransitionAnimation);
+    }
+
+    int getCurrentTransitionAnimation() {
+        return mViewModel.getCurrentTransitionScreen();
+    }
+
     private AbstractProvisioningController getController(ProvisioningParams params) {
         return mFactory.createProvisioningController(mContext, params, this);
     }
diff --git a/src/com/android/managedprovisioning/provisioning/ProvisioningViewModel.java b/src/com/android/managedprovisioning/provisioning/ProvisioningViewModel.java
new file mode 100644
index 0000000..c746767
--- /dev/null
+++ b/src/com/android/managedprovisioning/provisioning/ProvisioningViewModel.java
@@ -0,0 +1,53 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.managedprovisioning.provisioning;
+
+import androidx.lifecycle.ViewModel;
+import androidx.lifecycle.ViewModelProvider;
+
+/**
+ * An {@link ViewModel} which maintains data related to provisioning.
+ */
+public class ProvisioningViewModel extends ViewModel {
+
+    private int mCurrentTransitionScreen;
+
+    /**
+     * Returns the index of the last playing transition screen.
+     */
+    public int getCurrentTransitionScreen() {
+        return mCurrentTransitionScreen;
+    }
+
+    /**
+     * Sets the index of the currently playing transition screen.
+     */
+    public void setCurrentTransitionScreen(int currentTransitionScreen) {
+        mCurrentTransitionScreen = currentTransitionScreen;
+    }
+
+    static class ProvisioningViewModelFactory implements ViewModelProvider.Factory {
+        @Override
+        public <T extends ViewModel> T create(Class<T> modelClass) {
+            if (!ProvisioningViewModel.class.isAssignableFrom(modelClass)) {
+                throw new IllegalArgumentException("Invalid class for creating a "
+                        + "PreProvisioningViewModel: " + modelClass);
+            }
+            return (T) new ProvisioningViewModel();
+        }
+    }
+}
diff --git a/src/com/android/managedprovisioning/provisioning/TransitionAnimationHelper.java b/src/com/android/managedprovisioning/provisioning/TransitionAnimationHelper.java
index 495617d..39d22e1 100644
--- a/src/com/android/managedprovisioning/provisioning/TransitionAnimationHelper.java
+++ b/src/com/android/managedprovisioning/provisioning/TransitionAnimationHelper.java
@@ -47,6 +47,8 @@
 
     interface TransitionAnimationCallback {
         void onAllTransitionsShown();
+
+        void onTransitionStart(int screenIndex, AnimatedVectorDrawable animatedVectorDrawable);
     }
 
     @VisibleForTesting
@@ -90,7 +92,9 @@
 
     TransitionAnimationHelper(@ProvisioningMode int provisioningMode,
             boolean adminCanGrantSensorsPermissions,
-            AnimationComponents animationComponents, TransitionAnimationCallback callback) {
+            AnimationComponents animationComponents,
+            TransitionAnimationCallback callback,
+            int currentTransitionIndex) {
         mAnimationComponents = checkNotNull(animationComponents);
         mCallback = checkNotNull(callback);
         mProvisioningModeWrapper = getProvisioningModeWrapper(provisioningMode,
@@ -98,6 +102,9 @@
         mCrossFadeHelper = getCrossFadeHelper();
         mShowAnimations = shouldShowAnimations();
 
+        // TODO(b/182824327): Gracefully pause/resume edu screen animations rather than restarting
+        mCurrentTransitionIndex = currentTransitionIndex;
+
         applyContentDescription();
         updateUiValues(mCurrentTransitionIndex);
     }
@@ -166,6 +173,7 @@
         boolean shouldLoop = getTransitionForIndex(mCurrentTransitionIndex).shouldLoop;
         mRepeatingVectorAnimation = new RepeatingVectorAnimation(vectorDrawable, shouldLoop);
         mRepeatingVectorAnimation.start();
+        mCallback.onTransitionStart(mCurrentTransitionIndex, vectorDrawable);
     }
 
     @VisibleForTesting
diff --git a/tests/instrumentation/src/com/android/managedprovisioning/provisioning/ProvisioningManagerTest.java b/tests/instrumentation/src/com/android/managedprovisioning/provisioning/ProvisioningManagerTest.java
index 38dc145..829e54e 100644
--- a/tests/instrumentation/src/com/android/managedprovisioning/provisioning/ProvisioningManagerTest.java
+++ b/tests/instrumentation/src/com/android/managedprovisioning/provisioning/ProvisioningManagerTest.java
@@ -91,7 +91,10 @@
                     msg.getCallback().run();
                     return null;
                 });
-        mManager = new ProvisioningManager(mContext, mUiHandler, mFactory, mAnalyticsTracker,
+        mManager = new ProvisioningManager(
+                mContext,
+                mFactory,
+                mAnalyticsTracker,
                 mTimeLogger);
         when(mFactory.createProvisioningController(mContext, TEST_PARAMS, mManager))
                 .thenReturn(mController);
diff --git a/tests/instrumentation/src/com/android/managedprovisioning/provisioning/ProvisioningViewModelTest.java b/tests/instrumentation/src/com/android/managedprovisioning/provisioning/ProvisioningViewModelTest.java
new file mode 100644
index 0000000..1f12fa1
--- /dev/null
+++ b/tests/instrumentation/src/com/android/managedprovisioning/provisioning/ProvisioningViewModelTest.java
@@ -0,0 +1,43 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.managedprovisioning.provisioning;
+
+
+import static com.google.common.truth.Truth.assertThat;
+
+import androidx.test.filters.SmallTest;
+
+import org.junit.Test;
+
+@SmallTest
+public class ProvisioningViewModelTest {
+
+    private static final int CURRENT_TRANSITION_SCREEN = 4;
+
+    private final ProvisioningViewModel mViewModel = new ProvisioningViewModel();
+
+    @Test
+    public void getCurrentTransitionScreen_defaultsToZero() {
+        assertThat(mViewModel.getCurrentTransitionScreen()).isEqualTo(0);
+    }
+
+    @Test
+    public void setCurrentTransitionScreen_works() {
+        mViewModel.setCurrentTransitionScreen(CURRENT_TRANSITION_SCREEN);
+        assertThat(mViewModel.getCurrentTransitionScreen()).isEqualTo(CURRENT_TRANSITION_SCREEN);
+    }
+}
diff --git a/tests/robotests/src/com/android/managedprovisioning/provisioning/ProvisioningActivityRoboTest.java b/tests/robotests/src/com/android/managedprovisioning/provisioning/ProvisioningActivityRoboTest.java
index 371cd6c..41f406d 100644
--- a/tests/robotests/src/com/android/managedprovisioning/provisioning/ProvisioningActivityRoboTest.java
+++ b/tests/robotests/src/com/android/managedprovisioning/provisioning/ProvisioningActivityRoboTest.java
@@ -258,7 +258,7 @@
         final ProvisioningActivity activity =
                 Robolectric.buildActivity(ProvisioningActivity.class, PROFILE_OWNER_INTENT)
                         .setup().get();
-        activity.mProvisioningManager = mMockProvisioningManager;
+        activity.setProvisioningManager(mMockProvisioningManager);
 
         activity.onBackPressed();
         final Fragment dialog =
@@ -289,7 +289,7 @@
         final ProvisioningActivity activity =
                 Robolectric.buildActivity(ProvisioningActivity.class, PROFILE_OWNER_INTENT)
                         .setup().get();
-        activity.mProvisioningManager = mMockProvisioningManager;
+        activity.setProvisioningManager(mMockProvisioningManager);
 
         activity.onBackPressed();
         final Fragment dialog =