Replace the main user look up with UserManager#getMainUser

Test: atest SystemUITests:com.android.systemui.notetask.shortcut.LaunchNoteTaskActivityTest
Fix: 276904770
Change-Id: I564b4dd69d4c61193268a0cd1d88b459e8787af5
diff --git a/packages/SystemUI/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivity.kt b/packages/SystemUI/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivity.kt
index 44855fb..0f38d32 100644
--- a/packages/SystemUI/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivity.kt
+++ b/packages/SystemUI/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivity.kt
@@ -18,9 +18,11 @@
 
 import android.content.Context
 import android.content.Intent
-import android.content.pm.UserInfo
+import android.os.Build
 import android.os.Bundle
+import android.os.UserHandle
 import android.os.UserManager
+import android.util.Log
 import androidx.activity.ComponentActivity
 import com.android.systemui.notetask.NoteTaskController
 import com.android.systemui.notetask.NoteTaskEntryPoint
@@ -63,9 +65,13 @@
         //        | Bubble#showOrHideAppBubble |   <--------------
         //        |      (with WP user ID)     |
         //         ----------------------------
-        val mainUser: UserInfo? = userTracker.userProfiles.firstOrNull { it.isMain }
-        if (userManager.isManagedProfile && mainUser != null) {
-            controller.startNoteTaskProxyActivityForUser(mainUser.userHandle)
+        val mainUser: UserHandle? = userManager.mainUser
+        if (userManager.isManagedProfile) {
+            if (mainUser == null) {
+                logDebug { "Can't find the main user. Skipping the notes app launch." }
+            } else {
+                controller.startNoteTaskProxyActivityForUser(mainUser)
+            }
         } else {
             controller.showNoteTask(entryPoint = NoteTaskEntryPoint.WIDGET_PICKER_SHORTCUT)
         }
@@ -83,3 +89,8 @@
         }
     }
 }
+
+/** [Log.println] a [Log.DEBUG] message, only when [Build.IS_DEBUGGABLE]. */
+private inline fun Any.logDebug(message: () -> String) {
+    if (Build.IS_DEBUGGABLE) Log.d(this::class.java.simpleName.orEmpty(), message())
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivityTest.kt b/packages/SystemUI/tests/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivityTest.kt
index c96853d..a0c376f 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivityTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/notetask/shortcut/LaunchNoteTaskActivityTest.kt
@@ -30,6 +30,7 @@
 import com.android.systemui.notetask.NoteTaskController
 import com.android.systemui.notetask.NoteTaskEntryPoint
 import com.android.systemui.settings.FakeUserTracker
+import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.eq
 import com.android.systemui.util.mockito.whenever
 import org.junit.After
@@ -38,6 +39,7 @@
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
+import org.mockito.Mockito.never
 import org.mockito.MockitoAnnotations
 
 @RunWith(AndroidTestingRunner::class)
@@ -86,15 +88,28 @@
 
     @Test
     fun startActivityOnWorkProfileUser_shouldLaunchProxyActivity() {
+        val mainUserHandle: UserHandle = mainUser.userHandle
         userTracker.set(listOf(mainUser, workProfileUser), selectedUserIndex = 1)
         whenever(userManager.isManagedProfile).thenReturn(true)
+        whenever(userManager.mainUser).thenReturn(mainUserHandle)
 
         activityRule.launchActivity(/* startIntent= */ null)
 
-        val mainUserHandle: UserHandle = mainUser.userHandle
         verify(noteTaskController).startNoteTaskProxyActivityForUser(eq(mainUserHandle))
     }
 
+    @Test
+    fun startActivityOnWorkProfileUser_noMainUser_shouldNotLaunch() {
+        userTracker.set(listOf(mainUser, workProfileUser), selectedUserIndex = 1)
+        whenever(userManager.isManagedProfile).thenReturn(true)
+        whenever(userManager.mainUser).thenReturn(null)
+
+        activityRule.launchActivity(/* startIntent= */ null)
+
+        verify(noteTaskController, never()).showNoteTask(any())
+        verify(noteTaskController, never()).startNoteTaskProxyActivityForUser(any())
+    }
+
     private companion object {
         val mainUser = UserInfo(/* id= */ 0, /* name= */ "primary", /* flags= */ UserInfo.FLAG_MAIN)
         val workProfileUser =