Fixing CCT App Menu issues introduced by refactor

This CL undoes some of the changes that could impact CCT that was
introduced in the MenuButton refactor CL:
https://chromium-review.googlesource.com/c/chromium/src/+/1419203

Bug: 925271
Change-Id: Ifb9575f8cf52934981c12b1a6c5114993997cc5e
Reviewed-on: https://chromium-review.googlesource.com/c/1447834
Reviewed-by: Theresa <twellington@chromium.org>
Commit-Queue: Theresa <twellington@chromium.org>
Cr-Original-Commit-Position: refs/heads/master@{#628296}(cherry picked from commit a5003b9dbef6d18097a3e8d6067c73af92f606d5)
Reviewed-on: https://chromium-review.googlesource.com/c/1452459
Reviewed-by: Pedro Amaral <amaralp@chromium.org>
Cr-Commit-Position: refs/branch-heads/3683@{#167}
Cr-Branched-From: e51029943e0a38dd794b73caaf6373d5496ae783-refs/heads/master@{#625896}
diff --git a/chrome/android/java/src/org/chromium/chrome/browser/toolbar/MenuButton.java b/chrome/android/java/src/org/chromium/chrome/browser/toolbar/MenuButton.java
index 66c0750..8185485 100644
--- a/chrome/android/java/src/org/chromium/chrome/browser/toolbar/MenuButton.java
+++ b/chrome/android/java/src/org/chromium/chrome/browser/toolbar/MenuButton.java
@@ -27,7 +27,6 @@
 import org.chromium.chrome.browser.omaha.UpdateMenuItemHelper;
 import org.chromium.chrome.browser.omaha.UpdateMenuItemHelper.MenuButtonState;
 import org.chromium.chrome.browser.widget.PulseDrawable;
-import org.chromium.ui.UiUtils;
 import org.chromium.ui.interpolators.BakedBezierInterpolator;
 
 /**
@@ -108,6 +107,10 @@
                 ApiCompatibilityUtils.getDrawable(getResources(), drawable));
     }
 
+    /**
+     * Show the update badge on the app menu button.
+     * @param animate Whether to animate the showing of the update badge.
+     */
     public void showAppMenuUpdateBadge(boolean animate) {
         if (mUpdateBadgeView == null) return;
         mShowMenuBadge = true;
@@ -116,6 +119,10 @@
         setAppMenuUpdateBadgeToVisible(animate);
     }
 
+    /**
+     * Remove the update badge on the app menu button.
+     * @param animate Whether to animate the hiding of the update badge.
+     */
     public void removeAppMenuUpdateBadge(boolean animate) {
         if (mUpdateBadgeView == null) return;
         boolean wasShowingMenuBadge = mShowMenuBadge;
@@ -264,7 +271,6 @@
             mThemeColorProvider.removeTintObserver(this);
             mThemeColorProvider = null;
         }
-        UiUtils.removeViewFromParent(this);
     }
 
     /**
diff --git a/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/CustomTabToolbar.java b/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/CustomTabToolbar.java
index dab1986..386bc02 100644
--- a/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/CustomTabToolbar.java
+++ b/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/CustomTabToolbar.java
@@ -134,6 +134,7 @@
     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 = true;
@@ -181,6 +182,7 @@
         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);
     }
 
@@ -780,6 +782,9 @@
     public void setAutocompleteProfile(Profile profile) {}
 
     @Override
+    void showAppMenuUpdateBadge(boolean animate) {}
+
+    @Override
     boolean isShowingAppMenuUpdateBadge() {
         return false;
     }
@@ -794,8 +799,14 @@
     }
 
     @Override
+    ImageButton getMenuButton() {
+        return mMenuButton;
+    }
+
+    @Override
     void disableMenuButton() {
         super.disableMenuButton();
+        mMenuButton = null;
         // In addition to removing the menu button, we also need to remove the margin on the custom
         // action button.
         if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.JELLY_BEAN_MR1) {
diff --git a/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/ToolbarLayout.java b/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/ToolbarLayout.java
index 469013b..5c87dc8 100644
--- a/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/ToolbarLayout.java
+++ b/chrome/android/java/src/org/chromium/chrome/browser/toolbar/top/ToolbarLayout.java
@@ -111,6 +111,12 @@
         mToolbarTabController = tabController;
         if (mMenuButtonWrapper != null) {
             mMenuButtonWrapper.setAppMenuButtonHelper(appMenuButtonHelper);
+        } else {
+            final ImageButton menuButton = getMenuButton();
+            if (menuButton != null) {
+                menuButton.setOnTouchListener(appMenuButtonHelper);
+                menuButton.setAccessibilityDelegate(appMenuButtonHelper);
+            }
         }
     }
 
@@ -139,6 +145,8 @@
      * instance vars.
      */
     void disableMenuButton() {
+        UiUtils.removeViewFromParent(getMenuButtonWrapper());
+
         if (mMenuButtonWrapper != null) {
             mMenuButtonWrapper.destroy();
             mMenuButtonWrapper = null;