Tests: waiting for nav bar to match nav mode

Change-Id: Ib17fce4577e1e2c585108f0e95a98fad710e5eb2
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 90763b8..3b35c86 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -159,13 +159,20 @@
                     }
 
                     for (int i = 0; i != 100; ++i) {
-                        if (mLauncher.getNavigationModel() == expectedMode) {
-                            Thread.sleep(5000);
-                            return;
-                        }
+                        if (mLauncher.getNavigationModel() == expectedMode) break;
                         Thread.sleep(100);
                     }
-                    Assert.fail("Couldn't switch to " + overlayPackage);
+                    Assert.assertTrue("Couldn't switch to " + overlayPackage,
+                            mLauncher.getNavigationModel() == expectedMode);
+
+                    for (int i = 0; i != 100; ++i) {
+                        if (mLauncher.getNavigationModeMismatchError() == null) break;
+                        Thread.sleep(100);
+                    }
+                    final String error = mLauncher.getNavigationModeMismatchError();
+                    Assert.assertTrue("Switching nav mode: " + error, error == null);
+
+                    Thread.sleep(5000);
                 }
 
                 private void setOverlayPackageEnabled(String overlayPackage, boolean enable)
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 22b04fb..9336806 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -315,18 +315,26 @@
         mExpectedRotation = expectedRotation;
     }
 
-    private UiObject2 verifyContainerType(ContainerType containerType) {
-        assertEquals("Unexpected display rotation",
-                mExpectedRotation, mDevice.getDisplayRotation());
+    public String getNavigationModeMismatchError() {
         final NavigationModel navigationModel = getNavigationModel();
         final boolean hasRecentsButton = hasSystemUiObject("recent_apps");
         final boolean hasHomeButton = hasSystemUiObject("home");
-        assertTrue("Presence of recents button doesn't match the interaction mode, mode="
-                        + navigationModel.name() + ", hasRecents=" + hasRecentsButton,
-                (navigationModel == NavigationModel.THREE_BUTTON) == hasRecentsButton);
-        assertTrue("Presence of home button doesn't match the interaction mode, mode="
-                        + navigationModel.name() + ", hasHome=" + hasHomeButton,
-                (navigationModel != NavigationModel.ZERO_BUTTON) == hasHomeButton);
+        if ((navigationModel == NavigationModel.THREE_BUTTON) != hasRecentsButton) {
+            return "Presence of recents button doesn't match the interaction mode, mode="
+                    + navigationModel.name() + ", hasRecents=" + hasRecentsButton;
+        }
+        if ((navigationModel != NavigationModel.ZERO_BUTTON) != hasHomeButton) {
+            return "Presence of home button doesn't match the interaction mode, mode="
+                    + navigationModel.name() + ", hasHome=" + hasHomeButton;
+        }
+        return null;
+    }
+
+    private UiObject2 verifyContainerType(ContainerType containerType) {
+        assertEquals("Unexpected display rotation",
+                mExpectedRotation, mDevice.getDisplayRotation());
+        final String error = getNavigationModeMismatchError();
+        assertTrue(error, error == null);
         log("verifyContainerType: " + containerType);
 
         try (Closable c = addContextLayer(