Keep track of requestedVisibility when control unavailable

When control is unavailable, but the client still wants to hide or
show something, we keep track of it and execute the animation as
soon as we get control.

Also fix an issue with a crash where control was lost before
animation was started.

Test: InsetsControllerTest, InsetsSourceConsumerTest
Change-Id: I31c5177d9281161b8a607b24ddf094dbba640da3

8 files changed, 251 insertions, 55 deletions

diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index 23ba097f2b6d..f2263693897b 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -17,6 +17,7 @@
 package android.view;
 
 import static android.view.InsetsState.ITYPE_IME;
+import static android.view.InsetsState.toPublicType;
 
 import android.annotation.Nullable;
 import android.inputmethodservice.InputMethodService;
@@ -44,6 +45,12 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
      */
     private boolean mShowOnNextImeRender;
 
+    /**
+     * Tracks whether we have an outstanding request from the IME to show, but weren't able to
+     * execute it because we didn't have control yet.
+     */
+    private boolean mImeRequestedShow;
+
     public ImeInsetsSourceConsumer(
             InsetsState state, Supplier<Transaction> transactionSupplier,
             InsetsController controller) {
@@ -81,13 +88,14 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     public void onWindowFocusLost() {
         super.onWindowFocusLost();
         getImm().unregisterImeConsumer(this);
+        mImeRequestedShow = false;
     }
 
     @Override
-    public void setControl(@Nullable InsetsSourceControl control) {
-        super.setControl(control);
-        if (control == null) {
-            hide();
+    public void show(boolean fromIme) {
+        super.show(fromIme);
+        if (fromIme) {
+            mImeRequestedShow = true;
         }
     }
 
@@ -99,7 +107,11 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
     public @ShowResult int requestShow(boolean fromIme) {
         // TODO: ResultReceiver for IME.
         // TODO: Set mShowOnNextImeRender to automatically show IME and guard it with a flag.
-        if (fromIme) {
+
+        // If we had a request before to show from IME (tracked with mImeRequestedShow), reaching
+        // this code here means that we now got control, so we can start the animation immediately.
+        if (fromIme || mImeRequestedShow) {
+            mImeRequestedShow = false;
             return ShowResult.SHOW_IMMEDIATELY;
         }
 
@@ -115,6 +127,15 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
         getImm().notifyImeHidden();
     }
 
+    @Override
+    public void setControl(@Nullable InsetsSourceControl control, int[] showTypes,
+            int[] hideTypes) {
+        super.setControl(control, showTypes, hideTypes);
+        if (control == null) {
+            hide();
+        }
+    }
+
     private boolean isDummyOrEmptyEditor(EditorInfo info) {
         // TODO(b/123044812): Handle dummy input gracefully in IME Insets API
         return info == null || (info.fieldId <= 0 && info.inputType <= 0);
diff --git a/core/java/android/view/InsetsAnimationControlImpl.java b/core/java/android/view/InsetsAnimationControlImpl.java
index 5c24047a2c19..dad7696397ee 100644
--- a/core/java/android/view/InsetsAnimationControlImpl.java
+++ b/core/java/android/view/InsetsAnimationControlImpl.java
@@ -210,6 +210,10 @@ public class InsetsAnimationControlImpl implements WindowInsetsAnimationControll
         mListener.onCancelled();
     }
 
+    public boolean isCancelled() {
+        return mCancelled;
+    }
+
     InsetsAnimation getAnimation() {
         return mAnimation;
     }
diff --git a/core/java/android/view/InsetsController.java b/core/java/android/view/InsetsController.java
index 6f76497572d7..2785d216d24d 100644
--- a/core/java/android/view/InsetsController.java
+++ b/core/java/android/view/InsetsController.java
@@ -39,6 +39,7 @@ import android.util.Log;
 import android.util.Pair;
 import android.util.Property;
 import android.util.SparseArray;
+import android.view.InputDevice.MotionRange;
 import android.view.InsetsSourceConsumer.ShowResult;
 import android.view.InsetsState.InternalInsetsType;
 import android.view.SurfaceControl.Transaction;
@@ -273,7 +274,7 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
     private final String TAG = "InsetsControllerImpl";
 
     private final InsetsState mState = new InsetsState();
-    private final InsetsState mTmpState = new InsetsState();
+    private final InsetsState mLastDispachedState = new InsetsState();
 
     private final Rect mFrame = new Rect();
     private final BiFunction<InsetsController, Integer, InsetsSourceConsumer> mConsumerCreator;
@@ -367,15 +368,20 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
         return mState;
     }
 
-    boolean onStateChanged(InsetsState state) {
-        if (mState.equals(state)) {
+    public InsetsState getLastDispatchedState() {
+        return mLastDispachedState;
+    }
+
+    @VisibleForTesting
+    public boolean onStateChanged(InsetsState state) {
+        if (mState.equals(state) && mLastDispachedState.equals(state)) {
             return false;
         }
         mState.set(state);
-        mTmpState.set(state, true /* copySources */);
+        mLastDispachedState.set(state, true /* copySources */);
         applyLocalVisibilityOverride();
         mViewRoot.notifyInsetsChanged();
-        if (!mState.equals(mTmpState)) {
+        if (!mState.equals(mLastDispachedState)) {
             sendStateToWindowManager();
         }
         return true;
@@ -419,6 +425,9 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
             }
         }
 
+        int[] showTypes = new int[1];
+        int[] hideTypes = new int[1];
+
         // Ensure to update all existing source consumers
         for (int i = mSourceConsumers.size() - 1; i >= 0; i--) {
             final InsetsSourceConsumer consumer = mSourceConsumers.valueAt(i);
@@ -426,15 +435,23 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
 
             // control may be null, but we still need to update the control to null if it got
             // revoked.
-            consumer.setControl(control);
+            consumer.setControl(control, showTypes, hideTypes);
         }
 
         // Ensure to create source consumers if not available yet.
         for (int i = mTmpControlArray.size() - 1; i >= 0; i--) {
             final InsetsSourceControl control = mTmpControlArray.valueAt(i);
-            getSourceConsumer(control.getType()).setControl(control);
+            InsetsSourceConsumer consumer = getSourceConsumer(control.getType());
+            consumer.setControl(control, showTypes, hideTypes);
+
         }
         mTmpControlArray.clear();
+        if (showTypes[0] != 0) {
+            applyAnimation(showTypes[0], true /* show */, false /* fromIme */);
+        }
+        if (hideTypes[0] != 0) {
+            applyAnimation(hideTypes[0], false /* show */, false /* fromIme */);
+        }
     }
 
     @Override
@@ -465,7 +482,7 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
             @InternalInsetsType int internalType = internalTypes.valueAt(i);
             @AnimationType int animationType = getAnimationType(internalType);
             InsetsSourceConsumer consumer = getSourceConsumer(internalType);
-            if (mState.getSource(internalType).isVisible() && animationType == ANIMATION_TYPE_NONE
+            if (consumer.isRequestedVisible() && animationType == ANIMATION_TYPE_NONE
                     || animationType == ANIMATION_TYPE_SHOW) {
                 // no-op: already shown or animating in (because window visibility is
                 // applied before starting animation).
@@ -488,7 +505,7 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
             @InternalInsetsType int internalType = internalTypes.valueAt(i);
             @AnimationType int animationType = getAnimationType(internalType);
             InsetsSourceConsumer consumer = getSourceConsumer(internalType);
-            if (!mState.getSource(internalType).isVisible() && animationType == ANIMATION_TYPE_NONE
+            if (!consumer.isRequestedVisible() && animationType == ANIMATION_TYPE_NONE
                     || animationType == ANIMATION_TYPE_HIDE) {
                 // no-op: already hidden or animating out.
                 continue;
@@ -535,7 +552,7 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
         final SparseArray<InsetsSourceControl> controls = new SparseArray<>();
 
         Pair<Integer, Boolean> typesReadyPair = collectSourceControls(
-                fromIme, internalTypes, controls);
+                fromIme, internalTypes, controls, animationType);
         int typesReady = typesReadyPair.first;
         boolean imeReady = typesReadyPair.second;
         if (!imeReady) {
@@ -562,17 +579,20 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
      * @return Pair of (types ready to animate, IME ready to animate).
      */
     private Pair<Integer, Boolean> collectSourceControls(boolean fromIme,
-            ArraySet<Integer> internalTypes, SparseArray<InsetsSourceControl> controls) {
+            ArraySet<Integer> internalTypes, SparseArray<InsetsSourceControl> controls,
+            @AnimationType int animationType) {
         int typesReady = 0;
         boolean imeReady = true;
         for (int i = internalTypes.size() - 1; i >= 0; i--) {
             InsetsSourceConsumer consumer = getSourceConsumer(internalTypes.valueAt(i));
-            boolean setVisible = !consumer.isRequestedVisible();
-            if (setVisible) {
+            boolean show = animationType == ANIMATION_TYPE_SHOW
+                    || animationType == ANIMATION_TYPE_USER;
+            boolean canRun = false;
+            if (show) {
                 // Show request
                 switch(consumer.requestShow(fromIme)) {
                     case ShowResult.SHOW_IMMEDIATELY:
-                        typesReady |= InsetsState.toPublicType(consumer.getType());
+                        canRun = true;
                         break;
                     case ShowResult.IME_SHOW_DELAYED:
                         imeReady = false;
@@ -589,11 +609,22 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
                 if (!fromIme) {
                     consumer.notifyHidden();
                 }
-                typesReady |= InsetsState.toPublicType(consumer.getType());
+                canRun = true;
+            }
+            if (!canRun) {
+                continue;
             }
             final InsetsSourceControl control = consumer.getControl();
             if (control != null) {
                 controls.put(consumer.getType(), control);
+                typesReady |= toPublicType(consumer.getType());
+            } else if (animationType == ANIMATION_TYPE_SHOW) {
+
+                // We don't have a control at the moment. However, we still want to update requested
+                // visibility state such that in case we get control, we can apply show animation.
+                consumer.show(fromIme);
+            } else if (animationType == ANIMATION_TYPE_HIDE) {
+                consumer.hide();
             }
         }
         return new Pair<>(typesReady, imeReady);
@@ -808,7 +839,7 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
     private void showDirectly(@InsetsType int types) {
         final ArraySet<Integer> internalTypes = InsetsState.toInternalType(types);
         for (int i = internalTypes.size() - 1; i >= 0; i--) {
-            getSourceConsumer(internalTypes.valueAt(i)).show();
+            getSourceConsumer(internalTypes.valueAt(i)).show(false /* fromIme */);
         }
     }
 
@@ -840,6 +871,9 @@ public class InsetsController implements WindowInsetsController, InsetsAnimation
             @Override
             public boolean onPreDraw() {
                 mViewRoot.mView.getViewTreeObserver().removeOnPreDrawListener(this);
+                if (controller.isCancelled()) {
+                    return true;
+                }
                 mViewRoot.mView.dispatchWindowInsetsAnimationStart(animation, bounds);
                 listener.onReady(controller, types);
                 return true;
diff --git a/core/java/android/view/InsetsSourceConsumer.java b/core/java/android/view/InsetsSourceConsumer.java
index 9901d053184c..e6497c00c8dd 100644
--- a/core/java/android/view/InsetsSourceConsumer.java
+++ b/core/java/android/view/InsetsSourceConsumer.java
@@ -17,11 +17,14 @@
 package android.view;
 
 import static android.view.InsetsController.ANIMATION_TYPE_NONE;
+import static android.view.InsetsState.toPublicType;
 
 import android.annotation.IntDef;
 import android.annotation.Nullable;
+import android.util.MutableShort;
 import android.view.InsetsState.InternalInsetsType;
 import android.view.SurfaceControl.Transaction;
+import android.view.WindowInsets.Type.InsetsType;
 
 import com.android.internal.annotations.VisibleForTesting;
 
@@ -71,18 +74,48 @@ public class InsetsSourceConsumer {
         mRequestedVisible = InsetsState.getDefaultVisibility(type);
     }
 
-    public void setControl(@Nullable InsetsSourceControl control) {
+    /**
+     * Updates the control delivered from the server.
+
+     * @param showTypes An integer array with a single entry that determines which types a show
+     *                  animation should be run after setting the control.
+     * @param hideTypes An integer array with a single entry that determines which types a hide
+     *                  animation should be run after setting the control.
+     */
+    public void setControl(@Nullable InsetsSourceControl control,
+            @InsetsType int[] showTypes, @InsetsType int[] hideTypes) {
         if (mSourceControl == control) {
             return;
         }
         mSourceControl = control;
-        applyHiddenToControl();
-        if (applyLocalVisibilityOverride()) {
-            mController.notifyVisibilityChanged();
-        }
+
+        // We are loosing control
         if (mSourceControl == null) {
             mController.notifyControlRevoked(this);
+
+            // Restore server visibility.
+            mState.getSource(getType()).setVisible(
+                    mController.getLastDispatchedState().getSource(getType()).isVisible());
+            applyLocalVisibilityOverride();
+            return;
+        }
+
+        // We are gaining control, and need to run an animation since previous state didn't match
+        if (mRequestedVisible != mState.getSource(mType).isVisible()) {
+            if (mRequestedVisible) {
+                showTypes[0] |= toPublicType(getType());
+            } else {
+                hideTypes[0] |= toPublicType(getType());
+            }
+            return;
+        }
+
+        // We are gaining control, but don't need to run an animation. However make sure that the
+        // leash visibility is still up to date.
+        if (applyLocalVisibilityOverride()) {
+            mController.notifyVisibilityChanged();
         }
+        applyHiddenToControl();
     }
 
     @VisibleForTesting
@@ -95,7 +128,7 @@ public class InsetsSourceConsumer {
     }
 
     @VisibleForTesting
-    public void show() {
+    public void show(boolean fromIme) {
         setRequestedVisible(true);
     }
 
@@ -172,17 +205,13 @@ public class InsetsSourceConsumer {
      * the moment.
      */
     private void setRequestedVisible(boolean requestedVisible) {
-        if (mRequestedVisible == requestedVisible) {
-            return;
-        }
         mRequestedVisible = requestedVisible;
-        applyLocalVisibilityOverride();
-        mController.notifyVisibilityChanged();
+        if (applyLocalVisibilityOverride()) {
+            mController.notifyVisibilityChanged();
+        }
     }
 
     private void applyHiddenToControl() {
-
-        // TODO: Handle case properly when animation is running already (it shouldn't!)
         if (mSourceControl == null || mSourceControl.getLeash() == null) {
             return;
         }
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index 51ea30b41741..2b4b71f01aa5 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -4887,8 +4887,14 @@ public final class ViewRootImpl implements ViewParent,
                     break;
                 case MSG_INSETS_CONTROL_CHANGED: {
                     SomeArgs args = (SomeArgs) msg.obj;
-                    mInsetsController.onControlsChanged((InsetsSourceControl[]) args.arg2);
+
+                    // Deliver state change before control change, such that:
+                    // a) When gaining control, controller can compare with server state to evaluate
+                    // whether it needs to run animation.
+                    // b) When loosing control, controller can restore server state by taking last
+                    // dispatched state as truth.
                     mInsetsController.onStateChanged((InsetsState) args.arg1);
+                    mInsetsController.onControlsChanged((InsetsSourceControl[]) args.arg2);
                     break;
                 }
                 case MSG_SHOW_INSETS: {
diff --git a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
index 169716f99dea..bdb802195d8b 100644
--- a/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
+++ b/core/tests/coretests/src/android/view/InsetsAnimationControlImplTest.java
@@ -109,13 +109,14 @@ public class InsetsAnimationControlImplTest {
         InsetsSourceConsumer topConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, mInsetsState,
                 () -> mMockTransaction, mMockController);
         topConsumer.setControl(
-                new InsetsSourceControl(ITYPE_STATUS_BAR, mTopLeash, new Point(0, 0)));
+                new InsetsSourceControl(ITYPE_STATUS_BAR, mTopLeash, new Point(0, 0)),
+                new int[1], new int[1]);
 
         InsetsSourceConsumer navConsumer = new InsetsSourceConsumer(ITYPE_NAVIGATION_BAR,
                 mInsetsState, () -> mMockTransaction, mMockController);
-        navConsumer.hide();
         navConsumer.setControl(new InsetsSourceControl(ITYPE_NAVIGATION_BAR, mNavLeash,
-                new Point(400, 0)));
+                new Point(400, 0)), new int[1], new int[1]);
+        navConsumer.hide();
 
         SparseArray<InsetsSourceControl> controls = new SparseArray<>();
         controls.put(ITYPE_STATUS_BAR, topConsumer.getControl());
diff --git a/core/tests/coretests/src/android/view/InsetsControllerTest.java b/core/tests/coretests/src/android/view/InsetsControllerTest.java
index e68c4b8d2ab3..75fd2183946f 100644
--- a/core/tests/coretests/src/android/view/InsetsControllerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsControllerTest.java
@@ -121,9 +121,20 @@ public class InsetsControllerTest {
                 if (type == ITYPE_IME) {
                     return new InsetsSourceConsumer(type, controller.getState(),
                             Transaction::new, controller) {
+
+                        private boolean mImeRequestedShow;
+
+                        @Override
+                        public void show(boolean fromIme) {
+                            super.show(fromIme);
+                            if (fromIme) {
+                                mImeRequestedShow = true;
+                            }
+                        }
+
                         @Override
                         public int requestShow(boolean fromController) {
-                            if (fromController) {
+                            if (fromController || mImeRequestedShow) {
                                 return SHOW_IMMEDIATELY;
                             } else {
                                 return IME_SHOW_DELAYED;
@@ -399,6 +410,84 @@ public class InsetsControllerTest {
     }
 
     @Test
+    public void testRestoreStartsAnimation() {
+        InsetsSourceControl control =
+                new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point());
+        mController.onControlsChanged(new InsetsSourceControl[]{control});
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            mController.hide(Type.statusBars());
+            mController.cancelExistingAnimation();
+            assertFalse(mController.getSourceConsumer(ITYPE_STATUS_BAR).isRequestedVisible());
+            assertFalse(mController.getState().getSource(ITYPE_STATUS_BAR).isVisible());
+
+            // Loosing control
+            InsetsState state = new InsetsState(mController.getState());
+            state.setSourceVisible(ITYPE_STATUS_BAR, true);
+            mController.onStateChanged(state);
+            mController.onControlsChanged(new InsetsSourceControl[0]);
+            assertFalse(mController.getSourceConsumer(ITYPE_STATUS_BAR).isRequestedVisible());
+            assertTrue(mController.getState().getSource(ITYPE_STATUS_BAR).isVisible());
+
+            // Gaining control
+            mController.onControlsChanged(new InsetsSourceControl[]{control});
+            assertEquals(ANIMATION_TYPE_HIDE, mController.getAnimationType(ITYPE_STATUS_BAR));
+            mController.cancelExistingAnimation();
+            assertFalse(mController.getSourceConsumer(ITYPE_STATUS_BAR).isRequestedVisible());
+            assertFalse(mController.getState().getSource(ITYPE_STATUS_BAR).isVisible());
+        });
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
+
+    @Test
+    public void testStartImeAnimationAfterGettingControl() {
+        InsetsSourceControl control =
+                new InsetsSourceControl(ITYPE_IME, mLeash, new Point());
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+
+            mController.show(ime());
+            assertFalse(mController.getState().getSource(ITYPE_IME).isVisible());
+
+            // Pretend IME is calling
+            mController.show(ime(), true /* fromIme */);
+
+            // Gaining control shortly after
+            mController.onControlsChanged(new InsetsSourceControl[]{control});
+
+            assertEquals(ANIMATION_TYPE_SHOW, mController.getAnimationType(ITYPE_IME));
+            mController.cancelExistingAnimation();
+            assertTrue(mController.getSourceConsumer(ITYPE_IME).isRequestedVisible());
+            assertTrue(mController.getState().getSource(ITYPE_IME).isVisible());
+        });
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
+
+    @Test
+    public void testStartImeAnimationAfterGettingControl_imeLater() {
+        InsetsSourceControl control =
+                new InsetsSourceControl(ITYPE_IME, mLeash, new Point());
+
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+
+            mController.show(ime());
+            assertFalse(mController.getState().getSource(ITYPE_IME).isVisible());
+
+            // Gaining control shortly after
+            mController.onControlsChanged(new InsetsSourceControl[]{control});
+
+            // Pretend IME is calling
+            mController.show(ime(), true /* fromIme */);
+
+            assertEquals(ANIMATION_TYPE_SHOW, mController.getAnimationType(ITYPE_IME));
+            mController.cancelExistingAnimation();
+            assertTrue(mController.getSourceConsumer(ITYPE_IME).isRequestedVisible());
+            assertTrue(mController.getState().getSource(ITYPE_IME).isVisible());
+        });
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync();
+    }
+
+    @Test
     public void testAnimationEndState_controller() throws Exception {
         InsetsSourceControl control =
                 new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point());
diff --git a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
index 492c03653990..5e9e2f0065ed 100644
--- a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
@@ -18,6 +18,8 @@ package android.view;
 
 import static android.view.InsetsState.ITYPE_STATUS_BAR;
 
+import static android.view.WindowInsets.Type.statusBars;
+import static junit.framework.Assert.assertEquals;
 import static junit.framework.TestCase.assertFalse;
 import static junit.framework.TestCase.assertTrue;
 
@@ -90,34 +92,44 @@ public class InsetsSourceConsumerTest {
         });
         instrumentation.waitForIdleSync();
 
-        mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()));
+        mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()),
+                new int[1], new int[1]);
     }
 
     @Test
     public void testHide() {
-        mConsumer.hide();
-        assertFalse("Consumer should not be visible", mConsumer.isRequestedVisible());
-        verify(mSpyInsetsSource).setVisible(eq(false));
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            mConsumer.hide();
+            assertFalse("Consumer should not be visible", mConsumer.isRequestedVisible());
+            verify(mSpyInsetsSource).setVisible(eq(false));
+        });
+
     }
 
     @Test
     public void testShow() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            // Insets source starts out visible
+            mConsumer.hide();
+            mConsumer.show(false /* fromIme */);
+            assertTrue("Consumer should be visible", mConsumer.isRequestedVisible());
+            verify(mSpyInsetsSource).setVisible(eq(false));
+            verify(mSpyInsetsSource).setVisible(eq(true));
+        });
 
-        // Insets source starts out visible
-        mConsumer.hide();
-        mConsumer.show();
-        assertTrue("Consumer should be visible", mConsumer.isRequestedVisible());
-        verify(mSpyInsetsSource).setVisible(eq(false));
-        verify(mSpyInsetsSource).setVisible(eq(true));
     }
 
     @Test
     public void testRestore() {
-        mConsumer.setControl(null);
-        reset(mMockTransaction);
-        mConsumer.hide();
-        verifyZeroInteractions(mMockTransaction);
-        mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()));
-        verify(mMockTransaction).hide(eq(mLeash));
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            mConsumer.setControl(null, new int[1], new int[1]);
+            reset(mMockTransaction);
+            mConsumer.hide();
+            verifyZeroInteractions(mMockTransaction);
+            int[] hideTypes = new int[1];
+            mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()),
+                    new int[1], hideTypes);
+            assertEquals(statusBars(), hideTypes[0]);
+        });
     }
 }