From d1c43eaa15813aa239df737075d16631fbbe9134 Mon Sep 17 00:00:00 2001
From: omarmt <omarmt@google.com>
Date: Fri, 21 Oct 2022 14:39:43 +0000
Subject: Add support for Predictive Back in CreateUserActivity

Test: atest CreateUserActivityTest
Bug: 254840316
Change-Id: I9d00c114ae38e72f11678dab93d3feaf222463d1
---
 packages/SystemUI/Android.bp                       |  1 +
 .../android/systemui/user/CreateUserActivity.java  | 19 +++++++-
 packages/SystemUI/tests/AndroidManifest.xml        |  6 +++
 .../systemui/user/CreateUserActivityTest.kt        | 55 ++++++++++++++++++++++
 4 files changed, 80 insertions(+), 1 deletion(-)
 create mode 100644 packages/SystemUI/tests/src/com/android/systemui/user/CreateUserActivityTest.kt

diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 6d542700fceb..540ca049885c 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -223,6 +223,7 @@ android_library {
         "androidx.test.rules",
         "androidx.test.uiautomator",
         "mockito-target-extended-minus-junit4",
+        "androidx.test.ext.junit",
         "testables",
         "truth-prebuilt",
         "monet",
diff --git a/packages/SystemUI/src/com/android/systemui/user/CreateUserActivity.java b/packages/SystemUI/src/com/android/systemui/user/CreateUserActivity.java
index f01712653e1b..b56c4034936f 100644
--- a/packages/SystemUI/src/com/android/systemui/user/CreateUserActivity.java
+++ b/packages/SystemUI/src/com/android/systemui/user/CreateUserActivity.java
@@ -25,6 +25,8 @@ import android.graphics.drawable.Drawable;
 import android.os.Bundle;
 import android.os.RemoteException;
 import android.util.Log;
+import android.window.OnBackInvokedCallback;
+import android.window.OnBackInvokedDispatcher;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -59,6 +61,7 @@ public class CreateUserActivity extends Activity {
     private final ActivityStarter mActivityStarter;
 
     private Dialog mSetupUserDialog;
+    private final OnBackInvokedCallback mBackCallback = this::onBackInvoked;
 
     @Inject
     public CreateUserActivity(UserCreator userCreator,
@@ -82,6 +85,10 @@ public class CreateUserActivity extends Activity {
 
         mSetupUserDialog = createDialog();
         mSetupUserDialog.show();
+
+        getOnBackInvokedDispatcher().registerOnBackInvokedCallback(
+                        OnBackInvokedDispatcher.PRIORITY_DEFAULT,
+                        mBackCallback);
     }
 
     @Override
@@ -125,10 +132,20 @@ public class CreateUserActivity extends Activity {
 
     @Override
     public void onBackPressed() {
-        super.onBackPressed();
+        onBackInvoked();
+    }
+
+    private void onBackInvoked() {
         if (mSetupUserDialog != null) {
             mSetupUserDialog.dismiss();
         }
+        finish();
+    }
+
+    @Override
+    protected void onDestroy() {
+        getOnBackInvokedDispatcher().unregisterOnBackInvokedCallback(mBackCallback);
+        super.onDestroy();
     }
 
     private void addUserNow(String userName, Drawable userIcon) {
diff --git a/packages/SystemUI/tests/AndroidManifest.xml b/packages/SystemUI/tests/AndroidManifest.xml
index 1b404a82145b..e51236be2ffb 100644
--- a/packages/SystemUI/tests/AndroidManifest.xml
+++ b/packages/SystemUI/tests/AndroidManifest.xml
@@ -106,6 +106,12 @@
                   android:finishOnCloseSystemDialogs="true"
                   android:excludeFromRecents="true" />
 
+        <activity android:name=".user.CreateUserActivityTest$CreateUserActivityTestable"
+            android:exported="false"
+            android:theme="@style/Theme.SystemUI.Dialog.Alert"
+            android:finishOnCloseSystemDialogs="true"
+            android:excludeFromRecents="true" />
+
         <provider
             android:name="androidx.startup.InitializationProvider"
             tools:replace="android:authorities"
diff --git a/packages/SystemUI/tests/src/com/android/systemui/user/CreateUserActivityTest.kt b/packages/SystemUI/tests/src/com/android/systemui/user/CreateUserActivityTest.kt
new file mode 100644
index 000000000000..51afbcb37442
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/user/CreateUserActivityTest.kt
@@ -0,0 +1,55 @@
+package com.android.systemui.user
+
+import android.app.Dialog
+import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
+import androidx.test.ext.junit.rules.ActivityScenarioRule
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.util.mockito.mock
+import com.android.systemui.util.mockito.nullable
+import com.android.systemui.util.mockito.whenever
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+@TestableLooper.RunWithLooper
+class CreateUserActivityTest : SysuiTestCase() {
+    open class CreateUserActivityTestable :
+        CreateUserActivity(
+            /* userCreator = */ mock(),
+            /* editUserInfoController = */ mock {
+                val dialog: Dialog = mock()
+                whenever(
+                        createDialog(
+                            /* activity = */ nullable(),
+                            /* activityStarter = */ nullable(),
+                            /* oldUserIcon = */ nullable(),
+                            /* defaultUserName = */ nullable(),
+                            /* title = */ nullable(),
+                            /* successCallback = */ nullable(),
+                            /* cancelCallback = */ nullable()
+                        )
+                    )
+                    .thenReturn(dialog)
+            },
+            /* activityManager = */ mock(),
+            /* activityStarter = */ mock(),
+        )
+
+    @get:Rule val activityRule = ActivityScenarioRule(CreateUserActivityTestable::class.java)
+
+    @Test
+    fun onBackPressed_finishActivity() {
+        activityRule.scenario.onActivity { activity ->
+            assertThat(activity.isFinishing).isFalse()
+
+            activity.onBackPressed()
+
+            assertThat(activity.isFinishing).isTrue()
+        }
+    }
+}
-- 
cgit v1.2.3-59-g8ed1b