[E2E] Use Supplier for LayoutManager to create EdgeToEdgeController

In CCT, the LayoutManager is created a little later than
EdgeToEdgeController due to the sequence of
#onFinishNativeInitialization. Use supplier to pass in the LayoutManager
to resolve the timing.

Bug: 369699256
Change-Id: I59c8fd1ecc275b727bd5a4025c62d93d06d25969
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5894007
Commit-Queue: Wenyu Fu <wenyufu@chromium.org>
Reviewed-by: Charles Hager <clhager@google.com>
Cr-Commit-Position: refs/heads/main@{#1361228}
diff --git a/chrome/android/java/src/org/chromium/chrome/browser/ui/RootUiCoordinator.java b/chrome/android/java/src/org/chromium/chrome/browser/ui/RootUiCoordinator.java
index 1b865f7d..8d075ac 100644
--- a/chrome/android/java/src/org/chromium/chrome/browser/ui/RootUiCoordinator.java
+++ b/chrome/android/java/src/org/chromium/chrome/browser/ui/RootUiCoordinator.java
@@ -4,7 +4,6 @@
 
 package org.chromium.chrome.browser.ui;
 
-import android.app.Activity;
 import android.app.Fragment;
 import android.content.ComponentName;
 import android.content.Intent;
@@ -894,12 +893,7 @@
         initScrollCapture();
 
         // TODO(crbug.com/350610430) Potentially create the E2EController earlier during startup
-        initializeEdgeToEdgeController(
-                mActivity,
-                mActivityTabProvider,
-                mEdgeToEdgeControllerSupplier,
-                mBrowserControlsManager,
-                mLayoutManager);
+        initializeEdgeToEdgeController();
         initBoardingPassDetector();
 
         if (EphemeralTabCoordinator.isSupported()) {
@@ -1954,13 +1948,9 @@
     }
 
     /** Setup drawing using Android Edge-to-Edge. */
-    protected void initializeEdgeToEdgeController(
-            Activity activity,
-            ActivityTabProvider activityTabProvider,
-            ObservableSupplierImpl<EdgeToEdgeController> supplier,
-            BrowserControlsManager browserControlsManager,
-            LayoutManager layoutManager) {
-        boolean eligible = EdgeToEdgeUtils.recordEligibility(activity);
+    @CallSuper
+    protected void initializeEdgeToEdgeController() {
+        boolean eligible = EdgeToEdgeUtils.recordEligibility(mActivity);
 
         UmaSessionStats.registerSyntheticFieldTrial(
                 "EdgeToEdgeChinEligibility", eligible ? "Eligible" : "Not Eligible");
@@ -1968,13 +1958,13 @@
         if (supportsEdgeToEdge()) {
             mEdgeToEdgeController =
                     EdgeToEdgeControllerFactory.create(
-                            activity,
+                            mActivity,
                             mWindowAndroid,
-                            activityTabProvider,
-                            browserControlsManager,
-                            layoutManager,
+                            mActivityTabProvider,
+                            mBrowserControlsManager,
+                            mLayoutManagerSupplier,
                             mFullscreenManager);
-            supplier.set(mEdgeToEdgeController);
+            mEdgeToEdgeControllerSupplier.set(mEdgeToEdgeController);
 
             if (EdgeToEdgeUtils.isEdgeToEdgeBottomChinEnabled()) {
                 mEdgeToEdgeBottomChin = createEdgeToEdgeBottomChin();
diff --git a/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerFactory.java b/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerFactory.java
index 697d1571..31a62794 100644
--- a/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerFactory.java
+++ b/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerFactory.java
@@ -46,7 +46,8 @@
      *     changes.
      * @param browserControlsStateProvider Provides the state of the BrowserControls so we can tell
      *     if the Toolbar is changing.
-     * @param layoutManager The {@link LayoutManager} for checking the active layout type.
+     * @param layoutManagerSupplier The supplier of {@link LayoutManager} for checking the active
+     *     layout type.
      * @return An EdgeToEdgeController to control drawing under System Bars, or {@code null} if this
      *     version of Android does not support the APIs needed.
      */
@@ -55,7 +56,7 @@
             WindowAndroid windowAndroid,
             @NonNull ObservableSupplier<Tab> tabObservableSupplier,
             BrowserControlsStateProvider browserControlsStateProvider,
-            LayoutManager layoutManager,
+            ObservableSupplier<LayoutManager> layoutManagerSupplier,
             FullscreenManager fullscreenManager) {
         if (Build.VERSION.SDK_INT < VERSION_CODES.R) return null;
         assert isSupportedConfiguration(activity);
@@ -65,7 +66,7 @@
                 tabObservableSupplier,
                 null,
                 browserControlsStateProvider,
-                layoutManager,
+                layoutManagerSupplier,
                 fullscreenManager);
     }
 
diff --git a/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerImpl.java b/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerImpl.java
index 84c9970..b16fab06 100644
--- a/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerImpl.java
+++ b/chrome/browser/ui/android/edge_to_edge/internal/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerImpl.java
@@ -17,14 +17,17 @@
 import androidx.core.graphics.Insets;
 import androidx.core.view.WindowInsetsCompat;
 
+import org.chromium.base.Callback;
 import org.chromium.base.Log;
 import org.chromium.base.ObserverList;
+import org.chromium.base.ValueChangedCallback;
 import org.chromium.base.supplier.ObservableSupplier;
 import org.chromium.chrome.browser.browser_controls.BrowserControlsStateProvider;
 import org.chromium.chrome.browser.fullscreen.FullscreenManager;
 import org.chromium.chrome.browser.fullscreen.FullscreenOptions;
 import org.chromium.chrome.browser.layouts.LayoutManager;
 import org.chromium.chrome.browser.layouts.LayoutStateProvider;
+import org.chromium.chrome.browser.layouts.LayoutType;
 import org.chromium.chrome.browser.tab.EmptyTabObserver;
 import org.chromium.chrome.browser.tab.Tab;
 import org.chromium.chrome.browser.tab.TabObserver;
@@ -58,7 +61,9 @@
     private final ObserverList<ChangeObserver> mEdgeChangeObservers = new ObserverList<>();
     private final @NonNull TabObserver mTabObserver;
     private final BrowserControlsStateProvider mBrowserControlsStateProvider;
-    private final LayoutManager mLayoutManager;
+    private final ObservableSupplier<LayoutManager> mLayoutManagerSupplier;
+    private final Callback<LayoutManager> mOnLayoutManagerCallback =
+            new ValueChangedCallback<>(this::updateLayoutStateProvider);
     private final FullscreenManager mFullscreenManager;
 
     // Cached rects used for adding under fullscreen.
@@ -68,7 +73,8 @@
     /** Multiplier to convert from pixels to DPs. */
     private final float mPxToDp;
 
-    private @NonNull EdgeToEdgeOSWrapper mEdgeToEdgeOSWrapper;
+    private @NonNull EdgeToEdgeOSWrapper mEdgeToEdgeOsWrapper;
+    private @Nullable LayoutManager mLayoutManager;
 
     private Tab mCurrentTab;
     private WebContentsObserver mWebContentsObserver;
@@ -103,24 +109,25 @@
      * @param windowAndroid The current {@link WindowAndroid} to allow drawing under System Bars.
      * @param tabObservableSupplier A supplier for Tab changes so this implementation can adjust
      *     whether to draw under or not for each page.
-     * @param edgeToEdgeOSWrapper An optional wrapper for OS calls for testing etc.
+     * @param edgeToEdgeOsWrapper An optional wrapper for OS calls for testing etc.
      * @param browserControlsStateProvider Provides the state of the BrowserControls for Totally
      *     Edge to Edge.
-     * @param layoutManager The {@link LayoutManager} for checking the active layout type.
+     * @param layoutManagerSupplier The supplier to {@link LayoutManager} for checking the active
+     *     layout type.
      * @param fullscreenManager The {@link FullscreenManager} for checking the fullscreen state.
      */
     public EdgeToEdgeControllerImpl(
-            Activity activity,
-            WindowAndroid windowAndroid,
-            ObservableSupplier<Tab> tabObservableSupplier,
-            @Nullable EdgeToEdgeOSWrapper edgeToEdgeOSWrapper,
-            BrowserControlsStateProvider browserControlsStateProvider,
-            LayoutManager layoutManager,
-            FullscreenManager fullscreenManager) {
+            @NonNull Activity activity,
+            @NonNull WindowAndroid windowAndroid,
+            @NonNull ObservableSupplier<Tab> tabObservableSupplier,
+            @Nullable EdgeToEdgeOSWrapper edgeToEdgeOsWrapper,
+            @NonNull BrowserControlsStateProvider browserControlsStateProvider,
+            @NonNull ObservableSupplier<LayoutManager> layoutManagerSupplier,
+            @NonNull FullscreenManager fullscreenManager) {
         mActivity = activity;
         mWindowAndroid = windowAndroid;
-        mEdgeToEdgeOSWrapper =
-                edgeToEdgeOSWrapper == null ? new EdgeToEdgeOSWrapperImpl() : edgeToEdgeOSWrapper;
+        mEdgeToEdgeOsWrapper =
+                edgeToEdgeOsWrapper == null ? new EdgeToEdgeOSWrapperImpl() : edgeToEdgeOsWrapper;
         mPxToDp = 1.f / mActivity.getResources().getDisplayMetrics().density;
         mTabSupplierObserver =
                 new TabSupplierObserver(tabObservableSupplier) {
@@ -148,8 +155,14 @@
         mInsetObserver = mWindowAndroid.getInsetObserver();
         mBrowserControlsStateProvider = browserControlsStateProvider;
         mBrowserControlsStateProvider.addObserver(this);
-        mLayoutManager = layoutManager;
-        mLayoutManager.addObserver(this);
+
+        mLayoutManagerSupplier = layoutManagerSupplier;
+        mLayoutManagerSupplier.addObserver(mOnLayoutManagerCallback);
+        mLayoutManager = layoutManagerSupplier.get();
+        if (mLayoutManager != null) {
+            mLayoutManager.addObserver(this);
+        }
+
         mFullscreenManager = fullscreenManager;
         mFullscreenManager.addObserver(this);
 
@@ -160,7 +173,7 @@
                 : "The inset observer should have non-null insets by the time the"
                         + " EdgeToEdgeControllerImpl is initialized.";
         mSystemInsets = getSystemInsets(mInsetObserver.getLastRawWindowInsets());
-        mEdgeToEdgeOSWrapper.setDecorFitsSystemWindows(mActivity.getWindow(), false);
+        mEdgeToEdgeOsWrapper.setDecorFitsSystemWindows(mActivity.getWindow(), false);
         drawToEdge(
                 EdgeToEdgeUtils.isPageOptedIntoEdgeToEdge(mCurrentTab),
                 /* changedWindowState= */ true);
@@ -303,6 +316,20 @@
         Log.i(TAG, "E2E_Up Tab '%s'", tab.getTitle());
     }
 
+    private void updateLayoutStateProvider(
+            @Nullable LayoutManager newValue, @Nullable LayoutManager oldValue) {
+        if (oldValue != null) {
+            oldValue.removeObserver(this);
+        }
+        if (newValue != null) {
+            newValue.addObserver(this);
+        }
+        mLayoutManager = newValue;
+        drawToEdge(
+                EdgeToEdgeUtils.isPageOptedIntoEdgeToEdge(mCurrentTab),
+                /* changedWindowState= */ false);
+    }
+
     /**
      * Conditionally draws the given View ToEdge or ToNormal based on the {@code toEdge} param.
      *
@@ -312,11 +339,12 @@
      */
     @VisibleForTesting
     void drawToEdge(boolean pageOptedIntoEdgeToEdge, boolean changedWindowState) {
+        @LayoutType
+        int currentLayoutType =
+                mLayoutManager != null ? mLayoutManager.getActiveLayoutType() : LayoutType.NONE;
         boolean shouldDrawToEdge =
                 EdgeToEdgeUtils.shouldDrawToEdge(
-                        pageOptedIntoEdgeToEdge,
-                        mLayoutManager.getActiveLayoutType(),
-                        mSystemInsets.bottom);
+                        pageOptedIntoEdgeToEdge, currentLayoutType, mSystemInsets.bottom);
         boolean changedPageOptedIn = pageOptedIntoEdgeToEdge != mIsPageOptedIntoEdgeToEdge;
         boolean changedDrawToEdge = shouldDrawToEdge != mIsDrawingToEdge;
         mIsPageOptedIntoEdgeToEdge = pageOptedIntoEdgeToEdge;
@@ -450,7 +478,7 @@
                 Insets.of(mSystemInsets.left, topPadding, mSystemInsets.right, bottomPadding);
         if (!newPaddings.equals(mAppliedContentViewPadding)) {
             mAppliedContentViewPadding = newPaddings;
-            mEdgeToEdgeOSWrapper.setPadding(
+            mEdgeToEdgeOsWrapper.setPadding(
                     contentView,
                     newPaddings.left,
                     newPaddings.top,
@@ -484,8 +512,12 @@
         if (mBrowserControlsStateProvider != null) {
             mBrowserControlsStateProvider.removeObserver(this);
         }
+        if (mOnLayoutManagerCallback != null) {
+            mLayoutManagerSupplier.removeObserver(mOnLayoutManagerCallback);
+        }
         if (mLayoutManager != null) {
             mLayoutManager.removeObserver(this);
+            mLayoutManager = null;
         }
         if (mFullscreenManager != null) {
             mFullscreenManager.removeObserver(this);
@@ -493,7 +525,7 @@
     }
 
     public void setOsWrapperForTesting(EdgeToEdgeOSWrapper testOsWrapper) {
-        mEdgeToEdgeOSWrapper = testOsWrapper;
+        mEdgeToEdgeOsWrapper = testOsWrapper;
     }
 
     @VisibleForTesting
diff --git a/chrome/browser/ui/android/edge_to_edge/internal/junit/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerTest.java b/chrome/browser/ui/android/edge_to_edge/internal/junit/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerTest.java
index 76cbb8206..e75468d 100644
--- a/chrome/browser/ui/android/edge_to_edge/internal/junit/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerTest.java
+++ b/chrome/browser/ui/android/edge_to_edge/internal/junit/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeControllerTest.java
@@ -39,12 +39,14 @@
 import org.hamcrest.Matchers;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 import org.robolectric.Robolectric;
 import org.robolectric.annotation.Config;
 import org.robolectric.annotation.Implementation;
@@ -112,10 +114,14 @@
     private Activity mActivity;
     private EdgeToEdgeControllerImpl mEdgeToEdgeControllerImpl;
 
-    private ObservableSupplierImpl<Tab> mTabProvider;
+    private final ObservableSupplierImpl<Tab> mTabProvider = new ObservableSupplierImpl<>();
+    private final ObservableSupplierImpl<LayoutManager> mLayoutManagerSupplier =
+            new ObservableSupplierImpl<>();
 
     private UserDataHost mTabDataHost = new UserDataHost();
 
+    @Rule public MockitoRule mMockitoRule = MockitoJUnit.rule();
+
     @Mock private WindowAndroid mWindowAndroid;
     @Mock private InsetObserver mInsetObserver;
     @Mock private Tab mTab;
@@ -149,12 +155,11 @@
         ChromeFeatureList.sEdgeToEdgeWebOptIn.setForTesting(true);
         ChromeFeatureList.sDrawNativeEdgeToEdge.setForTesting(false);
 
-        MockitoAnnotations.openMocks(this);
         when(mWindowAndroid.getInsetObserver()).thenReturn(mInsetObserver);
         when(mInsetObserver.getLastRawWindowInsets()).thenReturn(SYSTEM_BARS_WINDOW_INSETS);
 
         mActivity = Robolectric.buildActivity(AppCompatActivity.class).setup().get();
-        mTabProvider = new ObservableSupplierImpl<>();
+        mLayoutManagerSupplier.set(mLayoutManager);
 
         doNothing().when(mTab).addObserver(any());
         when(mTab.getUserDataHost()).thenReturn(mTabDataHost);
@@ -179,7 +184,7 @@
                         mTabProvider,
                         mOsWrapper,
                         mBrowserControlsStateProvider,
-                        mLayoutManager,
+                        mLayoutManagerSupplier,
                         mFullscreenManager);
         assertNotNull(mEdgeToEdgeControllerImpl);
         verify(mOsWrapper, times(1)).setDecorFitsSystemWindows(any(), eq(false));
@@ -198,7 +203,6 @@
     public void tearDown() {
         mEdgeToEdgeControllerImpl.destroy();
         mEdgeToEdgeControllerImpl = null;
-        mTabProvider = null;
     }
 
     @Test
@@ -361,7 +365,7 @@
                                 mWindowAndroid,
                                 liveSupplier,
                                 mBrowserControlsStateProvider,
-                                mLayoutManager,
+                                mLayoutManagerSupplier,
                                 mFullscreenManager);
         assertNotNull(liveController);
         liveController.setIsOptedIntoEdgeToEdgeForTesting(true);
@@ -471,6 +475,21 @@
     }
 
     @Test
+    @Features.EnableFeatures(ChromeFeatureList.EDGE_TO_EDGE_BOTTOM_CHIN)
+    public void testLayoutManagerChanged() {
+        mEdgeToEdgeControllerImpl.setIsOptedIntoEdgeToEdgeForTesting(false);
+        mEdgeToEdgeControllerImpl.setIsDrawingToEdgeForTesting(true);
+        mEdgeToEdgeControllerImpl.setSystemInsetsForTesting(SYSTEM_INSETS);
+
+        doReturn(LayoutType.BROWSING).when(mLayoutManager).getActiveLayoutType();
+        mEdgeToEdgeControllerImpl.onStartedShowing(LayoutType.BROWSING);
+        assertToEdgeExpectations();
+
+        mLayoutManagerSupplier.set(null);
+        assertToNormalExpectations();
+    }
+
+    @Test
     public void switchFullscreenMode_NoStatusBarNoNavBar() {
         doReturn(true).when(mFullscreenManager).getPersistentFullscreenMode();
         mEdgeToEdgeControllerImpl.onEnterFullscreen(mTab, new FullscreenOptions(false, false));
diff --git a/chrome/browser/ui/android/edge_to_edge/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeUtils.java b/chrome/browser/ui/android/edge_to_edge/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeUtils.java
index 333d4d4..4bfd456 100644
--- a/chrome/browser/ui/android/edge_to_edge/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeUtils.java
+++ b/chrome/browser/ui/android/edge_to_edge/java/src/org/chromium/chrome/browser/ui/edge_to_edge/EdgeToEdgeUtils.java
@@ -229,6 +229,9 @@
         if (tab == null || tab.isNativePage()) {
             return ChromeFeatureList.sDrawNativeEdgeToEdge.isEnabled();
         }
+        if (sAlwaysDrawWebEdgeToEdgeForTesting) {
+            return true;
+        }
         if (!isEdgeToEdgeWebOptInEnabled()) {
             return false;
         }
diff --git a/chrome/browser/ui/android/toolbar/java/src/org/chromium/chrome/browser/toolbar/bottom/BottomControlsMediatorTest.java b/chrome/browser/ui/android/toolbar/java/src/org/chromium/chrome/browser/toolbar/bottom/BottomControlsMediatorTest.java
index 6744943..5d38d69 100644
--- a/chrome/browser/ui/android/toolbar/java/src/org/chromium/chrome/browser/toolbar/bottom/BottomControlsMediatorTest.java
+++ b/chrome/browser/ui/android/toolbar/java/src/org/chromium/chrome/browser/toolbar/bottom/BottomControlsMediatorTest.java
@@ -24,10 +24,12 @@
 import androidx.core.view.WindowInsetsCompat;
 
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.Mock;
-import org.mockito.MockitoAnnotations;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 import org.robolectric.Robolectric;
 import org.robolectric.annotation.Config;
 
@@ -69,6 +71,8 @@
                     .setInsets(WindowInsetsCompat.Type.statusBars(), STATUS_BAR_INSETS)
                     .build();
 
+    @Rule public MockitoRule mMockitoRule = MockitoJUnit.rule();
+
     @Mock BottomControlsStacker mBottomControlsStacker;
     @Mock BrowserControlsStateProvider mBrowserControlsStateProvider;
     @Mock BrowserStateBrowserControlsVisibilityDelegate mBrowserControlsVisibilityDelegate;
@@ -91,7 +95,6 @@
 
     @Before
     public void setUp() {
-        MockitoAnnotations.initMocks(this);
         doReturn(mKeyboardDelegate).when(mWindowAndroid).getKeyboardDelegate();
         doReturn(SYSTEM_BARS_WINDOW_INSETS).when(mInsetObserver).getLastRawWindowInsets();
         doReturn(mInsetObserver).when(mWindowAndroid).getInsetObserver();
@@ -168,7 +171,7 @@
                         mTabObservableSupplier,
                         null,
                         mBrowserControlsStateProvider,
-                        mLayoutManager,
+                        new ObservableSupplierImpl<>(mLayoutManager),
                         mFullscreenManager);
         BottomControlsMediator plainMediator =
                 new BottomControlsMediator(
@@ -199,7 +202,7 @@
                         mTabObservableSupplier,
                         null,
                         mBrowserControlsStateProvider,
-                        mLayoutManager,
+                        new ObservableSupplierImpl<>(mLayoutManager),
                         mFullscreenManager);
         new BottomControlsMediator(
                 mWindowAndroid,