How to fix WebContentsImplBrowserTest.NoResetOnBeforeUnloadCanceledOnCommit.

content/browser/web_contents/web_contents_impl_browsertest.cc

Index: content/browser/web_contents/web_contents_impl_browsertest.cc
diff --git a/content/browser/web_contents/web_contents_impl_browsertest.cc b/content/browser/web_contents/web_contents_impl_browsertest.cc
index 924873625295deb322cde30e702c5bd3ffef35a5..2eba0a83f1155cb2a0305882b58f1d69ebb5dc8d 100644
--- a/content/browser/web_contents/web_contents_impl_browsertest.cc
+++ b/content/browser/web_contents/web_contents_impl_browsertest.cc
@@ -41,6 +41,8 @@
 
 namespace content {
 
+namespace {
+
 void ResizeWebContentsView(Shell* shell, const gfx::Size& size,
                            bool set_start_page) {
   // Shell::SizeTo is not implemented on Aura; WebContentsView::SizeContents
@@ -198,6 +200,74 @@ class LoadingStateChangedDelegate : public WebContentsDelegate {
   int loadingStateToDifferentDocumentCount_;
 };
 
+class TestJavaScriptDialogManager : public JavaScriptDialogManager,
+                                    public WebContentsDelegate {
+ public:
+  TestJavaScriptDialogManager() : message_loop_runner_(new MessageLoopRunner) {}
+  ~TestJavaScriptDialogManager() override {}
+
+  void Wait() {
+    message_loop_runner_->Run();
+    message_loop_runner_ = new MessageLoopRunner;
+  }
+
+  std::string last_message() { return last_message_; }
+  bool did_get_dialog() { return did_get_dialog_; }
+  bool did_get_beforeunload() { return did_get_beforeunload_; }
+
+  // WebContentsDelegate
+
+  JavaScriptDialogManager* GetJavaScriptDialogManager(
+      WebContents* source) override {
+    return this;
+  }
+
+  // JavaScriptDialogManager
+
+  void RunJavaScriptDialog(WebContents* web_contents,
+                           const GURL& origin_url,
+                           JavaScriptDialogType dialog_type,
+                           const base::string16& message_text,
+                           const base::string16& default_prompt_text,
+                           const DialogClosedCallback& callback,
+                           bool* did_suppress_message) override {
+    did_get_dialog_ = true;
+    last_message_ = base::UTF16ToUTF8(message_text);
+    *did_suppress_message = true;
+
+    message_loop_runner_->Quit();
+  };
+
+  void RunBeforeUnloadDialog(WebContents* web_contents,
+                             bool is_reload,
+                             const DialogClosedCallback& callback) override {
+    did_get_beforeunload_ = true;
+
+    message_loop_runner_->Quit();
+  }
+
+  bool HandleJavaScriptDialog(WebContents* web_contents,
+                              bool accept,
+                              const base::string16* prompt_override) override {
+    return true;
+  }
+
+  void CancelDialogs(WebContents* web_contents, bool reset_state) override {}
+
+ private:
+  std::string last_message_;
+
+  bool did_get_dialog_ = false;
+  bool did_get_beforeunload_ = false;
+
+  // The MessageLoopRunner used to spin the message loop.
+  scoped_refptr<MessageLoopRunner> message_loop_runner_;
+
+  DISALLOW_COPY_AND_ASSIGN(TestJavaScriptDialogManager);
+};
+
+}  // namespace
+
 // Test that DidStopLoading includes the correct URL in the details.
 IN_PROC_BROWSER_TEST_F(WebContentsImplBrowserTest, DidStopLoadingDetails) {
   ASSERT_TRUE(embedded_test_server()->Start());
@@ -980,13 +1050,15 @@ IN_PROC_BROWSER_TEST_F(WebContentsImplBrowserTest, NewNamedWindow) {
   }
 }
 
-// TODO(clamy): Make the test work on Windows and on Mac. On Mac and Windows,
-// there seem to be an issue with the ShellJavascriptDialogManager.
-// Flaky on all platforms: https://crbug.com/655628
 // Test that if a BeforeUnload dialog is destroyed due to the commit of a
 // cross-site navigation, it will not reset the loading state.
 IN_PROC_BROWSER_TEST_F(WebContentsImplBrowserTest,
-                       DISABLED_NoResetOnBeforeUnloadCanceledOnCommit) {
+                       NoResetOnBeforeUnloadCanceledOnCommit) {
+  WebContentsImpl* contents =
+      static_cast<WebContentsImpl*>(shell()->web_contents());
+  TestJavaScriptDialogManager dialog_manager;
+  contents->SetDelegate(&dialog_manager);
+
   ASSERT_TRUE(embedded_test_server()->Start());
   const GURL kStartURL(
       embedded_test_server()->GetURL("/hang_before_unload.html"));
@@ -995,10 +1067,12 @@ IN_PROC_BROWSER_TEST_F(WebContentsImplBrowserTest,
 
   // Navigate to a first web page with a BeforeUnload event listener.
   EXPECT_TRUE(NavigateToURL(shell(), kStartURL));
+  // Disable the hang monitor, otherwise there will be a race between the
+  // beforeunload dialog and the beforeunload hang timer.
+  contents->GetMainFrame()->DisableBeforeUnloadHangMonitorForTesting();
 
   // Start a cross-site navigation that will not commit for the moment.
-  TestNavigationManager cross_site_delayer(shell()->web_contents(),
-                                           kCrossSiteURL);
+  TestNavigationManager cross_site_delayer(contents, kCrossSiteURL);
   shell()->LoadURL(kCrossSiteURL);
   EXPECT_TRUE(cross_site_delayer.WaitForRequestStart());
 
@@ -1009,14 +1083,17 @@ IN_PROC_BROWSER_TEST_F(WebContentsImplBrowserTest,
   // schedule it for afterwards. Since the BeforeUnload event is synchronous,
   // clicking on the link right away would cause the ExecuteScript to never
   // return.
-  SetShouldProceedOnBeforeUnload(shell(), false);
   EXPECT_TRUE(ExecuteScript(shell(), "clickLinkSoon()"));
-  WaitForAppModalDialog(shell());
+  dialog_manager.Wait();
+  EXPECT_TRUE(dialog_manager.did_get_beforeunload());
 
   // Have the cross-site navigation commit. The main RenderFrameHost should
   // still be loading after that.
   cross_site_delayer.WaitForNavigationFinished();
   EXPECT_TRUE(shell()->web_contents()->IsLoading());
+
+  contents->SetDelegate(nullptr);
+  contents->SetJavaScriptDialogManagerForTesting(nullptr);
 }
 
 namespace {
@@ -1042,67 +1119,6 @@ IN_PROC_BROWSER_TEST_F(WebContentsImplBrowserTest, OnBeforeUnload) {
   NavigateToDataURLAndExpectBeforeUnload(shell(), BEFORE_UNLOAD_HTML, true);
 }
 
-namespace {
-
-class TestJavaScriptDialogManager : public JavaScriptDialogManager,
-                                    public WebContentsDelegate {
- public:
-  TestJavaScriptDialogManager() : message_loop_runner_(new MessageLoopRunner) {}
-  ~TestJavaScriptDialogManager() override {}
-
-  void Wait() {
-    message_loop_runner_->Run();
-    message_loop_runner_ = new MessageLoopRunner;
-  }
-
-  std::string last_message() { return last_message_; }
-
-  // WebContentsDelegate
-
-  JavaScriptDialogManager* GetJavaScriptDialogManager(
-      WebContents* source) override {
-    return this;
-  }
-
-  // JavaScriptDialogManager
-
-  void RunJavaScriptDialog(WebContents* web_contents,
-                           const GURL& origin_url,
-                           JavaScriptDialogType dialog_type,
-                           const base::string16& message_text,
-                           const base::string16& default_prompt_text,
-                           const DialogClosedCallback& callback,
-                           bool* did_suppress_message) override {
-    last_message_ = base::UTF16ToUTF8(message_text);
-    *did_suppress_message = true;
-
-    message_loop_runner_->Quit();
-  };
-
-  void RunBeforeUnloadDialog(WebContents* web_contents,
-                             bool is_reload,
-                             const DialogClosedCallback& callback) override {}
-
-  bool HandleJavaScriptDialog(WebContents* web_contents,
-                              bool accept,
-                              const base::string16* prompt_override) override {
-    return true;
-  }
-
-  void CancelDialogs(WebContents* web_contents,
-                     bool reset_state) override {}
-
- private:
-  std::string last_message_;
-
-  // The MessageLoopRunner used to spin the message loop.
-  scoped_refptr<MessageLoopRunner> message_loop_runner_;
-
-  DISALLOW_COPY_AND_ASSIGN(TestJavaScriptDialogManager);
-};
-
-}  // namespace
-
 IN_PROC_BROWSER_TEST_F(WebContentsImplBrowserTest,
                        JavaScriptDialogsInMainAndSubframes) {
   WebContentsImpl* wc = static_cast<WebContentsImpl*>(shell()->web_contents());