Fix back presses dismissing bubbles when IME is up

If bubbles is expanded & shade is not use the virtual display id of the
expanded activity view (if it's valid, otherwise do whats normal).

Test: manual  1) have expanded bubble with IME up
              2) hit back button
              => note that the IME goes away but the bubble remains

              1) have expanded bubble up
              2) pull down shade
              3) hit back button
              => shade goes away

              Do above with gesture nav turned on & do back gesture

Test: atest KeyButtonViewTest
Fixes: 122535136
Change-Id: Iabaace66cfb4d3d80b1a3a3c7b0773deb48fbcaa

8 files changed, 119 insertions, 13 deletions

diff --git a/api/test-current.txt b/api/test-current.txt
index fcad996a67fa..6ed0e1859996 100644
--- a/api/test-current.txt
+++ b/api/test-current.txt
@@ -107,6 +107,7 @@ package android.app {
     ctor public ActivityView(android.content.Context, android.util.AttributeSet);
     ctor public ActivityView(android.content.Context, android.util.AttributeSet, int);
     ctor public ActivityView(android.content.Context, android.util.AttributeSet, int, boolean);
+    method public int getVirtualDisplayId();
     method public void onLayout(boolean, int, int, int, int);
     method public void onLocationChanged();
     method public void performBackPress();
diff --git a/core/java/android/app/ActivityView.java b/core/java/android/app/ActivityView.java
index 8ec5e3a43218..0ccaf62d45b2 100644
--- a/core/java/android/app/ActivityView.java
+++ b/core/java/android/app/ActivityView.java
@@ -19,6 +19,7 @@ package android.app;
 import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_DESTROY_CONTENT_ON_REMOVAL;
 import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY;
 import static android.hardware.display.DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC;
+import static android.view.Display.INVALID_DISPLAY;
 
 import android.annotation.NonNull;
 import android.annotation.TestApi;
@@ -377,6 +378,16 @@ public class ActivityView extends ViewGroup {
     }
 
     /**
+     * @return the display id of the virtual display.
+     */
+    public int getVirtualDisplayId() {
+        if (mVirtualDisplay != null) {
+            return mVirtualDisplay.getDisplay().getDisplayId();
+        }
+        return INVALID_DISPLAY;
+    }
+
+    /**
      * Injects a pair of down/up key events with keycode {@link KeyEvent#KEYCODE_BACK} to the
      * virtual display.
      */
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
index 93effed5958e..13bdb99e7175 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleController.java
@@ -16,6 +16,8 @@
 
 package com.android.systemui.bubbles;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.Display.INVALID_DISPLAY;
 import static android.view.View.INVISIBLE;
 import static android.view.View.VISIBLE;
 import static android.view.ViewGroup.LayoutParams.MATCH_PARENT;
@@ -530,6 +532,21 @@ public class BubbleController implements BubbleExpandedView.OnBubbleBlockedListe
         return mTempRect;
     }
 
+    /**
+     * The display id of the expanded view, if the stack is expanded and not occluded by the
+     * status bar, otherwise returns {@link Display#INVALID_DISPLAY}.
+     */
+    public int getExpandedDisplayId(Context context) {
+        boolean defaultDisplay = context.getDisplay() != null
+                && context.getDisplay().getDisplayId() == DEFAULT_DISPLAY;
+        Bubble b = mStackView.getExpandedBubble();
+        if (defaultDisplay && b != null && isStackExpanded()
+                && !mStatusBarWindowController.getPanelExpanded()) {
+            return b.expandedView.getVirtualDisplayId();
+        }
+        return INVALID_DISPLAY;
+    }
+
     @VisibleForTesting
     BubbleStackView getStackView() {
         return mStackView;
diff --git a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
index 285d4aab4f66..990331417c60 100644
--- a/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
+++ b/packages/SystemUI/src/com/android/systemui/bubbles/BubbleExpandedView.java
@@ -20,6 +20,7 @@ import static android.content.pm.ActivityInfo.DOCUMENT_LAUNCH_ALWAYS;
 import static android.util.StatsLogInternal.BUBBLE_DEVELOPER_ERROR_REPORTED__ERROR__ACTIVITY_INFO_MISSING;
 import static android.util.StatsLogInternal.BUBBLE_DEVELOPER_ERROR_REPORTED__ERROR__ACTIVITY_INFO_NOT_RESIZABLE;
 import static android.util.StatsLogInternal.BUBBLE_DEVELOPER_ERROR_REPORTED__ERROR__DOCUMENT_LAUNCH_NOT_ALWAYS;
+import static android.view.Display.INVALID_DISPLAY;
 
 import android.animation.LayoutTransition;
 import android.animation.ObjectAnimator;
@@ -598,6 +599,16 @@ public class BubbleExpandedView extends LinearLayout implements View.OnClickList
         return mBubbleIntent != null && mActivityView != null;
     }
 
+    /**
+     * @return the display id of the virtual display.
+     */
+    public int getVirtualDisplayId() {
+        if (usingActivityView()) {
+            return mActivityView.getVirtualDisplayId();
+        }
+        return INVALID_DISPLAY;
+    }
+
     private void applyRowState(ExpandableNotificationRow view) {
         view.reset();
         view.setHeadsUp(false);
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java
index 79bf6b3e49f6..8aa50a7a31e4 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/EdgeBackGestureHandler.java
@@ -15,6 +15,8 @@
  */
 package com.android.systemui.statusbar.phone;
 
+import static android.view.Display.INVALID_DISPLAY;
+
 import android.content.Context;
 import android.content.pm.ParceledListSlice;
 import android.content.res.Resources;
@@ -46,7 +48,9 @@ import android.view.ViewConfiguration;
 import android.view.WindowManager;
 import android.view.WindowManagerGlobal;
 
+import com.android.systemui.Dependency;
 import com.android.systemui.R;
+import com.android.systemui.bubbles.BubbleController;
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.shared.system.InputChannelCompat.InputEventReceiver;
 import com.android.systemui.shared.system.QuickStepContract;
@@ -358,6 +362,13 @@ public class EdgeBackGestureHandler implements DisplayListener {
                 0 /* metaState */, KeyCharacterMap.VIRTUAL_KEYBOARD, 0 /* scancode */,
                 KeyEvent.FLAG_FROM_SYSTEM | KeyEvent.FLAG_VIRTUAL_HARD_KEY,
                 InputDevice.SOURCE_KEYBOARD);
+
+        // Bubble controller will give us a valid display id if it should get the back event
+        BubbleController bubbleController = Dependency.get(BubbleController.class);
+        int bubbleDisplayId = bubbleController.getExpandedDisplayId(mContext);
+        if (code == KeyEvent.KEYCODE_BACK && bubbleDisplayId != INVALID_DISPLAY) {
+            ev.setDisplayId(bubbleDisplayId);
+        }
         InputManager.getInstance().injectInputEvent(ev, InputManager.INJECT_INPUT_EVENT_MODE_ASYNC);
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java
index 8d71ab810710..984ab0563fea 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowController.java
@@ -523,12 +523,19 @@ public class StatusBarWindowController implements Callback, Dumpable, Configurat
     }
 
     /**
-     * The bubble is shown in expanded state for the status bar.
+     * Whether the bubble is shown in expanded state for the status bar.
      */
     public boolean getBubbleExpanded() {
         return mCurrentState.bubbleExpanded;
     }
 
+    /**
+     * Whether the status bar panel is expanded or not.
+     */
+    public boolean getPanelExpanded() {
+        return mCurrentState.panelExpanded;
+    }
+
     public void setStateListener(OtherwisedCollapsedListener listener) {
         mListener = listener;
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonView.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonView.java
index 06fc745b0832..9923ec9d8a75 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/KeyButtonView.java
@@ -48,10 +48,12 @@ import android.view.accessibility.AccessibilityEvent;
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.ImageView;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
+import com.android.systemui.bubbles.BubbleController;
 import com.android.systemui.plugins.statusbar.phone.NavBarButtonProvider.ButtonInterface;
 import com.android.systemui.recents.OverviewProxyService;
 import com.android.systemui.shared.system.NavigationBarCompat;
@@ -73,6 +75,7 @@ public class KeyButtonView extends ImageView implements ButtonInterface {
     private final KeyButtonRipple mRipple;
     private final OverviewProxyService mOverviewProxyService;
     private final MetricsLogger mMetricsLogger = Dependency.get(MetricsLogger.class);
+    private final InputManager mInputManager;
 
     private final Runnable mCheckLongPress = new Runnable() {
         public void run() {
@@ -96,6 +99,11 @@ public class KeyButtonView extends ImageView implements ButtonInterface {
     }
 
     public KeyButtonView(Context context, AttributeSet attrs, int defStyle) {
+        this(context, attrs, defStyle, InputManager.getInstance());
+    }
+
+    @VisibleForTesting
+    public KeyButtonView(Context context, AttributeSet attrs, int defStyle, InputManager manager) {
         super(context, attrs);
 
         TypedArray a = context.obtainStyledAttributes(attrs, R.styleable.KeyButtonView,
@@ -117,6 +125,7 @@ public class KeyButtonView extends ImageView implements ButtonInterface {
 
         mRipple = new KeyButtonRipple(context, this);
         mOverviewProxyService = Dependency.get(OverviewProxyService.class);
+        mInputManager = manager;
         setBackground(mRipple);
         forceHasOverlappingRendering(false);
     }
@@ -318,16 +327,23 @@ public class KeyButtonView extends ImageView implements ButtonInterface {
                 0, KeyCharacterMap.VIRTUAL_KEYBOARD, 0,
                 flags | KeyEvent.FLAG_FROM_SYSTEM | KeyEvent.FLAG_VIRTUAL_HARD_KEY,
                 InputDevice.SOURCE_KEYBOARD);
-        //Make KeyEvent work on multi-display environment
-        if (getDisplay() != null) {
-            final int displayId = getDisplay().getDisplayId();
 
-            if (displayId != INVALID_DISPLAY) {
-                ev.setDisplayId(displayId);
-            }
+        int displayId = INVALID_DISPLAY;
+
+        // Make KeyEvent work on multi-display environment
+        if (getDisplay() != null) {
+            displayId = getDisplay().getDisplayId();
+        }
+        // Bubble controller will give us a valid display id if it should get the back event
+        BubbleController bubbleController = Dependency.get(BubbleController.class);
+        int bubbleDisplayId = bubbleController.getExpandedDisplayId(mContext);
+        if (mCode == KeyEvent.KEYCODE_BACK && bubbleDisplayId != INVALID_DISPLAY) {
+            displayId = bubbleDisplayId;
+        }
+        if (displayId != INVALID_DISPLAY) {
+            ev.setDisplayId(displayId);
         }
-        InputManager.getInstance().injectInputEvent(ev,
-                InputManager.INJECT_INPUT_EVENT_MODE_ASYNC);
+        mInputManager.injectInputEvent(ev, InputManager.INJECT_INPUT_EVENT_MODE_ASYNC);
     }
 
     @Override
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyButtonViewTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyButtonViewTest.java
index 39bdf204cd65..d16dc168d74c 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyButtonViewTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/KeyButtonViewTest.java
@@ -18,8 +18,16 @@ import static com.android.internal.logging.nano.MetricsProto.MetricsEvent.ACTION
 import static com.android.internal.logging.nano.MetricsProto.MetricsEvent.FIELD_FLAGS;
 import static com.android.internal.logging.nano.MetricsProto.MetricsEvent.FIELD_NAV_ACTION;
 
+import static junit.framework.Assert.assertEquals;
+
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.argThat;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import android.hardware.input.InputManager;
 import android.metrics.LogMaker;
 import android.testing.AndroidTestingRunner;
 import android.testing.TestableLooper;
@@ -31,12 +39,15 @@ import androidx.test.filters.SmallTest;
 import com.android.internal.logging.MetricsLogger;
 import com.android.internal.logging.nano.MetricsProto.MetricsEvent;
 import com.android.systemui.SysuiTestCase;
+import com.android.systemui.bubbles.BubbleController;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.ArgumentMatcher;
-import org.mockito.Mockito;
+import org.mockito.Captor;
+import org.mockito.MockitoAnnotations;
 
 import java.util.Objects;
 
@@ -47,12 +58,19 @@ public class KeyButtonViewTest extends SysuiTestCase {
 
     private KeyButtonView mKeyButtonView;
     private MetricsLogger mMetricsLogger;
+    private BubbleController mBubbleController;
+    private InputManager mInputManager = mock(InputManager.class);
+    @Captor
+    private ArgumentCaptor<KeyEvent> mInputEventCaptor;
 
     @Before
     public void setup() throws Exception {
+        MockitoAnnotations.initMocks(this);
         mMetricsLogger = mDependency.injectMockDependency(MetricsLogger.class);
-        TestableLooper.get(this).runWithLooper(() ->
-                mKeyButtonView = new KeyButtonView(mContext, null));
+        mBubbleController = mDependency.injectMockDependency(BubbleController.class);
+        TestableLooper.get(this).runWithLooper(() -> {
+            mKeyButtonView = new KeyButtonView(mContext, null, 0, mInputManager);
+        });
     }
 
     @Test
@@ -63,7 +81,7 @@ public class KeyButtonViewTest extends SysuiTestCase {
         mKeyButtonView.setCode(code);
         mKeyButtonView.sendEvent(action, flags);
 
-        Mockito.verify(mMetricsLogger).write(argThat(new ArgumentMatcher<LogMaker>() {
+        verify(mMetricsLogger).write(argThat(new ArgumentMatcher<LogMaker>() {
             public String mReason;
 
             @Override
@@ -91,4 +109,18 @@ public class KeyButtonViewTest extends SysuiTestCase {
         }));
     }
 
+    @Test
+    public void testBubbleEvents_bubbleExpanded() {
+        when(mBubbleController.getExpandedDisplayId(mContext)).thenReturn(3);
+
+        int action = KeyEvent.ACTION_DOWN;
+        int flags = 0;
+        int code = KeyEvent.KEYCODE_BACK;
+        mKeyButtonView.setCode(code);
+        mKeyButtonView.sendEvent(action, flags);
+
+        verify(mInputManager, times(1)).injectInputEvent(mInputEventCaptor.capture(),
+                anyInt());
+        assertEquals(3, mInputEventCaptor.getValue().getDisplayId());
+    }
 }
\ No newline at end of file