Swap orders of GAIA signin and lock setup during Private space setup

Bug: 312892127
Test: Verified manually Lock setup is before GAIA login
Change-Id: I05b7c61306e4988a59511988f88e26cdefbed194
diff --git a/res/navigation/privatespace_main_context_nav.xml b/res/navigation/privatespace_main_context_nav.xml
index ffc63ec..7849458 100644
--- a/res/navigation/privatespace_main_context_nav.xml
+++ b/res/navigation/privatespace_main_context_nav.xml
@@ -33,8 +33,8 @@
             android:id="@+id/action_advance_profile_error"
             app:destination="@id/ps_profile_error_fragment"/>
         <action
-            android:id="@+id/action_advance_login_error"
-            app:destination="@id/ps_account_error_fragment"/>
+            android:id="@+id/action_set_lock_fragment"
+            app:destination="@id/ps_profile_lock_fragment"/>
     </fragment>
     <fragment android:id="@+id/ps_profile_error_fragment"
               android:name="com.android.settings.privatespace.PrivateProfileCreationError"
@@ -55,9 +55,11 @@
     </fragment>
     <fragment android:id="@+id/ps_profile_lock_fragment"
               android:name="com.android.settings.privatespace.PrivateSpaceSetLockFragment"
-              android:label="fragment_ps_lock"/>
+              android:label="fragment_ps_lock">
+        <action
+            android:id="@+id/action_advance_login_error"
+            app:destination="@id/ps_account_error_fragment"/>
+    </fragment>
     <action android:id="@+id/action_success_fragment"
             app:destination="@id/ps_profile_success_fragment"/>
-    <action android:id="@+id/action_set_lock_fragment"
-            app:destination="@id/ps_profile_lock_fragment"/>
 </navigation>
diff --git a/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java b/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java
index f9c4c4e..77d57f7 100644
--- a/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java
+++ b/src/com/android/settings/privatespace/AutoAdvanceSetupFragment.java
@@ -16,21 +16,15 @@
 
 package com.android.settings.privatespace;
 
-import static com.android.settings.privatespace.PrivateSpaceSetupActivity.ACCOUNT_LOGIN_ACTION;
-import static com.android.settings.privatespace.PrivateSpaceSetupActivity.EXTRA_ACTION_TYPE;
-
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.AnimatorSet;
 import android.animation.ObjectAnimator;
 import android.animation.ValueAnimator;
-import android.annotation.SuppressLint;
 import android.app.settings.SettingsEnums;
-import android.content.Intent;
 import android.os.Bundle;
 import android.os.Handler;
 import android.os.Looper;
-import android.os.UserHandle;
 import android.util.Log;
 import android.util.Pair;
 import android.view.LayoutInflater;
@@ -78,27 +72,20 @@
                         if (++mScreenTitleIndex < HEADER_IMAGE_PAIRS.size()) {
                             startFadeOutAnimation();
                             mHandler.postDelayed(mUpdateScreenResources, DELAY_BETWEEN_SCREENS);
+                        } else if (PrivateSpaceMaintainer.getInstance(getActivity())
+                                .doesPrivateSpaceExist()) {
+                            mMetricsFeatureProvider.action(
+                                    getContext(),
+                                    SettingsEnums.ACTION_PRIVATE_SPACE_SETUP_SPACE_CREATED,
+                                    true);
+                            NavHostFragment.findNavController(AutoAdvanceSetupFragment.this)
+                                    .navigate(R.id.action_set_lock_fragment);
                         } else {
-                            PrivateSpaceMaintainer privateSpaceMaintainer =
-                                    PrivateSpaceMaintainer.getInstance(getActivity());
-                            UserHandle userHandle;
-                            if (privateSpaceMaintainer.doesPrivateSpaceExist()
-                                    && (userHandle =
-                                                    privateSpaceMaintainer
-                                                            .getPrivateProfileHandle())
-                                            != null) {
-                                mMetricsFeatureProvider.action(
-                                        getContext(),
-                                        SettingsEnums.ACTION_PRIVATE_SPACE_SETUP_SPACE_CREATED,
-                                        true);
-                                startActivityInPrivateUser(userHandle);
-                            } else {
-                                mMetricsFeatureProvider.action(
-                                        getContext(),
-                                        SettingsEnums.ACTION_PRIVATE_SPACE_SETUP_SPACE_CREATED,
-                                        false);
-                                showPrivateSpaceErrorScreen();
-                            }
+                            mMetricsFeatureProvider.action(
+                                    getContext(),
+                                    SettingsEnums.ACTION_PRIVATE_SPACE_SETUP_SPACE_CREATED,
+                                    false);
+                            showPrivateSpaceErrorScreen();
                         }
                     }
                 }
@@ -161,14 +148,6 @@
         return SettingsEnums.PRIVATE_SPACE_SETUP_SPACE_CREATION;
     }
 
-    @SuppressLint("MissingPermission")
-    private void startActivityInPrivateUser(UserHandle userHandle) {
-        /* Start new activity in private profile which is needed to set private profile lock */
-        Intent intent = new Intent(getContext(), PrivateProfileContextHelperActivity.class);
-        intent.putExtra(EXTRA_ACTION_TYPE, ACCOUNT_LOGIN_ACTION);
-        getActivity().startActivityForResultAsUser(intent, ACCOUNT_LOGIN_ACTION, userHandle);
-    }
-
     private void showPrivateSpaceErrorScreen() {
         NavHostFragment.findNavController(AutoAdvanceSetupFragment.this)
                 .navigate(R.id.action_advance_profile_error);
diff --git a/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java b/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java
index 56f48a8..dd4bc7f 100644
--- a/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java
+++ b/src/com/android/settings/privatespace/PrivateSpaceSetLockFragment.java
@@ -16,21 +16,21 @@
 
 package com.android.settings.privatespace;
 
+import static com.android.settings.privatespace.PrivateSpaceSetupActivity.ACCOUNT_LOGIN_ACTION;
 import static com.android.settings.privatespace.PrivateSpaceSetupActivity.EXTRA_ACTION_TYPE;
 import static com.android.settings.privatespace.PrivateSpaceSetupActivity.SET_LOCK_ACTION;
 
-import android.annotation.SuppressLint;
 import android.app.settings.SettingsEnums;
 import android.content.Intent;
 import android.os.Bundle;
 import android.os.UserHandle;
+import android.util.Log;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.ViewGroup;
 
 import androidx.activity.OnBackPressedCallback;
 import androidx.annotation.Nullable;
-import androidx.navigation.fragment.NavHostFragment;
 
 import com.android.settings.R;
 import com.android.settings.core.InstrumentedFragment;
@@ -44,6 +44,7 @@
  * separate private profile lock.
  */
 public class PrivateSpaceSetLockFragment extends InstrumentedFragment {
+    private static final String TAG = "PrivateSpaceSetLockFrag";
 
     @Override
     public View onCreateView(
@@ -96,8 +97,7 @@
             mMetricsFeatureProvider.action(
                     getContext(), SettingsEnums.ACTION_PRIVATE_SPACE_SETUP_USE_SCREEN_LOCK);
             // Simply Use default screen lock. No need to handle
-            NavHostFragment.findNavController(PrivateSpaceSetLockFragment.this)
-                    .navigate(R.id.action_success_fragment);
+            launchActivityForAction(ACCOUNT_LOGIN_ACTION);
         };
     }
 
@@ -105,21 +105,19 @@
         return v -> {
             mMetricsFeatureProvider.action(
                     getContext(), SettingsEnums.ACTION_PRIVATE_SPACE_SETUP_NEW_LOCK);
-            createPrivateSpaceLock();
+            launchActivityForAction(SET_LOCK_ACTION);
         };
     }
 
-    @SuppressLint("MissingPermission")
-    private void createPrivateSpaceLock() {
-        PrivateSpaceMaintainer privateSpaceMaintainer = PrivateSpaceMaintainer
-                .getInstance(getActivity());
-        UserHandle userHandle;
-        if (privateSpaceMaintainer.doesPrivateSpaceExist() && (userHandle =
-                privateSpaceMaintainer.getPrivateProfileHandle()) != null) {
+    private void launchActivityForAction(int action) {
+        UserHandle userHandle =
+                PrivateSpaceMaintainer.getInstance(getActivity()).getPrivateProfileHandle();
+        if (userHandle != null) {
             Intent intent = new Intent(getContext(), PrivateProfileContextHelperActivity.class);
-            intent.putExtra(EXTRA_ACTION_TYPE, SET_LOCK_ACTION);
-            getActivity().startActivityForResultAsUser(intent, SET_LOCK_ACTION,
-                    userHandle);
+            intent.putExtra(EXTRA_ACTION_TYPE, action);
+            getActivity().startActivityForResultAsUser(intent, action, userHandle);
+        } else {
+            Log.w(TAG, "Private profile user handle is null");
         }
     }
 }
diff --git a/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java b/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java
index d6e1a3f..21224ff 100644
--- a/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java
+++ b/src/com/android/settings/privatespace/PrivateSpaceSetupActivity.java
@@ -19,6 +19,8 @@
 import android.app.settings.SettingsEnums;
 import android.content.Intent;
 import android.os.Bundle;
+import android.os.UserHandle;
+import android.util.Log;
 
 import androidx.annotation.Nullable;
 import androidx.fragment.app.FragmentActivity;
@@ -33,6 +35,7 @@
 
 /** Activity class that helps in setting up of private space */
 public class PrivateSpaceSetupActivity extends FragmentActivity {
+    private static final String TAG = "PrivateSpaceSetupAct";
     public static final int SET_LOCK_ACTION = 1;
     public static final int ACCOUNT_LOGIN_ACTION = 2;
     public static final String EXTRA_ACTION_TYPE = "action_type";
@@ -57,12 +60,21 @@
     @Override
     protected void onActivityResult(int requestCode, int resultCode, @Nullable Intent data) {
         if (requestCode == SET_LOCK_ACTION && resultCode == RESULT_OK) {
-            mNavHostFragment.getNavController().navigate(R.id.action_success_fragment);
+            /* Start new activity in private profile to add an account to private profile */
+            UserHandle userHandle =
+                    PrivateSpaceMaintainer.getInstance(this).getPrivateProfileHandle();
+            if (userHandle != null) {
+                Intent intent = new Intent(this, PrivateProfileContextHelperActivity.class);
+                intent.putExtra(EXTRA_ACTION_TYPE, ACCOUNT_LOGIN_ACTION);
+                startActivityForResultAsUser(intent, ACCOUNT_LOGIN_ACTION, userHandle);
+            } else {
+                Log.w(TAG, "Private profile user handle is null");
+            }
         } else if (requestCode == ACCOUNT_LOGIN_ACTION) {
             if (resultCode == RESULT_OK) {
                 mMetricsFeatureProvider.action(
                         this, SettingsEnums.ACTION_PRIVATE_SPACE_SETUP_ACCOUNT_LOGIN_SUCCESS, true);
-                mNavHostFragment.getNavController().navigate(R.id.action_set_lock_fragment);
+                mNavHostFragment.getNavController().navigate(R.id.action_success_fragment);
             } else {
                 mMetricsFeatureProvider.action(
                         this,