Fix getTaskId when BubbleTaskViewListener is enabled.

When Flags.enableBubbleTaskViewListener() is enabled, the task ID is
stored in BubbleTaskViewListener rather than directly in
BubbleExpandedView's mTaskId field. This caused
BubbleExpandedView.getTaskId() to return an uninitialized task ID, which
broke onActivityRestartAttempt in the BubbleController.

The issue occurs because onActivityRestartAttempt in BubbleController
retrieves existing bubbles from BubbleData and calls Bubble.getTaskId()
on each to find matches. Bubble.getTaskId() depends on
BubbleExpandedView.getTaskId(), which was returning an uninitialized
value when the flag was enabled.

Bug: 272102927
Flag: com.android.wm.shell.enable_bubble_task_view_listener
Test: atest WMShellRobolectricTests:BubbleExpandedViewTest
Test: atest WMShellMultivalentTestsOnDevice:BubbleExpandedViewTest
Change-Id: I0648b2865f9d82f33eac3b1e38b242c7405909f3

3 files changed, 80 insertions, 0 deletions

diff --git a/libs/WindowManager/Shell/multivalentTests/Android.bp b/libs/WindowManager/Shell/multivalentTests/Android.bp
index 03076c0940a4..50666911e313 100644
--- a/libs/WindowManager/Shell/multivalentTests/Android.bp
+++ b/libs/WindowManager/Shell/multivalentTests/Android.bp
@@ -51,6 +51,7 @@ android_robolectric_test {
         "androidx.test.ext.junit",
         "mockito-robolectric-prebuilt",
         "mockito-kotlin2",
+        "platform-parametric-runner-lib",
         "truth",
         "flag-junit-base",
         "flag-junit",
@@ -74,6 +75,7 @@ android_test {
         "frameworks-base-testutils",
         "mockito-kotlin2",
         "mockito-target-extended-minus-junit4",
+        "platform-parametric-runner-lib",
         "truth",
         "platform-test-annotations",
         "platform-test-rules",
diff --git a/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleExpandedViewTest.kt b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleExpandedViewTest.kt
new file mode 100644
index 000000000000..bdfaef2c6960
--- /dev/null
+++ b/libs/WindowManager/Shell/multivalentTests/src/com/android/wm/shell/bubbles/BubbleExpandedViewTest.kt
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2025 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.wm.shell.bubbles
+
+import android.content.ComponentName
+import android.content.Context
+import android.platform.test.flag.junit.FlagsParameterization
+import android.platform.test.flag.junit.SetFlagsRule
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.filters.SmallTest
+import com.android.wm.shell.Flags
+import com.android.wm.shell.taskview.TaskView
+import com.google.common.truth.Truth.assertThat
+import com.google.common.util.concurrent.MoreExecutors.directExecutor
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
+
+/** Tests for [BubbleExpandedView] */
+@SmallTest
+@RunWith(ParameterizedAndroidJunit4::class)
+class BubbleExpandedViewTest(flags: FlagsParameterization) {
+
+    @get:Rule
+    val setFlagsRule = SetFlagsRule(flags)
+
+    private val context = ApplicationProvider.getApplicationContext<Context>()
+    private val componentName = ComponentName(context, "TestClass")
+
+    @Test
+    fun getTaskId_onTaskCreated_returnsCorrectTaskId() {
+        val bubbleTaskView = BubbleTaskView(mock<TaskView>(), directExecutor())
+        val expandedView = BubbleExpandedView(context).apply {
+            initialize(
+                mock<BubbleExpandedViewManager>(),
+                mock<BubbleStackView>(),
+                mock<BubblePositioner>(),
+                false /* isOverflow */,
+                bubbleTaskView,
+            )
+            setAnimating(true) // Skips setContentVisibility for testing.
+        }
+
+        bubbleTaskView.listener.onTaskCreated(123, componentName)
+
+        assertThat(expandedView.getTaskId()).isEqualTo(123)
+    }
+
+    companion object {
+        @JvmStatic
+        @Parameters(name = "{0}")
+        fun getParams() = FlagsParameterization.allCombinationsOf(
+            Flags.FLAG_ENABLE_BUBBLE_TASK_VIEW_LISTENER,
+        )
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
index 2c2451cab999..20f24bd62209 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/bubbles/BubbleExpandedView.java
@@ -467,6 +467,11 @@ public class BubbleExpandedView extends LinearLayout {
                         new BubbleTaskViewListener.Callback() {
                             @Override
                             public void onTaskCreated() {
+                                // The taskId is saved to use for removeTask,
+                                // preventing appearance in recent tasks.
+                                mTaskId = ((BubbleTaskViewListener) mCurrentTaskViewListener)
+                                    .getTaskId();
+
                                 setContentVisibility(true);
                             }