diff --git a/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
index 6b3f699..9e5d958 100644
--- a/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
+++ b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
@@ -17,8 +17,6 @@
 
 package com.android.quickstep;
 
-import static android.view.Display.DEFAULT_DISPLAY;
-
 import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
 
 import static com.android.launcher3.util.DisplayController.NavigationMode.NO_BUTTON;
@@ -27,17 +25,16 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.when;
 
-import android.content.Context;
 import android.content.res.Resources;
 import android.graphics.Point;
-import android.hardware.display.DisplayManager;
+import android.graphics.Rect;
+import android.util.ArrayMap;
 import android.util.DisplayMetrics;
+import android.util.Size;
 import android.view.Display;
 import android.view.MotionEvent;
 import android.view.Surface;
@@ -47,6 +44,10 @@
 
 import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.util.DisplayController;
+import com.android.launcher3.util.RotationUtils;
+import com.android.launcher3.util.WindowBounds;
+import com.android.launcher3.util.window.CachedDisplayInfo;
+import com.android.launcher3.util.window.WindowManagerProxy;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -56,18 +57,9 @@
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class OrientationTouchTransformerTest {
-    static class ScreenSize {
-        int mHeight;
-        int mWidth;
 
-        ScreenSize(int height, int width) {
-            mHeight = height;
-            mWidth = width;
-        }
-    }
-
-    private static final ScreenSize NORMAL_SCREEN_SIZE = new ScreenSize(2280, 1080);
-    private static final ScreenSize LARGE_SCREEN_SIZE = new ScreenSize(3280, 1080);
+    private static final Size NORMAL_SCREEN_SIZE = new Size(1080, 2280);
+    private static final Size LARGE_SCREEN_SIZE = new Size(1080, 3280);
     private static final float DENSITY_DISPLAY_METRICS = 3.0f;
 
     private OrientationTouchTransformer mTouchTransformer;
@@ -75,7 +67,6 @@
     Resources mResources;
     private DisplayController.Info mInfo;
 
-
     @Before
     public void setup() {
         MockitoAnnotations.initMocks(this);
@@ -296,33 +287,24 @@
         assertTrue(mTouchTransformer.touchInValidSwipeRegions(inRegion2.getX(), inRegion2.getY()));
     }
 
-    private DisplayController.Info createDisplayInfo(ScreenSize screenSize, int rotation) {
-        Context context = getApplicationContext();
-        Display display = spy(context.getSystemService(DisplayManager.class)
-                .getDisplay(DEFAULT_DISPLAY));
-
-        Point p = new Point(screenSize.mWidth, screenSize.mHeight);
-        if (rotation == Surface.ROTATION_90 || rotation == Surface.ROTATION_270) {
-            p.set(screenSize.mHeight, screenSize.mWidth);
-        }
-
-        doReturn(rotation).when(display).getRotation();
-        doAnswer(i -> {
-            ((Point) i.getArgument(0)).set(p.x, p.y);
-            return null;
-        }).when(display).getRealSize(any(Point.class));
-        doAnswer(i -> {
-            ((Point) i.getArgument(0)).set(p.x, p.y);
-            ((Point) i.getArgument(1)).set(p.x, p.y);
-            return null;
-        }).when(display).getCurrentSizeRange(any(Point.class), any(Point.class));
-        return new DisplayController.Info(context, display);
+    private DisplayController.Info createDisplayInfo(Size screenSize, int rotation) {
+        Point displaySize = new Point(screenSize.getWidth(), screenSize.getHeight());
+        RotationUtils.rotateSize(displaySize, rotation);
+        CachedDisplayInfo cdi = new CachedDisplayInfo(displaySize, rotation);
+        WindowBounds wm = new WindowBounds(
+                new Rect(0, 0, displaySize.x, displaySize.y),
+                new Rect());
+        WindowManagerProxy wmProxy = mock(WindowManagerProxy.class);
+        doReturn(cdi).when(wmProxy).getDisplayInfo(any(), any());
+        doReturn(wm).when(wmProxy).getRealBounds(any(), any(), any());
+        return new DisplayController.Info(
+                getApplicationContext(), mock(Display.class), wmProxy, new ArrayMap<>());
     }
 
-    private float generateTouchRegionHeight(ScreenSize screenSize, int rotation) {
-        float height = screenSize.mHeight;
+    private float generateTouchRegionHeight(Size screenSize, int rotation) {
+        float height = screenSize.getHeight();
         if (rotation == Surface.ROTATION_90 || rotation == Surface.ROTATION_270) {
-            height = screenSize.mWidth;
+            height = screenSize.getWidth();
         }
         return height - ResourceUtils.DEFAULT_NAVBAR_VALUE * DENSITY_DISPLAY_METRICS;
     }
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
index c6cdafc..7d414f4 100644
--- a/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
+++ b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
@@ -164,7 +164,7 @@
                 }
 
                 WindowManagerProxy wmProxy = mock(WindowManagerProxy.class);
-                doReturn(cdi).when(wmProxy).getDisplayInfo(any());
+                doReturn(cdi).when(wmProxy).getDisplayInfo(any(), any());
                 doReturn(wm).when(wmProxy).getRealBounds(any(), any(), any());
 
                 ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> perDisplayBoundsCache =
diff --git a/src/com/android/launcher3/util/DisplayController.java b/src/com/android/launcher3/util/DisplayController.java
index 8b4ff85..d68aab6 100644
--- a/src/com/android/launcher3/util/DisplayController.java
+++ b/src/com/android/launcher3/util/DisplayController.java
@@ -306,7 +306,7 @@
         public Info(Context context, Display display,
                 WindowManagerProxy wmProxy,
                 ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> perDisplayBoundsCache) {
-            CachedDisplayInfo displayInfo = wmProxy.getDisplayInfo(display);
+            CachedDisplayInfo displayInfo = wmProxy.getDisplayInfo(context, display);
             rotation = displayInfo.rotation;
             currentSize = displayInfo.size;
             displayId = displayInfo.id;
@@ -372,7 +372,7 @@
         pw.println("  id=" + info.displayId);
         pw.println("  rotation=" + info.rotation);
         pw.println("  fontScale=" + info.fontScale);
-        pw.println("  densityDpi=" + info.displayId);
+        pw.println("  densityDpi=" + info.densityDpi);
         pw.println("  navigationMode=" + info.navigationMode.name());
         pw.println("  currentSize=" + info.currentSize);
         pw.println("  supportedBounds=" + info.supportedBounds);
diff --git a/src/com/android/launcher3/util/window/WindowManagerProxy.java b/src/com/android/launcher3/util/window/WindowManagerProxy.java
index ba3d981..96f3ef6 100644
--- a/src/com/android/launcher3/util/window/WindowManagerProxy.java
+++ b/src/com/android/launcher3/util/window/WindowManagerProxy.java
@@ -15,6 +15,9 @@
  */
 package com.android.launcher3.util.window;
 
+import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION;
+
 import static com.android.launcher3.ResourceUtils.INVALID_RESOURCE_HANDLE;
 import static com.android.launcher3.ResourceUtils.NAVBAR_HEIGHT;
 import static com.android.launcher3.ResourceUtils.NAVBAR_HEIGHT_LANDSCAPE;
@@ -47,7 +50,6 @@
 import com.android.launcher3.R;
 import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.Utilities;
-import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.ResourceBasedOverride;
 import com.android.launcher3.util.WindowBounds;
@@ -91,7 +93,10 @@
         ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> result = new ArrayMap<>();
         for (Display display : displays) {
             if (isInternalDisplay(display)) {
-                CachedDisplayInfo info = getDisplayInfo(display).normalize();
+                Context displayContext = Utilities.ATLEAST_S
+                        ? context.createWindowContext(display, TYPE_APPLICATION, null)
+                        : context.createDisplayContext(display);
+                CachedDisplayInfo info = getDisplayInfo(displayContext, display).normalize();
                 WindowBounds[] bounds = estimateWindowBounds(context, info);
                 result.put(info.id, Pair.create(info, bounds));
             }
@@ -122,7 +127,7 @@
         }
 
         WindowMetrics wm = windowContext.getSystemService(WindowManager.class)
-                .getCurrentWindowMetrics();
+                .getMaximumWindowMetrics();
 
         Rect insets = new Rect();
         normalizeWindowInsets(windowContext, wm.getWindowInsets(), insets);
@@ -268,21 +273,23 @@
      * Returns a CachedDisplayInfo initialized for the current display
      */
     @TargetApi(Build.VERSION_CODES.S)
-    public CachedDisplayInfo getDisplayInfo(Display display) {
-        int rotation = display.getRotation();
-
-        Point size = new Point();
-        display.getRealSize(size);
-
+    public CachedDisplayInfo getDisplayInfo(Context displayContext, Display display) {
+        int rotation = getRotation(displayContext);
         Rect cutoutRect = new Rect();
+        Point size = new Point();
         if (Utilities.ATLEAST_S) {
-            DisplayCutout cutout = display.getCutout();
+            WindowMetrics wm = displayContext.getSystemService(WindowManager.class)
+                    .getMaximumWindowMetrics();
+            DisplayCutout cutout = wm.getWindowInsets().getDisplayCutout();
             if (cutout != null) {
                 cutoutRect.set(cutout.getSafeInsetLeft(), cutout.getSafeInsetTop(),
                         cutout.getSafeInsetRight(), cutout.getSafeInsetBottom());
             }
-        }
 
+            size.set(wm.getBounds().right, wm.getBounds().bottom);
+        } else {
+            display.getRealSize(size);
+        }
         return new CachedDisplayInfo(getDisplayId(display), size, rotation, cutoutRect);
     }
 
@@ -305,7 +312,9 @@
                 // Ignore
             }
         }
-        return d == null ? DisplayController.INSTANCE.get(context).getInfo().rotation
-                : d.getRotation();
+        if (d == null) {
+            d = context.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY);
+        }
+        return d.getRotation();
     }
 }