Update progress LiveData when setup completed

Test: atest DeviceLockControllerRoboTests
Fix: 279939367
Change-Id: I59b410c740645a0340882a7b14c1b76efd8b7804
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/activities/DevicePoliciesFragment.java b/DeviceLockController/src/com/android/devicelockcontroller/activities/DevicePoliciesFragment.java
index 69517e2..65cc925 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/activities/DevicePoliciesFragment.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/activities/DevicePoliciesFragment.java
@@ -35,6 +35,7 @@
 
 import com.android.devicelockcontroller.R;
 import com.android.devicelockcontroller.policy.PolicyObjectsInterface;
+import com.android.devicelockcontroller.policy.SetupController;
 import com.android.devicelockcontroller.util.LogUtil;
 
 import com.google.common.util.concurrent.FutureCallback;
@@ -89,18 +90,21 @@
         checkNotNull(footerTextView);
         viewModel.mFooterTextIdLiveData.observe(getViewLifecycleOwner(), footerTextView::setText);
 
+        SetupController setupController =
+                ((PolicyObjectsInterface) getActivity().getApplicationContext())
+                        .getSetupController();
+
         ProvisioningProgressViewModel provisioningProgressViewModel =
                 new ViewModelProvider(requireActivity()).get(ProvisioningProgressViewModel.class);
+        MutableLiveData<ProvisioningProgress> progressLiveData =
+                provisioningProgressViewModel.getProvisioningProgressMutableLiveData();
         Button button = view.findViewById(R.id.button_next);
         checkNotNull(button);
         button.setOnClickListener(
                 v -> {
-                    MutableLiveData<ProvisioningProgress> progressLiveData =
-                            provisioningProgressViewModel.getProvisioningProgressMutableLiveData();
                     progressLiveData.setValue(ProvisioningProgress.GETTING_DEVICE_READY);
                     Futures.addCallback(
-                            ((PolicyObjectsInterface) getActivity().getApplicationContext())
-                                    .getSetupController().startSetupFlow(getViewLifecycleOwner()),
+                            setupController.startSetupFlow(getViewLifecycleOwner()),
                             new FutureCallback<>() {
                                 @Override
                                 public void onSuccess(Void result) {
@@ -116,5 +120,19 @@
                                 }
                             }, MoreExecutors.directExecutor());
                 });
+
+        setupController.addListener(new SetupController.SetupUpdatesCallbacks() {
+            @Override
+            public void setupFailed(int reason) {
+                //TODO(b/279969959): Show the failure UI if we have one.
+                LogUtil.e(TAG, "Failed to finish setup flow!");
+            }
+
+            @Override
+            public void setupCompleted() {
+                LogUtil.i(TAG, "Successfully finished setup flow!");
+                progressLiveData.postValue(ProvisioningProgress.OPENING_KIOSK_APP);
+            }
+        });
     }
 }
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupController.java b/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupController.java
index 94c12bd..8ec42ff 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupController.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupController.java
@@ -64,12 +64,6 @@
     /** Triggers the setup flow process. */
     ListenableFuture<Void> startSetupFlow(LifecycleOwner owner);
 
-    /**
-     * Finishes the setup flow process. Triggers the appropriate actions based on whether setup was
-     * successful.
-     */
-    ListenableFuture<Void> finishSetup();
-
     /** Callback interface for updates on setup tasks */
     interface SetupUpdatesCallbacks {
 
diff --git a/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupControllerImpl.java b/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupControllerImpl.java
index 6efbedd..5324235 100644
--- a/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupControllerImpl.java
+++ b/DeviceLockController/src/com/android/devicelockcontroller/policy/SetupControllerImpl.java
@@ -39,6 +39,8 @@
 import static com.android.devicelockcontroller.policy.SetupController.SetupUpdatesCallbacks.FailureType.INSTALL_FAILED;
 import static com.android.devicelockcontroller.policy.SetupController.SetupUpdatesCallbacks.FailureType.VERIFICATION_FAILED;
 
+import static com.google.common.util.concurrent.Futures.transform;
+
 import android.content.Context;
 import android.content.pm.PackageManager.NameNotFoundException;
 import android.os.Build;
@@ -62,6 +64,7 @@
 import com.android.devicelockcontroller.setup.SetupParametersClient;
 import com.android.devicelockcontroller.util.LogUtil;
 
+import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.MoreExecutors;
@@ -157,7 +160,7 @@
     private ListenableFuture<Boolean> isKioskAppPreInstalled() {
         return Build.isDebuggable()
                 ? Futures.immediateFuture(false)
-                : Futures.transform(SetupParametersClient.getInstance().getKioskPackage(),
+                : transform(SetupParametersClient.getInstance().getKioskPackage(),
                         packageName -> {
                             try {
                                 mContext.getPackageManager().getPackageInfo(
@@ -336,45 +339,43 @@
                     if (areAllTasksSucceeded(workInfo)) {
                         setupFlowTaskSuccessCallbackHandler();
                     } else if (isAtLeastOneTaskFailedOrCancelled(workInfo)) {
-                        setupFlowTaskFailureCallbackHandler(
-                                getTaskFailureType(workInfo));
+                        setupFlowTaskFailureCallbackHandler(getTaskFailureType(workInfo));
                     }
                 });
     }
 
+    @VisibleForTesting
+    void finishSetup() {
+        if (mCurrentSetupState == SetupStatus.SETUP_FINISHED) {
+            Futures.addCallback(mPolicyController.launchActivityInLockedMode(),
+                    new FutureCallback<>() {
+                        @Override
+                        public void onSuccess(Boolean isLaunched) {
+                            if (!isLaunched) {
+                                onFailure(new IllegalStateException());
+                            }
+                            LogUtil.i(TAG, "Launched kiosk activity in lock task mode");
+                        }
 
-    @Override
-    public ListenableFuture<Void> finishSetup() {
-        ListenableFuture<Void> setupFailureTask = Futures.immediateVoidFuture();
-        try {
-            if (mCurrentSetupState == SetupStatus.SETUP_FINISHED) {
-                return Futures.transformAsync(
-                        mStateController.setNextStateForEvent(DeviceEvent.SETUP_COMPLETE),
-                        input -> Futures.transform(
-                                mPolicyController.launchActivityInLockedMode(),
-                                result -> {
-                                    if (!result) {
-                                        throw new IllegalStateException(
-                                                "Failed to launch kiosk Activity in lock mode!");
-                                    }
-                                    return null;
-                                }, MoreExecutors.directExecutor()),
-                        MoreExecutors.directExecutor());
-            }
-            setupFailureTask = mStateController.setNextStateForEvent(DeviceEvent.SETUP_FAILURE);
-        } catch (StateTransitionException e) {
-            LogUtil.e(TAG, "State transition failed!", e);
+                        @Override
+                        public void onFailure(Throwable t) {
+                            LogUtil.e(TAG, "Failed to launch kiosk activity in lock task mode!", t);
+                        }
+                    }, MoreExecutors.directExecutor());
+        } else {
+            Futures.addCallback(SetupParametersClient.getInstance().isProvisionMandatory(),
+                    new FutureCallback<>() {
+                        @Override
+                        public void onSuccess(Boolean isMandatory) {
+                            if (isMandatory) mPolicyController.wipeData();
+                        }
+
+                        @Override
+                        public void onFailure(Throwable t) {
+                            LogUtil.e(TAG, "Failed to know if Provision is mandatory", t);
+                        }
+                    }, MoreExecutors.directExecutor());
         }
-        LogUtil.e(TAG, "Setup failed!");
-        ListenableFuture<Boolean> provisionMandatoryTask =
-                SetupParametersClient.getInstance().isProvisionMandatory();
-        return Futures.whenAllComplete(setupFailureTask, provisionMandatoryTask).call(
-                () -> {
-                    if (Futures.getDone(provisionMandatoryTask)) {
-                        mPolicyController.wipeData();
-                    }
-                    return null;
-                }, MoreExecutors.directExecutor());
     }
 
     private void setupFlowTaskSuccessCallbackHandler() {
@@ -423,6 +424,7 @@
                 }
             }
         }
+        finishSetup();
     }
 
     @VisibleForTesting
diff --git a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/SetupControllerImplTest.java b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/SetupControllerImplTest.java
index bcc3627..c2fdb3b 100644
--- a/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/SetupControllerImplTest.java
+++ b/DeviceLockController/tests/robolectric/src/com/android/devicelockcontroller/policy/SetupControllerImplTest.java
@@ -64,6 +64,7 @@
 import androidx.work.WorkerParameters;
 import androidx.work.testing.WorkManagerTestInitHelper;
 
+import com.android.devicelockcontroller.TestDeviceLockControllerApplication;
 import com.android.devicelockcontroller.policy.DeviceStateController.DeviceEvent;
 import com.android.devicelockcontroller.policy.DeviceStateController.DeviceState;
 import com.android.devicelockcontroller.setup.SetupParametersClient;
@@ -108,35 +109,36 @@
     @Rule
     public final MockitoRule mMocks = MockitoJUnit.rule();
 
-    @Mock
     private DeviceStateController mMockStateController;
-    @Mock
     private DevicePolicyController mMockPolicyController;
     @Mock
     private SetupController.SetupUpdatesCallbacks mMockCbs;
     @Mock
     private LifecycleOwner mMockLifecycleOwner;
-
-    private Context mContext;
+    private TestDeviceLockControllerApplication mTestApplication;
     private String mFileLocation;
     private SetupParametersClient mSetupParametersClient;
     private TestWorkFactory mTestWorkFactory;
 
     @Before
     public void setUp() {
-        mContext = ApplicationProvider.getApplicationContext();
-        Shadows.shadowOf((Application) mContext).setComponentNameAndServiceForBindService(
-                new ComponentName(mContext, SetupParametersService.class),
+        mTestApplication = ApplicationProvider.getApplicationContext();
+        mMockStateController = mTestApplication.getMockStateController();
+        mMockPolicyController = mTestApplication.getMockPolicyController();
+        when(mMockPolicyController.launchActivityInLockedMode()).thenReturn(
+                Futures.immediateFuture(true));
+        Shadows.shadowOf((Application) mTestApplication).setComponentNameAndServiceForBindService(
+                new ComponentName(mTestApplication, SetupParametersService.class),
                 Robolectric.setupService(SetupParametersService.class).onBind(null));
         mSetupParametersClient = SetupParametersClient.getInstance(
-                mContext, TestingExecutors.sameThreadScheduledExecutor());
-        mFileLocation = mContext.getFilesDir() + "/TEST_FILE_NAME";
+                mTestApplication, TestingExecutors.sameThreadScheduledExecutor());
+        mFileLocation = mTestApplication.getFilesDir() + "/TEST_FILE_NAME";
         createTestFile(mFileLocation);
 
         mTestWorkFactory = new TestWorkFactory();
         Configuration config =
                 new Configuration.Builder().setWorkerFactory(mTestWorkFactory).build();
-        WorkManagerTestInitHelper.initializeTestWorkManager(mContext, config);
+        WorkManagerTestInitHelper.initializeTestWorkManager(mTestApplication, config);
     }
 
     @After
@@ -150,56 +152,28 @@
         b.putString(EXTRA_KIOSK_SETUP_ACTIVITY, TEST_SETUP_ACTIVITY);
         createParameters(b);
         when(mMockStateController.getState()).thenReturn(DeviceState.KIOSK_SETUP);
-        when(mMockStateController.setNextStateForEvent(DeviceEvent.SETUP_COMPLETE)).thenReturn(
-                Futures.immediateVoidFuture());
         SetupControllerImpl setupController =
                 new SetupControllerImpl(
-                        mContext, mMockStateController, mMockPolicyController);
+                        mTestApplication, mMockStateController, mMockPolicyController);
         assertThat(setupController.getSetupState()).isEqualTo(
                 SetupController.SetupStatus.SETUP_FINISHED);
         setupController.finishSetup();
-        verify(mMockStateController).setNextStateForEvent(eq(DeviceEvent.SETUP_COMPLETE));
         verify(mMockPolicyController).launchActivityInLockedMode();
         verify(mMockPolicyController, never()).wipeData();
     }
 
     @Test
-    public void testInitialState_SetupFinishedException() throws Exception {
-        Bundle b = new Bundle();
-        b.putString(EXTRA_KIOSK_SETUP_ACTIVITY, TEST_SETUP_ACTIVITY);
-        createParameters(b);
-        when(mMockStateController.getState()).thenReturn(DeviceState.KIOSK_SETUP);
-        SetupControllerImpl setupController =
-                new SetupControllerImpl(
-                        mContext, mMockStateController, mMockPolicyController);
-        assertThat(setupController.getSetupState()).isEqualTo(
-                SetupController.SetupStatus.SETUP_FINISHED);
-
-        doThrow(
-                new StateTransitionException(
-                        DeviceEvent.PROVISIONING_SUCCESS, DeviceState.UNPROVISIONED))
-                .when(mMockStateController)
-                .setNextStateForEvent(anyInt());
-        setupController.finishSetup();
-        verify(mMockPolicyController, never()).launchActivityInLockedMode();
-    }
-
-    @Test
-    public void testInitialState_mandatoryProvisioning_SetupFailed()
-            throws StateTransitionException {
+    public void testInitialState_mandatoryProvisioning_SetupFailed() {
         Bundle bundle = new Bundle();
         bundle.putBoolean(EXTRA_MANDATORY_PROVISION, true);
         createParameters(bundle);
         when(mMockStateController.getState()).thenReturn(DeviceState.SETUP_FAILED);
-        when(mMockStateController.setNextStateForEvent(DeviceEvent.SETUP_FAILURE)).thenReturn(
-                Futures.immediateVoidFuture());
         SetupControllerImpl setupController =
                 new SetupControllerImpl(
-                        mContext, mMockStateController, mMockPolicyController);
+                        mTestApplication, mMockStateController, mMockPolicyController);
         setupController.finishSetup();
         assertThat(setupController.getSetupState()).isEqualTo(
                 SetupController.SetupStatus.SETUP_FAILED);
-        verify(mMockStateController).setNextStateForEvent(eq(DeviceEvent.SETUP_FAILURE));
         verify(mMockPolicyController, never()).launchActivityInLockedMode();
         verify(mMockPolicyController).wipeData();
     }
@@ -219,8 +193,9 @@
         SetupControllerImpl setupController = createSetupControllerImpl(mMockCbs);
 
         // WHEN finish kiosk app setup
-        setupController.installKioskAppFromURL(WorkManager.getInstance(mContext),
-                mMockLifecycleOwner);
+        Futures.getUnchecked(
+                setupController.installKioskAppFromURL(WorkManager.getInstance(mTestApplication),
+                        mMockLifecycleOwner));
 
         // THEN setup succeeds
         verify(mMockStateController).setNextStateForEvent(eq(DeviceEvent.SETUP_SUCCESS));
@@ -237,8 +212,9 @@
         SetupControllerImpl setupController = createSetupControllerImpl(mMockCbs);
 
         // WHEN finish kiosk app setup
-        setupController.installKioskAppFromURL(WorkManager.getInstance(mContext),
-                mMockLifecycleOwner);
+        Futures.getUnchecked(
+                setupController.installKioskAppFromURL(WorkManager.getInstance(mTestApplication),
+                        mMockLifecycleOwner));
 
         // THEN verify task will fail
         verify(mMockStateController).setNextStateForEvent(eq(DeviceEvent.SETUP_FAILURE));
@@ -253,15 +229,16 @@
         final SetupControllerImpl setupController = createSetupControllerImpl(mMockCbs);
 
         // WHEN install kiosk app for secondary user
-        setupController.installKioskAppForSecondaryUser(WorkManager.getInstance(mContext),
-                mMockLifecycleOwner);
+        Futures.getUnchecked(setupController.installKioskAppForSecondaryUser(
+                WorkManager.getInstance(mTestApplication),
+                mMockLifecycleOwner));
 
         verify(mMockStateController).setNextStateForEvent(eq(DeviceEvent.SETUP_SUCCESS));
         verify(mMockCbs).setupCompleted();
     }
 
     @Test
-    public void installKioskAppForSecondaryUser_kioskAppNotInstalled_oneTaskFaied()
+    public void installKioskAppForSecondaryUser_kioskAppNotInstalled_oneTaskFails()
             throws StateTransitionException {
         // GIVEN verify install task is failed due to no installed package info
         whenVerifyInstallTaskFailed(ERROR_CODE_NO_PACKAGE_INFO);
@@ -270,8 +247,9 @@
         SetupControllerImpl setupController = createSetupControllerImpl(mMockCbs);
 
         // WHEN install kiosk app for secondary user
-        setupController.installKioskAppForSecondaryUser(WorkManager.getInstance(mContext),
-                mMockLifecycleOwner);
+        Futures.getUnchecked(setupController.installKioskAppForSecondaryUser(
+                WorkManager.getInstance(mTestApplication),
+                mMockLifecycleOwner));
 
         // THEN verify task will fail
         verify(mMockStateController).setNextStateForEvent(eq(DeviceEvent.SETUP_FAILURE));
@@ -283,7 +261,7 @@
         when(mMockStateController.getState()).thenReturn(DeviceState.SETUP_IN_PROGRESS);
         SetupControllerImpl setupController =
                 new SetupControllerImpl(
-                        mContext, mMockStateController, mMockPolicyController);
+                        mTestApplication, mMockStateController, mMockPolicyController);
         assertThat(setupController.getSetupState()).isEqualTo(
                 SetupController.SetupStatus.SETUP_NOT_STARTED);
     }
@@ -331,6 +309,39 @@
     }
 
     @Test
+    public void testSetupUpdatesCallbacks_stateTransitionException()
+            throws StateTransitionException {
+        AtomicBoolean result = new AtomicBoolean(true);
+        AtomicInteger reason = new AtomicInteger(-1);
+        SetupController.SetupUpdatesCallbacks callbacks =
+                new SetupController.SetupUpdatesCallbacks() {
+                    @Override
+                    public void setupCompleted() {
+                    }
+
+                    @Override
+                    public void setupFailed(int failReason) {
+                        result.set(false);
+                        reason.set(failReason);
+                    }
+                };
+        doThrow(
+                new StateTransitionException(
+                        DeviceEvent.PROVISIONING_SUCCESS, DeviceState.UNPROVISIONED))
+                .when(mMockStateController)
+                .setNextStateForEvent(anyInt());
+
+        SetupControllerImpl setupController = createSetupControllerImpl(callbacks);
+
+        setupController.setupFlowTaskCallbackHandler(/* result= */ true, /* failReason= */ -1);
+        assertThat(result.get()).isFalse();
+        assertThat(reason.get()).isEqualTo(
+                SetupController.SetupUpdatesCallbacks.FailureType.SETUP_FAILED);
+        assertThat(setupController.getSetupState()).isEqualTo(
+                SetupController.SetupStatus.SETUP_FAILED);
+    }
+
+    @Test
     public void testSetupUpdatesCallbacks_successCallback() {
         AtomicBoolean result = new AtomicBoolean(false);
         SetupController.SetupUpdatesCallbacks callbacks =
@@ -430,7 +441,8 @@
     private SetupControllerImpl createSetupControllerImpl(
             SetupController.SetupUpdatesCallbacks callbacks) {
         SetupControllerImpl setupController =
-                new SetupControllerImpl(mContext, mMockStateController, mMockPolicyController);
+                new SetupControllerImpl(mTestApplication, mMockStateController,
+                        mMockPolicyController);
         setupController.addListener(callbacks);
         return setupController;
     }