Wait for callback in AppMenuTest#testDragHelper_ClickItem

This test has seem some flakiness on asserting the item clicked call
count. Instead of asserting immediately, wait for the callback.

BUG=1062236

Change-Id: I37b3c2d431656ebe3ec7545d5e08a27f7eb3ed16
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2106389Reviewed-by: Lijin Shen <lazzzis@chromium.org>
Commit-Queue: Theresa  <twellington@chromium.org>
Auto-Submit: Theresa  <twellington@chromium.org>
Cr-Commit-Position: refs/heads/master@{#751438}

chrome/browser/ui/android/appmenu/internal/java/src/org/chromium/chrome/browser/ui/appmenu/AppMenuTest.java

@@ -202,8 +202,7 @@ public class AppMenuTest extends DummyUiActivityTestCase {
                         -> AppMenuTestSupport.callOnItemClick(
                                 mAppMenuCoordinator, R.id.menu_item_three));
 
-        Assert.assertEquals("Item selected callback should have been called.", 1,
-                mDelegate.itemSelectedCallbackHelper.getCallCount());
+        mDelegate.itemSelectedCallbackHelper.waitForCallback(0);
         Assert.assertEquals("Incorrect id for last selected item.", R.id.menu_item_three,
                 mDelegate.lastSelectedItemId);
     }
@@ -228,8 +227,7 @@ public class AppMenuTest extends DummyUiActivityTestCase {
         TestThreadUtils.runOnUiThreadBlocking(
                 () -> mAppMenuHandler.getAppMenu().onItemClick(null, null, 0, 0));
 
-        Assert.assertEquals("Item selected callback should have been called.", 1,
-                mDelegate.itemSelectedCallbackHelper.getCallCount());
+        mDelegate.itemSelectedCallbackHelper.waitForCallback(0);
         Assert.assertEquals("Incorrect id for last selected item.", R.id.menu_item_one,
                 mDelegate.lastSelectedItemId);
     }
@@ -648,8 +646,7 @@ public class AppMenuTest extends DummyUiActivityTestCase {
         sendMotionEventToButtonHelper(
                 buttonHelper, mTestMenuButtonDelegate.getMenuButtonView(), upMotionEvent);
 
-        Assert.assertEquals("Item selected callback should have been called.", 1,
-                mDelegate.itemSelectedCallbackHelper.getCallCount());
+        mDelegate.itemSelectedCallbackHelper.waitForCallback(0);
         Assert.assertEquals("Incorrect id for last selected item.", R.id.menu_item_one,
                 mDelegate.lastSelectedItemId);
     }