Properly attach InfoBarContainer when it is swapped to a new WebContents

chrome/android/javatests/src/org/chromium/chrome/browser/infobar/InfoBarTest.java

Index: chrome/android/javatests/src/org/chromium/chrome/browser/infobar/InfoBarTest.java
diff --git a/chrome/android/javatests/src/org/chromium/chrome/browser/infobar/InfoBarTest.java b/chrome/android/javatests/src/org/chromium/chrome/browser/infobar/InfoBarTest.java
index 25b784fd218d466c53f67352dfbc05eef6707f3d..162222646b6ef2ed7eee7c258ad9d0c6e2ab0dd4 100644
--- a/chrome/android/javatests/src/org/chromium/chrome/browser/infobar/InfoBarTest.java
+++ b/chrome/android/javatests/src/org/chromium/chrome/browser/infobar/InfoBarTest.java
@@ -9,8 +9,10 @@ import static org.chromium.base.test.util.ScalableTimeout.scaleTimeout;
 import android.test.suitebuilder.annotation.MediumTest;
 import android.test.suitebuilder.annotation.Smoke;
 
+import org.chromium.base.ThreadUtils;
 import org.chromium.base.test.util.Feature;
 import org.chromium.base.test.util.UrlUtils;
+import org.chromium.chrome.browser.ContentViewUtil;
 import org.chromium.chrome.browser.location.LocationSettingsTestUtil;
 import org.chromium.chrome.shell.ChromeShellTestBase;
 import org.chromium.chrome.test.util.InfoBarTestAnimationListener;
@@ -18,9 +20,11 @@ import org.chromium.chrome.test.util.InfoBarUtil;
 import org.chromium.chrome.test.util.TestHttpServerClient;
 import org.chromium.content.browser.test.util.Criteria;
 import org.chromium.content.browser.test.util.CriteriaHelper;
+import org.chromium.content_public.browser.WebContents;
 
 import java.util.List;
 
+/** Tests for the InfoBars. */
 public class InfoBarTest extends ChromeShellTestBase {
     private static final long MAX_TIMEOUT = scaleTimeout(2000);
     private static final int CHECK_INTERVAL = 500;
@@ -41,6 +45,14 @@ public class InfoBarTest extends ChromeShellTestBase {
         super.setUp();
 
         // Register for animation notifications
+        assertTrue(CriteriaHelper.pollForCriteria(new Criteria() {
+            @Override
+            public boolean isSatisfied() {
+                if (getActivity().getActiveTab() == null) return false;
+                if (getActivity().getActiveTab().getInfoBarContainer() == null) return false;
+                return true;
+            }
+        }));
         InfoBarContainer container = getActivity().getActiveTab().getInfoBarContainer();
         mListener =  new InfoBarTestAnimationListener();
         container.setAnimationListener(mListener);
@@ -128,4 +140,41 @@ public class InfoBarTest extends ChromeShellTestBase {
                 MAX_TIMEOUT, CHECK_INTERVAL);
         assertTrue("InfoBar not removed.", mListener.removeInfoBarAnimationFinished());
     }
+
+    /**
+     * Verify InfoBarContainers swap the WebContents they are monitoring properly.
+     */
+    @MediumTest
+    @Feature({"Browser", "Main"})
+    public void testInfoBarContainerSwapsWebContents() throws InterruptedException {
+        // Add an infobar.
+        LocationSettingsTestUtil.setSystemLocationSettingEnabled(true);
+        loadUrlWithSanitization(TestHttpServerClient.getUrl(GEOLOCATION_PAGE));
+        assertTrue("InfoBar not added", mListener.addInfoBarAnimationFinished());
+        List<InfoBar> infoBars = getActivity().getActiveTab().getInfoBarContainer().getInfoBars();
+        assertEquals("Wrong infobar count", 1, infoBars.size());
+
+        // Swap out the WebContents and send the user somewhere so that the InfoBar gets removed.
+        InfoBarTestAnimationListener removeListener = new InfoBarTestAnimationListener();
+        getActivity().getActiveTab().getInfoBarContainer().setAnimationListener(removeListener);
+        ThreadUtils.runOnUiThread(new Runnable() {
+            @Override
+            public void run() {
+                WebContents newContents = ContentViewUtil.createWebContents(false, false);
+                getActivity().getActiveTab().swapWebContents(newContents, false, false);
+            }
+        });
+        loadUrlWithSanitization(HELLO_WORLD_URL);
+        assertTrue("InfoBar not removed.", removeListener.removeInfoBarAnimationFinished());
+        infoBars = getActivity().getActiveTab().getInfoBarContainer().getInfoBars();
+        assertEquals("Wrong infobar count", 0, infoBars.size());
+
+        // Revisiting the original page should make the InfoBar reappear.
+        InfoBarTestAnimationListener addListener = new InfoBarTestAnimationListener();
+        getActivity().getActiveTab().getInfoBarContainer().setAnimationListener(addListener);
+        loadUrlWithSanitization(TestHttpServerClient.getUrl(GEOLOCATION_PAGE));
+        assertTrue("InfoBar not added", addListener.addInfoBarAnimationFinished());
+        infoBars = getActivity().getActiveTab().getInfoBarContainer().getInfoBars();
+        assertEquals("Wrong infobar count", 1, infoBars.size());
+    }
 }