Refresh display bounds when display layout updated

Previoulsy mDefaultDisplayBounds & mLastVisualDisplayBounds are
incorrect after device reboot without re-register DisplayAreaOrganizer.
Then the bounds does not set correctly for transition.
When AnimationController#getAnimator() get incorrect display bounds,
either surface bounds or rounded corner does not crop correctly.

Test: manual enable OHM, reboot device and trigger
     (Bounds/Rounded corner should works when trigger)
Test: atest WMShellUnitTests
Bug: 212651010
Change-Id: Ib30edf6ecc2244f7fec7e233b0fb048d5ee40f3c
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java
index 507204c..c9c73fd8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedController.java
@@ -277,6 +277,7 @@
         registerSettingObservers(mUserId);
         setupTimeoutListener();
         updateSettings();
+        updateDisplayLayout(mContext.getDisplayId());
 
         mAccessibilityManager = AccessibilityManager.getInstance(context);
         mAccessibilityManager.addAccessibilityStateChangeListener(
@@ -448,8 +449,13 @@
         onShortcutEnabledChanged();
     }
 
-    private void updateDisplayLayout(int displayId) {
+    @VisibleForTesting
+    void updateDisplayLayout(int displayId) {
         final DisplayLayout newDisplayLayout = mDisplayController.getDisplayLayout(displayId);
+        if (newDisplayLayout == null) {
+            Slog.w(TAG, "Failed to get new DisplayLayout.");
+            return;
+        }
         mDisplayAreaOrganizer.setDisplayLayout(newDisplayLayout);
         mTutorialHandler.onDisplayChanged(newDisplayLayout);
         mBackgroundPanelOrganizer.onDisplayChanged(newDisplayLayout);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
index 1b2f476..ec3ef5a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizer.java
@@ -123,9 +123,8 @@
             OneHandedBackgroundPanelOrganizer oneHandedBackgroundGradientOrganizer,
             ShellExecutor mainExecutor) {
         super(mainExecutor);
-        mDisplayLayout.set(displayLayout);
+        setDisplayLayout(displayLayout);
         mOneHandedSettingsUtil = oneHandedSettingsUtil;
-        updateDisplayBounds();
         mAnimationController = animationController;
         final int animationDurationConfig = context.getResources().getInteger(
                 R.integer.config_one_handed_translate_animation_duration);
@@ -282,6 +281,7 @@
     @VisibleForTesting
     void setDisplayLayout(@NonNull DisplayLayout displayLayout) {
         mDisplayLayout.set(displayLayout);
+        updateDisplayBounds();
     }
 
     @VisibleForTesting
@@ -289,6 +289,7 @@
         return mDisplayAreaTokenMap;
     }
 
+    @VisibleForTesting
     void updateDisplayBounds() {
         mDefaultDisplayBounds.set(0, 0, mDisplayLayout.width(), mDisplayLayout.height());
         mLastVisualDisplayBounds.set(mDefaultDisplayBounds);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedControllerTest.java
index 0a3a849..16bc5075 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedControllerTest.java
@@ -109,6 +109,7 @@
         mSpiedTransitionState = spy(new OneHandedState());
 
         when(mMockDisplayController.getDisplay(anyInt())).thenReturn(mDisplay);
+        when(mMockDisplayController.getDisplayLayout(anyInt())).thenReturn(null);
         when(mMockDisplayAreaOrganizer.getDisplayAreaTokenMap()).thenReturn(new ArrayMap<>());
         when(mMockDisplayAreaOrganizer.isReady()).thenReturn(true);
         when(mMockBackgroundOrganizer.isRegistered()).thenReturn(true);
@@ -153,6 +154,13 @@
     }
 
     @Test
+    public void testNullDisplayLayout() {
+        mSpiedOneHandedController.updateDisplayLayout(0);
+
+        verify(mMockDisplayAreaOrganizer, never()).setDisplayLayout(any());
+    }
+
+    @Test
     public void testStartOneHandedShouldTriggerScheduleOffset() {
         mSpiedTransitionState.setState(STATE_NONE);
         mSpiedOneHandedController.setOneHandedEnabled(true);
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java
index ef16fd3..1d92a48 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/onehanded/OneHandedDisplayAreaOrganizerTest.java
@@ -432,4 +432,11 @@
 
         assertThat(testSpiedDisplayAreaOrganizer.isReady()).isFalse();
     }
+
+    @Test
+    public void testDisplayArea_setDisplayLayout_should_updateDisplayBounds() {
+        mSpiedDisplayAreaOrganizer.setDisplayLayout(mDisplayLayout);
+
+        verify(mSpiedDisplayAreaOrganizer).updateDisplayBounds();
+    }
 }