diff --git a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
index 5777fb9..68159fa 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DepthController.java
@@ -26,6 +26,7 @@
 import android.os.IBinder;
 import android.os.SystemProperties;
 import android.util.FloatProperty;
+import android.view.AttachedSurfaceControl;
 import android.view.CrossWindowBlurListeners;
 import android.view.SurfaceControl;
 import android.view.View;
@@ -316,7 +317,12 @@
                 transaction.setEarlyWakeupEnd();
                 mInEarlyWakeUp = false;
             }
-            transaction.apply();
+
+            AttachedSurfaceControl rootSurfaceControl =
+                    mLauncher.getRootView().getRootSurfaceControl();
+            if (rootSurfaceControl != null) {
+                rootSurfaceControl.applyTransactionOnDraw(transaction);
+            }
         }
         return true;
     }
diff --git a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
index 35a851a..81e6917 100644
--- a/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
+++ b/quickstep/src/com/android/quickstep/OrientationTouchTransformer.java
@@ -462,7 +462,7 @@
                             + "mRotation: " + mRotation
                             + " this: " + this);
                 }
-                event.transform(mTmpMatrix);
+                event.applyTransform(mTmpMatrix);
                 return true;
             }
             mTmpPoint[0] = event.getX();
@@ -478,7 +478,7 @@
             }
 
             if (contains(mTmpPoint[0], mTmpPoint[1])) {
-                event.transform(mTmpMatrix);
+                event.applyTransform(mTmpMatrix);
                 return true;
             }
             return false;
diff --git a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
index ec51599..1f57e99 100644
--- a/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
+++ b/quickstep/src/com/android/quickstep/SwipeUpAnimationLogic.java
@@ -278,6 +278,8 @@
         RemoteAnimationTargetCompat primaryTaskTarget;
         RemoteAnimationTargetCompat secondaryTaskTarget;
 
+        // TODO(b/197568823) Determine if we need to exclude assistant as one of the targets we
+        //  animate
         if (!mIsSwipeForStagedSplit) {
             primaryTaskTarget = targets.findTask(mGestureState.getRunningTaskId());
             mRemoteTargetHandles[0].mTransformParams.setTargetSet(targets);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index dc7347e..2a19ca2 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -4019,12 +4019,14 @@
             recentsAnimationTargets.addReleaseCheck(mSyncTransactionApplier);
         }
 
-        // TODO Consolidate this shared code with SwipeUpAnimationLogic (or mabe just reuse
-        //  what that class has and pass it into here
-        mRemoteTargetHandles = new RemoteTargetHandle[recentsAnimationTargets.apps.length];
+        RemoteAnimationTargetCompat dividerTarget =
+                recentsAnimationTargets.getNonAppTargetOfType(TYPE_DOCK_DIVIDER);
+        // TODO Consolidate this shared code with SwipeUpAnimationLogic (or maybe just reuse
+        //  what that class has and pass it into here)
+        mRemoteTargetHandles = new RemoteTargetHandle[dividerTarget == null ? 1 : 2];
         TaskViewSimulator primaryTvs = createTaskViewSimulator();
         mRemoteTargetHandles[0] = new RemoteTargetHandle(primaryTvs, new TransformParams());
-        if (recentsAnimationTargets.apps.length == 1) {
+        if (dividerTarget == null) {
             mRemoteTargetHandles[0].mTaskViewSimulator
                     .setPreview(recentsAnimationTargets.apps[0], null);
             mRemoteTargetHandles[0].mTransformParams.setTargetSet(recentsAnimationTargets);
@@ -4034,8 +4036,6 @@
             secondaryTvs.recentsViewScale.value = 1;
 
             mRemoteTargetHandles[1] = new RemoteTargetHandle(secondaryTvs, new TransformParams());
-            RemoteAnimationTargetCompat dividerTarget =
-                    recentsAnimationTargets.getNonAppTargetOfType(TYPE_DOCK_DIVIDER);
             RemoteAnimationTargetCompat primaryTaskTarget = recentsAnimationTargets.apps[0];
             RemoteAnimationTargetCompat secondaryTaskTarget = recentsAnimationTargets.apps[1];
             SplitConfigurationOptions.StagedSplitBounds
diff --git a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
similarity index 98%
rename from quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
rename to quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
index 2b1b57c..159a51f 100644
--- a/quickstep/robolectric_tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
+++ b/quickstep/tests/src/com/android/quickstep/OrientationTouchTransformerTest.java
@@ -19,6 +19,8 @@
 
 import static android.view.Display.DEFAULT_DISPLAY;
 
+import static androidx.test.core.app.ApplicationProvider.getApplicationContext;
+
 import static com.android.quickstep.SysUINavigationMode.Mode.NO_BUTTON;
 
 import static org.junit.Assert.assertFalse;
@@ -40,6 +42,9 @@
 import android.view.MotionEvent;
 import android.view.Surface;
 
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.filters.SmallTest;
+
 import com.android.launcher3.ResourceUtils;
 import com.android.launcher3.util.DisplayController;
 
@@ -47,10 +52,9 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.MockitoAnnotations;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.RuntimeEnvironment;
 
-@RunWith(RobolectricTestRunner.class)
+@SmallTest
+@RunWith(AndroidJUnit4.class)
 public class OrientationTouchTransformerTest {
     static class ScreenSize {
         int mHeight;
@@ -293,7 +297,7 @@
     }
 
     private DisplayController.Info createDisplayInfo(ScreenSize screenSize, int rotation) {
-        Context context = RuntimeEnvironment.application;
+        Context context = getApplicationContext();
         Display display = spy(context.getSystemService(DisplayManager.class)
                 .getDisplay(DEFAULT_DISPLAY));
 
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 38beeea..42b2033 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -333,7 +333,8 @@
         int paddingBottom = grid.cellLayoutBottomPaddingPx;
 
         int panelCount = getPanelCount();
-        for (int i = mWorkspaceScreens.size() - 1; i >= 0; i--) {
+        int numberOfScreens = mScreenOrder.size();
+        for (int i = 0; i < numberOfScreens; i++) {
             int paddingLeft = paddingLeftRight;
             int paddingRight = paddingLeftRight;
             if (panelCount > 1) {
@@ -348,7 +349,9 @@
                     paddingRight = 0;
                 }
             }
-            mWorkspaceScreens.valueAt(i).setPadding(paddingLeft, 0, paddingRight, paddingBottom);
+            // SparseArrayMap doesn't keep the order
+            mWorkspaceScreens.get(mScreenOrder.get(i))
+                    .setPadding(paddingLeft, 0, paddingRight, paddingBottom);
         }
     }
 
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 9264b8a..f1fe508 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -248,6 +248,8 @@
         if (pm.getComponentEnabledSetting(cn) != COMPONENT_ENABLED_STATE_ENABLED) {
             if (TestHelpers.isInLauncherProcess()) {
                 pm.setComponentEnabledSetting(cn, COMPONENT_ENABLED_STATE_ENABLED, DONT_KILL_APP);
+                // b/195031154
+                SystemClock.sleep(5000);
             } else {
                 try {
                     final int userId = ContextUtils.getUserId(getContext());
@@ -1033,6 +1035,15 @@
         }
     }
 
+    List<UiObject2> getChildren(UiObject2 container) {
+        try {
+            return container.getChildren();
+        } catch (StaleObjectException e) {
+            fail("The container disappeared from screen");
+            return null;
+        }
+    }
+
     private boolean hasLauncherObject(String resId) {
         return mDevice.hasObject(getLauncherObjectSelector(resId));
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/Widgets.java b/tests/tapl/com/android/launcher3/tapl/Widgets.java
index 99d9889..6e7264a 100644
--- a/tests/tapl/com/android/launcher3/tapl/Widgets.java
+++ b/tests/tapl/com/android/launcher3/tapl/Widgets.java
@@ -116,9 +116,9 @@
                     "widget_preview");
             int i = 0;
             for (; ; ) {
-                final Collection<UiObject2> tableRows = widgetsContainer.getChildren();
+                final Collection<UiObject2> tableRows = mLauncher.getChildren(widgetsContainer);
                 for (UiObject2 row : tableRows) {
-                    final Collection<UiObject2> widgetCells = row.getChildren();
+                    final Collection<UiObject2> widgetCells = mLauncher.getChildren(row);
                     for (UiObject2 widget : widgetCells) {
                         final UiObject2 label = mLauncher.findObjectInContainer(widget,
                                 labelSelector);