Reland:[ToolbarMVC] Hide MenuButton interactions in toolbar.top behind MenuButtonCoordinator

Most method calls are translated 1:1 by forwarding the call to the equivalent
MenuButton or View method. A notable exception is handling of the tab switcher
animation, which currently requires a view in order to perform canvas
translation before drawing the button in its animated state. In this case,
we provide a direct accessor for the view.

Reland note: the NPE was due to the MenuButtonCoordinator not being set until
right after inflation finishes, which is when ToolbarTablet was trying to use
it.

Bug: 1086676
Change-Id: I60819a54bc4ad6d84fcb5b932d04798ba6561bce
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2391643Reviewed-by: Matthew Jones <mdjones@chromium.org>
Commit-Queue: Patrick Noland <pnoland@chromium.org>
Cr-Commit-Position: refs/heads/master@{#804408}

chrome/android/java/src/org/chromium/chrome/browser/toolbar/ToolbarManager.java

@@ -80,7 +80,6 @@ import org.chromium.chrome.browser.toolbar.bottom.BottomTabSwitcherActionMenuCoo
 import org.chromium.chrome.browser.toolbar.bottom.BottomToolbarConfiguration;
 import org.chromium.chrome.browser.toolbar.bottom.BottomToolbarVariationManager;
 import org.chromium.chrome.browser.toolbar.load_progress.LoadProgressCoordinator;
-import org.chromium.chrome.browser.toolbar.menu_button.MenuButton;
 import org.chromium.chrome.browser.toolbar.menu_button.MenuButtonCoordinator;
 import org.chromium.chrome.browser.toolbar.top.ActionModeController;
 import org.chromium.chrome.browser.toolbar.top.ActionModeController.ActionBarDelegate;
@@ -307,13 +306,12 @@ public class ToolbarManager implements UrlFocusChangeListener, ThemeColorObserve
         assert controlsVisibilityDelegate != null;
         mControlsVisibilityDelegate = controlsVisibilityDelegate;
 
-        MenuButton menuButton = mActivity.findViewById(R.id.menu_button_wrapper);
         mMenuButtonCoordinator = new MenuButtonCoordinator(appMenuCoordinatorSupplier,
                 mControlsVisibilityDelegate, mActivity,
                 (focus, type)
                         -> setUrlBarFocus(focus, type),
                 mActivity.getCompositorViewHolder()::requestFocus, shouldShowUpdateBadge,
-                mActivity::isInOverviewMode, menuButton);
+                mActivity::isInOverviewMode);
 
         mToolbar = new TopToolbarCoordinator(controlContainer, mActivity.findViewById(R.id.toolbar),
                 identityDiscController, mLocationBarModel, mToolbarTabController,

chrome/android/java/src/org/chromium/chrome/browser/toolbar/menu_button/MenuButtonCoordinator.java

@@ -5,10 +5,12 @@
 package org.chromium.chrome.browser.toolbar.menu_button;
 
 import android.app.Activity;
+import android.content.res.ColorStateList;
 
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
+import org.chromium.base.ApiCompatibilityUtils;
 import org.chromium.base.Callback;
 import org.chromium.base.metrics.RecordUserAction;
 import org.chromium.base.supplier.ObservableSupplier;
@@ -65,13 +67,12 @@ public class MenuButtonCoordinator implements AppMenuObserver {
      * @param shouldShowAppUpdateBadge Whether the app menu update badge should be shown if there is
      *         a pending update.
      * @param isInOverviewModeSupplier Supplier of overview mode state.
-     * @param menuButton View that presents the MenuButton.
      */
     public MenuButtonCoordinator(ObservableSupplier<AppMenuCoordinator> appMenuCoordinatorSupplier,
             BrowserStateBrowserControlsVisibilityDelegate controlsVisibilityDelegate,
             Activity activity, SetFocusFunction setUrlBarFocusFunction,
             Runnable requestRenderRunnable, boolean shouldShowAppUpdateBadge,
-            Supplier<Boolean> isInOverviewModeSupplier, MenuButton menuButton) {
+            Supplier<Boolean> isInOverviewModeSupplier) {
         mControlsVisibilityDelegate = controlsVisibilityDelegate;
         mActivity = activity;
         mSetUrlBarFocusFunction = setUrlBarFocusFunction;
@@ -81,7 +82,7 @@ public class MenuButtonCoordinator implements AppMenuObserver {
         mRequestRenderRunnable = requestRenderRunnable;
         mShouldShowAppUpdateBadge = shouldShowAppUpdateBadge;
         mIsInOverviewModeSupplier = isInOverviewModeSupplier;
-        mMenuButton = menuButton;
+        mMenuButton = mActivity.findViewById(R.id.menu_button_wrapper);
         mAppMenuButtonHelperSupplier = new ObservableSupplierImpl<>();
     }
 
@@ -108,6 +109,24 @@ public class MenuButtonCoordinator implements AppMenuObserver {
         }
     }
 
+    /**
+     * Get the underlying MenuButton view. Present for legacy reasons only; don't add new usages.
+     */
+    @Deprecated
+    public MenuButton getMenuButton() {
+        return mMenuButton;
+    }
+
+    /**
+     * Set the tint list on the underlying MenuButton view. Present for legacy reasons only; don't
+     * add new usages.
+     */
+    @Deprecated
+    public void setImageTintList(ColorStateList colorStateList) {
+        if (mMenuButton == null) return;
+        ApiCompatibilityUtils.setImageTintList(mMenuButton.getImageButton(), colorStateList);
+    }
+
     public void destroy() {
         if (mAppMenuButtonHelper != null) {
             mAppMenuHandler.removeObserver(this);
@@ -150,6 +169,23 @@ public class MenuButtonCoordinator implements AppMenuObserver {
         mMenuButton.setAppMenuUpdateBadgeSuppressed(isSuppressed);
     }
 
+    /**
+     * Set the visibility of the MenuButton controlled by this coordinator.
+     * @param visibility Visibility state flag, e.g. GONE or VISIBLE.
+     */
+    public void setVisibility(int visibility) {
+        if (mMenuButton == null) return;
+        mMenuButton.setVisibility(visibility);
+    }
+
+    /**
+     * @see MenuButton#setMenuButtonHighlightDrawable().
+     */
+    public void setMenuButtonHighlightDrawable() {
+        if (mMenuButton == null) return;
+        mMenuButton.setMenuButtonHighlightDrawable();
+    }
+
     @Override
     public void onMenuVisibilityChanged(boolean isVisible) {
         if (isVisible) {

chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/CustomTabToolbar.java

@@ -66,6 +66,7 @@ import org.chromium.chrome.browser.toolbar.IncognitoStateProvider;
 import org.chromium.chrome.browser.toolbar.ToolbarColors;
 import org.chromium.chrome.browser.toolbar.ToolbarDataProvider;
 import org.chromium.chrome.browser.toolbar.ToolbarTabController;
+import org.chromium.chrome.browser.toolbar.menu_button.MenuButtonCoordinator;
 import org.chromium.chrome.browser.ui.native_page.NativePage;
 import org.chromium.components.browser_ui.styles.ChromeColors;
 import org.chromium.components.browser_ui.widget.TintedDrawable;
@@ -147,7 +148,6 @@ public class CustomTabToolbar extends ToolbarLayout implements View.OnLongClickL
     private ImageButton mSecurityButton;
     private LinearLayout mCustomActionButtons;
     private ImageButton mCloseButton;
-    private ImageButton mMenuButton;
 
     // Whether dark tint should be applied to icons and text.
     private boolean mUseDarkColors;
@@ -205,14 +205,14 @@ public class CustomTabToolbar extends ToolbarLayout implements View.OnLongClickL
         mCustomActionButtons = findViewById(R.id.action_buttons);
         mCloseButton = findViewById(R.id.close_button);
         mCloseButton.setOnLongClickListener(this);
-        mMenuButton = findViewById(R.id.menu_button);
         mAnimDelegate = new CustomTabToolbarAnimationDelegate(
                 mSecurityButton, mTitleUrlContainer, R.dimen.location_bar_icon_width);
     }
 
     @Override
-    void initialize(ToolbarDataProvider toolbarDataProvider, ToolbarTabController tabController) {
-        super.initialize(toolbarDataProvider, tabController);
+    void initialize(ToolbarDataProvider toolbarDataProvider, ToolbarTabController tabController,
+            MenuButtonCoordinator menuButtonCoordinator) {
+        super.initialize(toolbarDataProvider, tabController, menuButtonCoordinator);
         mLocationBar.setToolbarDataProvider(toolbarDataProvider);
         mLocationBar.updateVisualsForState();
     }
@@ -373,10 +373,8 @@ public class CustomTabToolbar extends ToolbarLayout implements View.OnLongClickL
     }
 
     private void updateButtonsTint() {
-        if (getMenuButton() != null) {
-            ApiCompatibilityUtils.setImageTintList(
-                    getMenuButton(), mUseDarkColors ? mDarkModeTint : mLightModeTint);
-        }
+        getMenuButtonCoordinator().setImageTintList(
+                mUseDarkColors ? mDarkModeTint : mLightModeTint);
         updateButtonTint(mCloseButton);
         int numCustomActionButtons = mCustomActionButtons.getChildCount();
         for (int i = 0; i < numCustomActionButtons; i++) {
@@ -561,21 +559,9 @@ public class CustomTabToolbar extends ToolbarLayout implements View.OnLongClickL
         return url;
     }
 
-    @Override
-    View getMenuButtonWrapper() {
-        // This class has no menu button wrapper, so return the menu button instead.
-        return getMenuButton();
-    }
-
-    @Override
-    ImageButton getMenuButton() {
-        return mMenuButton;
-    }
-
     @Override
     void onMenuButtonDisabled() {
         super.onMenuButtonDisabled();
-        mMenuButton = null;
         // In addition to removing the menu button, we also need to remove the margin on the custom
         // action button.
         ViewGroup.MarginLayoutParams p =

chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/ToolbarLayout.java

@@ -47,7 +47,6 @@ import org.chromium.chrome.browser.toolbar.ToolbarColors;
 import org.chromium.chrome.browser.toolbar.ToolbarDataProvider;
 import org.chromium.chrome.browser.toolbar.ToolbarProgressBar;
 import org.chromium.chrome.browser.toolbar.ToolbarTabController;
-import org.chromium.chrome.browser.toolbar.menu_button.MenuButton;
 import org.chromium.chrome.browser.toolbar.menu_button.MenuButtonCoordinator;
 import org.chromium.chrome.browser.toolbar.top.TopToolbarCoordinator.UrlExpansionObserver;
 import org.chromium.chrome.browser.ui.appmenu.AppMenuButtonHelper;
@@ -68,11 +67,6 @@ public abstract class ToolbarLayout
             new ObserverList<>();
     private final int[] mTempPosition = new int[2];
 
-    /**
-     * The app menu button.
-     */
-    private MenuButton mMenuButtonWrapper;
-
     private final ColorStateList mDefaultTint;
 
     private ToolbarDataProvider mToolbarDataProvider;
@@ -89,6 +83,8 @@ public abstract class ToolbarLayout
     private boolean mFindInPageToolbarShowing;
 
     private ThemeColorProvider mThemeColorProvider;
+    private MenuButtonCoordinator mMenuButtonCoordinator;
+    private AppMenuButtonHelper mAppMenuButtonHelper;
 
     /**
      * Basic constructor for {@link ToolbarLayout}.
@@ -116,25 +112,21 @@ public abstract class ToolbarLayout
      * Initialize the external dependencies required for view interaction.
      * @param toolbarDataProvider The provider for toolbar data.
      * @param tabController       The controller that handles interactions with the tab.
+     * @param menuButtonCoordinator Coordinator for interacting with the MenuButton.
      */
-    void initialize(ToolbarDataProvider toolbarDataProvider, ToolbarTabController tabController) {
+    @CallSuper
+    void initialize(ToolbarDataProvider toolbarDataProvider, ToolbarTabController tabController,
+            MenuButtonCoordinator menuButtonCoordinator) {
         mToolbarDataProvider = toolbarDataProvider;
         mToolbarTabController = tabController;
+        mMenuButtonCoordinator = menuButtonCoordinator;
     }
 
     /**
      * @param appMenuButtonHelper The helper for managing menu button interactions.
      */
     void setAppMenuButtonHelper(AppMenuButtonHelper appMenuButtonHelper) {
-        if (mMenuButtonWrapper != null) {
-            mMenuButtonWrapper.setAppMenuButtonHelper(appMenuButtonHelper);
-        } else {
-            final ImageButton menuButton = getMenuButton();
-            if (menuButton != null) {
-                menuButton.setOnTouchListener(appMenuButtonHelper);
-                menuButton.setAccessibilityDelegate(appMenuButtonHelper.getAccessibilityDelegate());
-            }
-        }
+        mAppMenuButtonHelper = appMenuButtonHelper;
     }
 
     /**
@@ -224,8 +216,6 @@ public abstract class ToolbarLayout
     protected void onFinishInflate() {
         super.onFinishInflate();
 
-        mMenuButtonWrapper = findViewById(R.id.menu_button_wrapper);
-
         // Initialize the provider to an empty version to avoid null checking everywhere.
         mToolbarDataProvider = new ToolbarDataProvider() {
             @Override
@@ -313,12 +303,6 @@ public abstract class ToolbarLayout
                 return 0;
             }
         };
-
-        // Set menu button background in case it was previously called before inflation
-        // finished (i.e. mMenuButtonWrapper == null)
-        if (mMenuButtonWrapper != null) {
-            mMenuButtonWrapper.setMenuButtonHighlightDrawable();
-        }
     }
 
     @Override
@@ -371,29 +355,21 @@ public abstract class ToolbarLayout
     /**
      * @return The view containing the menu button and menu button badge.
      */
-    View getMenuButtonWrapper() {
-        return mMenuButtonWrapper;
+    MenuButtonCoordinator getMenuButtonCoordinator() {
+        return mMenuButtonCoordinator;
     }
 
     @VisibleForTesting
-    void setMenuButtonWrapperForTesting(MenuButton menuButton) {
-        mMenuButtonWrapper = menuButton;
+    void setMenuButtonCoordinatorForTesting(MenuButtonCoordinator menuButtonCoordinator) {
+        mMenuButtonCoordinator = menuButtonCoordinator;
     }
 
     /**
      * @return The {@link ImageButton} containing the menu button.
      */
     ImageButton getMenuButton() {
-        if (mMenuButtonWrapper == null) return null;
-        return mMenuButtonWrapper.getImageButton();
-    }
-
-    /**
-     * @return The view containing the menu badge.
-     */
-    View getMenuBadge() {
-        if (mMenuButtonWrapper == null) return null;
-        return mMenuButtonWrapper.getMenuBadge();
+        if (mMenuButtonCoordinator.getMenuButton() == null) return null;
+        return mMenuButtonCoordinator.getMenuButton().getImageButton();
     }
 
     /**
@@ -411,8 +387,7 @@ public abstract class ToolbarLayout
      * @return The helper for menu button UI interactions.
      */
     AppMenuButtonHelper getMenuButtonHelper() {
-        if (mMenuButtonWrapper == null) return null;
-        return mMenuButtonWrapper.getAppMenuButtonHelper();
+        return mAppMenuButtonHelper;
     }
 
     /**
@@ -710,9 +685,7 @@ public abstract class ToolbarLayout
 
     boolean shouldIgnoreSwipeGesture() {
         if (mUrlHasFocus || mFindInPageToolbarShowing) return true;
-        if (mMenuButtonWrapper == null) return false;
-        final AppMenuButtonHelper appMenuButtonHelper = mMenuButtonWrapper.getAppMenuButtonHelper();
-        return appMenuButtonHelper != null && appMenuButtonHelper.isAppMenuActive();
+        return mAppMenuButtonHelper != null && mAppMenuButtonHelper.isAppMenuActive();
     }
 
     /**
@@ -873,15 +846,6 @@ public abstract class ToolbarLayout
         return null;
     }
 
-    /**
-     * Sets the menu button's background depending on whether or not we are highlighting and whether
-     * or not we are using light or dark assets.
-     */
-    void setMenuButtonHighlightDrawable() {
-        if (mMenuButtonWrapper == null) return;
-        mMenuButtonWrapper.setMenuButtonHighlightDrawable();
-    }
-
     /**
      * Sets the current TabModelSelector so the toolbar can pass it into buttons that need access to
      * it.

chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/ToolbarPhone.java

@@ -374,8 +374,8 @@ public class ToolbarPhone extends ToolbarLayout implements Invalidator.Client, O
 
             setLayoutTransition(null);
 
-            if (getMenuButtonWrapper() != null) {
-                getMenuButtonWrapper().setVisibility(View.VISIBLE);
+            if (getMenuButtonCoordinator() != null) {
+                getMenuButtonCoordinator().setVisibility(View.VISIBLE);
             }
 
             inflateTabSwitchingResources();
@@ -1318,7 +1318,7 @@ public class ToolbarPhone extends ToolbarLayout implements Invalidator.Client, O
         }
 
         // Draw the menu button if necessary.
-        final MenuButton menuButton = (MenuButton) getMenuButtonWrapper();
+        final MenuButton menuButton = getMenuButtonCoordinator().getMenuButton();
         if (menuButton != null) {
             canvas.save();
             translateCanvasToView(mToolbarButtonsContainer, menuButton, canvas);
@@ -2015,7 +2015,7 @@ public class ToolbarPhone extends ToolbarLayout implements Invalidator.Client, O
         float toolbarButtonTranslationX =
                 MathUtils.flipSignIf(URL_FOCUS_TOOLBAR_BUTTONS_TRANSLATION_X_DP, isRtl) * density;
 
-        final View menuButtonWrapper = getMenuButtonWrapper();
+        final View menuButtonWrapper = getMenuButtonCoordinator().getMenuButton();
         if (menuButtonWrapper != null) {
             animator = ObjectAnimator.ofFloat(
                     menuButtonWrapper, TRANSLATION_X, toolbarButtonTranslationX);
@@ -2070,7 +2070,7 @@ public class ToolbarPhone extends ToolbarLayout implements Invalidator.Client, O
         animator.setInterpolator(BakedBezierInterpolator.TRANSFORM_CURVE);
         animators.add(animator);
 
-        final View menuButtonWrapper = getMenuButtonWrapper();
+        final View menuButtonWrapper = getMenuButtonCoordinator().getMenuButton();
         if (menuButtonWrapper != null) {
             animator = ObjectAnimator.ofFloat(menuButtonWrapper, TRANSLATION_X, 0);
             animator.setDuration(URL_FOCUS_TOOLBAR_BUTTONS_DURATION_MS);
@@ -2535,10 +2535,8 @@ public class ToolbarPhone extends ToolbarLayout implements Invalidator.Client, O
             updateNtpTransitionAnimation();
         }
 
-        if (getMenuButtonWrapper() != null) {
-            setMenuButtonHighlightDrawable();
-            if (!mIsBottomToolbarVisible) getMenuButtonWrapper().setVisibility(View.VISIBLE);
-        }
+        getMenuButtonCoordinator().setMenuButtonHighlightDrawable();
+        if (!mIsBottomToolbarVisible) getMenuButtonCoordinator().setVisibility(View.VISIBLE);
 
         DrawableCompat.setTint(mLocationBarBackground,
                 isIncognito() ? Color.WHITE

chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/ToolbarTablet.java

@@ -43,6 +43,8 @@ import org.chromium.chrome.browser.toolbar.KeyboardNavigationListener;
 import org.chromium.chrome.browser.toolbar.TabCountProvider;
 import org.chromium.chrome.browser.toolbar.TabCountProvider.TabCountObserver;
 import org.chromium.chrome.browser.toolbar.ToolbarColors;
+import org.chromium.chrome.browser.toolbar.ToolbarDataProvider;
+import org.chromium.chrome.browser.toolbar.ToolbarTabController;
 import org.chromium.chrome.browser.toolbar.menu_button.MenuButtonCoordinator;
 import org.chromium.chrome.browser.util.ChromeAccessibilityUtil;
 import org.chromium.components.browser_ui.styles.ChromeColors;
@@ -137,15 +139,6 @@ public class ToolbarTablet extends ToolbarLayout
 
         mBookmarkButton = findViewById(R.id.bookmark_button);
 
-        final View menuButtonWrapper = getMenuButtonWrapper();
-        menuButtonWrapper.setVisibility(View.VISIBLE);
-
-        if (mAccessibilitySwitcherButton.getVisibility() == View.GONE
-                && menuButtonWrapper.getVisibility() == View.GONE) {
-            ViewCompat.setPaddingRelative((View) menuButtonWrapper.getParent(), 0, 0,
-                    getResources().getDimensionPixelSize(R.dimen.tablet_toolbar_end_padding), 0);
-        }
-
         mSaveOfflineButton = findViewById(R.id.save_offline_button);
 
         // Initialize values needed for showing/hiding toolbar buttons when the activity size
@@ -502,6 +495,13 @@ public class ToolbarTablet extends ToolbarLayout
         }
     }
 
+    @Override
+    void initialize(ToolbarDataProvider toolbarDataProvider, ToolbarTabController tabController,
+        MenuButtonCoordinator menuButtonCoordinator) {
+        super.initialize(toolbarDataProvider, tabController, menuButtonCoordinator);
+        menuButtonCoordinator.setVisibility(View.VISIBLE);
+    }
+
     @Override
     public void onTabCountChanged(int numberOfTabs, boolean isIncognito) {
         mAccessibilitySwitcherButton.setContentDescription(getResources().getQuantityString(

chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/TopToolbarCoordinator.java

@@ -136,7 +136,7 @@ public class TopToolbarCoordinator implements Toolbar {
         }
         controlContainer.setToolbar(this);
         HomepageManager.getInstance().addListener(mHomepageStateListener);
-        mToolbarLayout.initialize(toolbarDataProvider, tabController);
+        mToolbarLayout.initialize(toolbarDataProvider, tabController, menuButtonCoordinator);
 
         final MenuButton menuButtonWrapper = getMenuButtonWrapper();
         if (menuButtonWrapper != null) {
@@ -264,9 +264,7 @@ public class TopToolbarCoordinator implements Toolbar {
      * @return The wrapper for the browsing mode toolbar's menu button.
      */
     public MenuButton getMenuButtonWrapper() {
-        View menuButtonWrapper = mToolbarLayout.getMenuButtonWrapper();
-        if (menuButtonWrapper instanceof MenuButton) return (MenuButton) menuButtonWrapper;
-        return null;
+        return mMenuButtonCoordinator.getMenuButton();
     }
 
     /**

chrome/android/javatests/src/org/chromium/chrome/browser/toolbar/top/ToolbarPhoneTest.java

@@ -4,6 +4,7 @@
 
 package org.chromium.chrome.browser.toolbar.top;
 
+import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
@@ -15,6 +16,7 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
@@ -24,6 +26,7 @@ import org.chromium.chrome.R;
 import org.chromium.chrome.browser.app.ChromeActivity;
 import org.chromium.chrome.browser.flags.ChromeSwitches;
 import org.chromium.chrome.browser.toolbar.menu_button.MenuButton;
+import org.chromium.chrome.browser.toolbar.menu_button.MenuButtonCoordinator;
 import org.chromium.chrome.test.ChromeActivityTestRule;
 import org.chromium.chrome.test.ChromeJUnit4ClassRunner;
 import org.chromium.content_public.browser.test.util.TestThreadUtils;
@@ -40,9 +43,12 @@ public class ToolbarPhoneTest {
     public ChromeActivityTestRule<ChromeActivity> mActivityTestRule =
             new ChromeActivityTestRule<>(ChromeActivity.class);
 
+    @Mock
+    private MenuButtonCoordinator mMenuButtonCoordinator;
+
     private Canvas mCanvas = new Canvas();
     private ToolbarPhone mToolbar;
-    private MenuButton mMenuButtonSpy;
+    private MenuButton mMenuButton;
 
     @Before
     public void setUp() {
@@ -50,8 +56,9 @@ public class ToolbarPhoneTest {
 
         mActivityTestRule.startMainActivityOnBlankPage();
         mToolbar = mActivityTestRule.getActivity().findViewById(R.id.toolbar);
-        mMenuButtonSpy = Mockito.spy((MenuButton) mToolbar.getMenuButtonWrapper());
-        mToolbar.setMenuButtonWrapperForTesting(mMenuButtonSpy);
+        mMenuButton = Mockito.spy(mToolbar.findViewById(R.id.menu_button_wrapper));
+        mToolbar.setMenuButtonCoordinatorForTesting(mMenuButtonCoordinator);
+        doReturn(mMenuButton).when(mMenuButtonCoordinator).getMenuButton();
     }
 
     @Test
@@ -59,11 +66,11 @@ public class ToolbarPhoneTest {
     public void testDrawTabSwitcherAnimation_menuButtonDrawn() {
         TestThreadUtils.runOnUiThreadBlocking(() -> {
             mToolbar.drawTabSwitcherAnimationOverlay(mCanvas, 0);
-            verify(mMenuButtonSpy).drawTabSwitcherAnimationOverlay(mCanvas, 255);
+            verify(mMenuButton).drawTabSwitcherAnimationOverlay(mCanvas, 255);
 
             mToolbar.setTextureCaptureMode(true);
             mToolbar.draw(mCanvas);
-            verify(mMenuButtonSpy, times(2)).drawTabSwitcherAnimationOverlay(mCanvas, 255);
+            verify(mMenuButton, times(2)).drawTabSwitcherAnimationOverlay(mCanvas, 255);
             mToolbar.setTextureCaptureMode(false);
         });
     }

chrome/android/junit/src/org/chromium/chrome/browser/toolbar/menu_button/MenuButtonCoordinatorTest.java

@@ -71,10 +71,13 @@ public class MenuButtonCoordinatorTest {
         mAppMenuSupplier = new ObservableSupplierImpl<>();
         mMenuUiState = new UpdateMenuItemHelper.MenuUiState();
         doReturn(mMenuUiState).when(mUpdateMenuItemHelper).getUiState();
+        doReturn(mMenuButton)
+                .when(mActivity)
+                .findViewById(org.chromium.chrome.R.id.menu_button_wrapper);
 
         mMenuButtonCoordinator =
                 new MenuButtonCoordinator(mAppMenuSupplier, mControlsVisibilityDelegate, mActivity,
-                        mFocusFunction, mRequestRenderRunnable, true, () -> false, mMenuButton);
+                        mFocusFunction, mRequestRenderRunnable, true, () -> false);
     }
 
     @Test
@@ -147,7 +150,7 @@ public class MenuButtonCoordinatorTest {
     public void testAppMenuUpdateBadge_activityShouldNotShow() {
         MenuButtonCoordinator newCoordinator =
                 new MenuButtonCoordinator(mAppMenuSupplier, mControlsVisibilityDelegate, mActivity,
-                        mFocusFunction, mRequestRenderRunnable, false, () -> false, mMenuButton);
+                        mFocusFunction, mRequestRenderRunnable, false, () -> false);
 
         doReturn(true).when(mActivity).isDestroyed();
         newCoordinator.updateStateChanged();