diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/robolectric_tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
deleted file mode 100644
index fd93d98..0000000
--- a/quickstep/robolectric_tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
+++ /dev/null
@@ -1,206 +0,0 @@
-/*
- * Copyright (C) 2020 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.quickstep.util;
-
-import static android.view.Display.DEFAULT_DISPLAY;
-
-import static org.mockito.Mockito.mock;
-
-import android.content.Context;
-import android.graphics.Rect;
-import android.graphics.RectF;
-import android.hardware.display.DisplayManager;
-import android.view.Surface;
-import android.view.SurfaceControl;
-
-import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.shadows.LShadowDisplay;
-import com.android.launcher3.util.DisplayController;
-import com.android.quickstep.LauncherActivityInterface;
-import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
-
-import org.hamcrest.Description;
-import org.hamcrest.TypeSafeMatcher;
-import org.junit.Assert;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.annotation.LooperMode;
-import org.robolectric.annotation.LooperMode.Mode;
-import org.robolectric.shadow.api.Shadow;
-import org.robolectric.shadows.ShadowDisplayManager;
-
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
-public class TaskViewSimulatorTest {
-
-    @Test
-    public void taskProperlyScaled_portrait_noRotation_sameInsets1() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(1200, 2450)
-                .withInsets(new Rect(0, 80, 0, 120))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_portrait_noRotation_sameInsets2() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(1200, 2450)
-                .withInsets(new Rect(55, 80, 55, 120))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_noRotation_sameInsets1() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(2450, 1250)
-                .withInsets(new Rect(0, 80, 0, 40))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_noRotation_sameInsets2() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(2450, 1250)
-                .withInsets(new Rect(0, 80, 120, 0))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_noRotation_sameInsets3() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(2450, 1250)
-                .withInsets(new Rect(55, 80, 55, 120))
-                .verifyNoTransforms();
-    }
-
-    @Test
-    public void taskProperlyScaled_landscape_rotated() {
-        new TaskMatrixVerifier()
-                .withLauncherSize(1200, 2450)
-                .withInsets(new Rect(0, 80, 0, 120))
-                .withAppBounds(
-                        new Rect(0, 0, 2450, 1200),
-                        new Rect(0, 80, 0, 120),
-                        Surface.ROTATION_90)
-                .verifyNoTransforms();
-    }
-
-    private static class TaskMatrixVerifier extends TransformParams {
-
-        private final Context mContext = RuntimeEnvironment.application;
-
-        private Rect mAppBounds = new Rect();
-        private Rect mLauncherInsets = new Rect();
-
-        private Rect mAppInsets;
-
-        private int mAppRotation = -1;
-        private DeviceProfile mDeviceProfile;
-
-        TaskMatrixVerifier withLauncherSize(int width, int height) {
-            ShadowDisplayManager.changeDisplay(DEFAULT_DISPLAY,
-                    String.format("w%sdp-h%sdp-mdpi", width, height));
-            if (mAppBounds.isEmpty()) {
-                mAppBounds.set(0, 0, width, height);
-            }
-            return this;
-        }
-
-        TaskMatrixVerifier withInsets(Rect insets) {
-            LShadowDisplay shadowDisplay = Shadow.extract(
-                    mContext.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY));
-            shadowDisplay.setInsets(insets);
-            mLauncherInsets.set(insets);
-            return this;
-        }
-
-        TaskMatrixVerifier withAppBounds(Rect bounds, Rect insets, int appRotation) {
-            mAppBounds.set(bounds);
-            mAppInsets = insets;
-            mAppRotation = appRotation;
-            return this;
-        }
-
-        void verifyNoTransforms() {
-            mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(mContext)
-                    .getDeviceProfile(mContext);
-            mDeviceProfile.updateInsets(mLauncherInsets);
-
-            TaskViewSimulator tvs = new TaskViewSimulator(mContext,
-                    LauncherActivityInterface.INSTANCE);
-            tvs.setDp(mDeviceProfile);
-
-            int launcherRotation = DisplayController.INSTANCE.get(mContext).getInfo().rotation;
-            if (mAppRotation < 0) {
-                mAppRotation = launcherRotation;
-            }
-            tvs.getOrientationState().update(launcherRotation, mAppRotation);
-            if (mAppInsets == null) {
-                mAppInsets = new Rect(mLauncherInsets);
-            }
-            tvs.setPreviewBounds(mAppBounds, mAppInsets);
-
-            tvs.fullScreenProgress.value = 1;
-            tvs.recentsViewScale.value = tvs.getFullScreenScale();
-            tvs.apply(this);
-        }
-
-        @Override
-        public SurfaceParams[] createSurfaceParams(BuilderProxy proxy) {
-            SurfaceParams.Builder builder = new SurfaceParams.Builder((SurfaceControl) null);
-            proxy.onBuildTargetParams(builder, mock(RemoteAnimationTargetCompat.class), this);
-            return new SurfaceParams[] {builder.build()};
-        }
-
-        @Override
-        public void applySurfaceParams(SurfaceParams[] params) {
-            // Verify that the task position remains the same
-            RectF newAppBounds = new RectF(mAppBounds);
-            params[0].matrix.mapRect(newAppBounds);
-            Assert.assertThat(newAppBounds, new AlmostSame(mAppBounds));
-
-            System.err.println("Bounds mapped: " + mAppBounds + " => " + newAppBounds);
-        }
-    }
-
-    private static class AlmostSame extends TypeSafeMatcher<RectF>  {
-
-        // Allow 1px error margin to account for float to int conversions
-        private final float mError = 1f;
-        private final Rect mExpected;
-
-        AlmostSame(Rect expected) {
-            mExpected = expected;
-        }
-
-        @Override
-        protected boolean matchesSafely(RectF item) {
-            return Math.abs(item.left - mExpected.left) < mError
-                    && Math.abs(item.top - mExpected.top) < mError
-                    && Math.abs(item.right - mExpected.right) < mError
-                    && Math.abs(item.bottom - mExpected.bottom) < mError;
-        }
-
-        @Override
-        public void describeTo(Description description) {
-            description.appendValue(mExpected);
-        }
-    }
-}
diff --git a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
index 841e578..a4db596 100644
--- a/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
+++ b/quickstep/src/com/android/quickstep/util/RecentsOrientedState.java
@@ -600,17 +600,7 @@
             width = Math.min(currentSize.x, currentSize.y);
             height = Math.max(currentSize.x, currentSize.y);
         }
-
-        DeviceProfile bestMatch = idp.supportedProfiles.get(0);
-        float minDiff = Float.MAX_VALUE;
-        for (DeviceProfile profile : idp.supportedProfiles) {
-            float diff = Math.abs(profile.widthPx - width) + Math.abs(profile.heightPx - height);
-            if (diff < minDiff) {
-                minDiff = diff;
-                bestMatch = profile;
-            }
-        }
-        return bestMatch;
+        return idp.getBestMatch(width, height);
     }
 
     private static String nameAndAddress(Object obj) {
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java b/quickstep/tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
similarity index 89%
rename from quickstep/robolectric_tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
rename to quickstep/tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
index 656379f..47ef13b 100644
--- a/quickstep/robolectric_tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
+++ b/quickstep/tests/src/com/android/quickstep/util/RecentsOrientedStateTest.java
@@ -19,33 +19,34 @@
 import static android.view.Surface.ROTATION_180;
 import static android.view.Surface.ROTATION_90;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotEquals;
 
 import android.content.Context;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.quickstep.FallbackActivityInterface;
 
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
-import org.robolectric.annotation.LooperMode;
-import org.robolectric.annotation.LooperMode.Mode;
 
 /**
  * Tests for {@link RecentsOrientedState}
  */
-@RunWith(RobolectricTestRunner.class)
-@LooperMode(Mode.PAUSED)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class RecentsOrientedStateTest {
 
     private RecentsOrientedState mR1, mR2;
 
     @Before
     public void setup() {
-        Context context = RuntimeEnvironment.application;
+        Context context = getApplicationContext();
         mR1 = new RecentsOrientedState(context, FallbackActivityInterface.INSTANCE, i -> { });
         mR2 = new RecentsOrientedState(context, FallbackActivityInterface.INSTANCE, i -> { });
         assertEquals(mR1.getStateId(), mR2.getStateId());
diff --git a/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
new file mode 100644
index 0000000..9c5cfcd
--- /dev/null
+++ b/quickstep/tests/src/com/android/quickstep/util/TaskViewSimulatorTest.java
@@ -0,0 +1,246 @@
+/*
+ * Copyright (C) 2020 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.quickstep.util;
+
+import static android.view.Display.DEFAULT_DISPLAY;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+import android.graphics.Point;
+import android.graphics.Rect;
+import android.graphics.RectF;
+import android.view.Display;
+import android.view.Surface;
+import android.view.SurfaceControl;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
+import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.util.DisplayController;
+import com.android.launcher3.util.DisplayController.Info;
+import com.android.launcher3.util.LauncherModelHelper;
+import com.android.launcher3.util.ReflectionHelpers;
+import com.android.quickstep.FallbackActivityInterface;
+import com.android.quickstep.SysUINavigationMode;
+import com.android.quickstep.SystemUiProxy;
+import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
+import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
+
+import org.hamcrest.Description;
+import org.hamcrest.TypeSafeMatcher;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@SmallTest
+@RunWith(AndroidJUnit4.class)
+public class TaskViewSimulatorTest {
+
+    @Test
+    public void taskProperlyScaled_portrait_noRotation_sameInsets1() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(1200, 2450)
+                .withInsets(new Rect(0, 80, 0, 120))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_portrait_noRotation_sameInsets2() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(1200, 2450)
+                .withInsets(new Rect(55, 80, 55, 120))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_noRotation_sameInsets1() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(2450, 1250)
+                .withInsets(new Rect(0, 80, 0, 40))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_noRotation_sameInsets2() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(2450, 1250)
+                .withInsets(new Rect(0, 80, 120, 0))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_noRotation_sameInsets3() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(2450, 1250)
+                .withInsets(new Rect(55, 80, 55, 120))
+                .verifyNoTransforms();
+    }
+
+    @Test
+    public void taskProperlyScaled_landscape_rotated() {
+        new TaskMatrixVerifier()
+                .withLauncherSize(1200, 2450)
+                .withInsets(new Rect(0, 80, 0, 120))
+                .withAppBounds(
+                        new Rect(0, 0, 2450, 1200),
+                        new Rect(0, 80, 0, 120),
+                        Surface.ROTATION_90)
+                .verifyNoTransforms();
+    }
+
+    private static class TaskMatrixVerifier extends TransformParams {
+
+        private Point mDisplaySize = new Point();
+        private Rect mDisplayInsets = new Rect();
+        private Rect mAppBounds = new Rect();
+        private Rect mLauncherInsets = new Rect();
+
+        private Rect mAppInsets;
+
+        private int mAppRotation = -1;
+        private DeviceProfile mDeviceProfile;
+
+        TaskMatrixVerifier withLauncherSize(int width, int height) {
+            mDisplaySize.set(width, height);
+            if (mAppBounds.isEmpty()) {
+                mAppBounds.set(0, 0, width, height);
+            }
+            return this;
+        }
+
+        TaskMatrixVerifier withInsets(Rect insets) {
+            mDisplayInsets.set(insets);
+            mLauncherInsets.set(insets);
+            return this;
+        }
+
+        TaskMatrixVerifier withAppBounds(Rect bounds, Rect insets, int appRotation) {
+            mAppBounds.set(bounds);
+            mAppInsets = insets;
+            mAppRotation = appRotation;
+            return this;
+        }
+
+        void verifyNoTransforms() {
+            LauncherModelHelper helper = new LauncherModelHelper();
+            try {
+                helper.sandboxContext.allow(SystemUiProxy.INSTANCE);
+                helper.sandboxContext.allow(SysUINavigationMode.INSTANCE);
+
+                Display display = mock(Display.class);
+                doReturn(DEFAULT_DISPLAY).when(display).getDisplayId();
+                doReturn(mDisplaySize.x > mDisplaySize.y ? Surface.ROTATION_90 : Surface.ROTATION_0)
+                        .when(display).getRotation();
+                doAnswer(i -> {
+                    ((Point) i.getArgument(0)).set(mDisplaySize.x, mDisplaySize.y);
+                    return null;
+                }).when(display).getRealSize(any());
+                doAnswer(i -> {
+                    Point smallestSize = i.getArgument(0);
+                    Point largestSize = i.getArgument(1);
+                    smallestSize.x = smallestSize.y = Math.min(mDisplaySize.x, mDisplaySize.y);
+                    largestSize.x = largestSize.y = Math.max(mDisplaySize.x, mDisplaySize.y);
+
+                    smallestSize.x -= mDisplayInsets.left + mDisplayInsets.right;
+                    largestSize.x -= mDisplayInsets.left + mDisplayInsets.right;
+
+                    smallestSize.y -= mDisplayInsets.top + mDisplayInsets.bottom;
+                    largestSize.y -= mDisplayInsets.top + mDisplayInsets.bottom;
+                    return null;
+                }).when(display).getCurrentSizeRange(any(), any());
+                DisplayController.Info mockInfo = new Info(helper.sandboxContext, display);
+
+                DisplayController controller =
+                        DisplayController.INSTANCE.get(helper.sandboxContext);
+                controller.close();
+                ReflectionHelpers.setField(controller, "mInfo", mockInfo);
+
+                mDeviceProfile = InvariantDeviceProfile.INSTANCE.get(helper.sandboxContext)
+                        .getBestMatch(mAppBounds.width(), mAppBounds.height());
+                mDeviceProfile.updateInsets(mLauncherInsets);
+
+                TaskViewSimulator tvs = new TaskViewSimulator(helper.sandboxContext,
+                        FallbackActivityInterface.INSTANCE);
+                tvs.setDp(mDeviceProfile);
+
+                int launcherRotation = mockInfo.rotation;
+                if (mAppRotation < 0) {
+                    mAppRotation = launcherRotation;
+                }
+
+                tvs.getOrientationState().update(launcherRotation, mAppRotation);
+                if (mAppInsets == null) {
+                    mAppInsets = new Rect(mLauncherInsets);
+                }
+                tvs.setPreviewBounds(mAppBounds, mAppInsets);
+
+                tvs.fullScreenProgress.value = 1;
+                tvs.recentsViewScale.value = tvs.getFullScreenScale();
+                tvs.apply(this);
+            } finally {
+                helper.destroy();
+            }
+        }
+
+        @Override
+        public SurfaceParams[] createSurfaceParams(BuilderProxy proxy) {
+            SurfaceParams.Builder builder = new SurfaceParams.Builder((SurfaceControl) null);
+            proxy.onBuildTargetParams(builder, mock(RemoteAnimationTargetCompat.class), this);
+            return new SurfaceParams[] {builder.build()};
+        }
+
+        @Override
+        public void applySurfaceParams(SurfaceParams[] params) {
+            // Verify that the task position remains the same
+            RectF newAppBounds = new RectF(mAppBounds);
+            params[0].matrix.mapRect(newAppBounds);
+            Assert.assertThat(newAppBounds, new AlmostSame(mAppBounds));
+
+            System.err.println("Bounds mapped: " + mAppBounds + " => " + newAppBounds);
+        }
+    }
+
+    private static class AlmostSame extends TypeSafeMatcher<RectF>  {
+
+        // Allow .1% error margin to account for float to int conversions
+        private final float mErrorFactor = .001f;
+        private final Rect mExpected;
+
+        AlmostSame(Rect expected) {
+            mExpected = expected;
+        }
+
+        @Override
+        protected boolean matchesSafely(RectF item) {
+            float errorWidth = mErrorFactor * mExpected.width();
+            float errorHeight = mErrorFactor * mExpected.height();
+            return Math.abs(item.left - mExpected.left) < errorWidth
+                    && Math.abs(item.top - mExpected.top) < errorHeight
+                    && Math.abs(item.right - mExpected.right) < errorWidth
+                    && Math.abs(item.bottom - mExpected.bottom) < errorHeight;
+        }
+
+        @Override
+        public void describeTo(Description description) {
+            description.appendValue(mExpected);
+        }
+    }
+}
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index a2189c9..244cb59 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -548,15 +548,18 @@
         Resources res = context.getResources();
         Configuration config = context.getResources().getConfiguration();
 
-        float availableWidth = config.screenWidthDp * res.getDisplayMetrics().density;
-        float availableHeight = config.screenHeightDp * res.getDisplayMetrics().density;
+        float screenWidth = config.screenWidthDp * res.getDisplayMetrics().density;
+        float screenHeight = config.screenHeightDp * res.getDisplayMetrics().density;
+        return getBestMatch(screenWidth, screenHeight);
+    }
 
+    public DeviceProfile getBestMatch(float screenWidth, float screenHeight) {
         DeviceProfile bestMatch = supportedProfiles.get(0);
         float minDiff = Float.MAX_VALUE;
 
         for (DeviceProfile profile : supportedProfiles) {
-            float diff = Math.abs(profile.availableWidthPx - availableWidth)
-                    + Math.abs(profile.availableHeightPx - availableHeight);
+            float diff = Math.abs(profile.widthPx - screenWidth)
+                    + Math.abs(profile.heightPx - screenHeight);
             if (diff < minDiff) {
                 minDiff = diff;
                 bestMatch = profile;