Allow asynchronous deferral in NavigationSimulator

content/public/test/navigation_simulator.cc

Index: content/public/test/navigation_simulator.cc
diff --git a/content/public/test/navigation_simulator.cc b/content/public/test/navigation_simulator.cc
index 9cc1399bcf7da6050a8b26b6c019fdb51174dec9..bf07dfd398396d0ad6e44c0a9a1102b96846685a 100644
--- a/content/public/test/navigation_simulator.cc
+++ b/content/public/test/navigation_simulator.cc
@@ -4,6 +4,9 @@
 
 #include "content/public/test/navigation_simulator.h"
 
+#include "base/bind.h"
+#include "base/memory/ptr_util.h"
+#include "base/run_loop.h"
 #include "content/browser/frame_host/navigation_handle_impl.h"
 #include "content/browser/frame_host/navigation_request.h"
 #include "content/common/frame_messages.h"
@@ -61,6 +64,7 @@ void NavigationSimulator::NavigateAndCommitFromDocument(
   NavigationSimulator simulator(
       original_url, static_cast<TestRenderFrameHost*>(render_frame_host));
   simulator.Commit();
+  simulator.WaitForThrottleChecksComplete();
 }
 
 // static
@@ -101,7 +105,7 @@ NavigationSimulator::NavigationSimulator(const GURL& original_url,
 
 NavigationSimulator::~NavigationSimulator() {}
 
-void NavigationSimulator::Start() {
+NavigationThrottle::ThrottleCheckResult NavigationSimulator::Start() {
   CHECK(state_ == INITIALIZATION)
       << "NavigationSimulator::Start should only be called once.";
   state_ = STARTED;
@@ -123,7 +127,9 @@ void NavigationSimulator::Start() {
         render_frame_host_->GetRoutingID(), common_params, begin_params));
     NavigationRequest* request =
         render_frame_host_->frame_tree_node()->navigation_request();
-    DCHECK(request);
+    // The request failed synchronously.
+    if (!request)
+      return last_throttle_check_result_.value();
     DCHECK_EQ(handle_, request->navigation_handle());
   } else {
     render_frame_host_->OnMessageReceived(
@@ -137,21 +143,26 @@ void NavigationSimulator::Start() {
         true /* user_gesture */, transition_, false /* is_external_protocol */,
         REQUEST_CONTEXT_TYPE_LOCATION,
         blink::WebMixedContentContextType::NotMixedContent,
-        base::Callback<void(NavigationThrottle::ThrottleCheckResult)>());
+        ThrottleChecksFinishedCallback());
   }
 
   CHECK(handle_);
+  return RunOrStoreThrottleCompleteCallback(base::Bind(
+      &NavigationSimulator::FinishStart, weak_factory_.GetWeakPtr()));
+}
 
-  // Make sure all NavigationThrottles have run.
-  // TODO(clamy): provide a non auto-advance mode if needed.
-  while (handle_->state_for_testing() == NavigationHandleImpl::DEFERRING_START)
-    handle_->Resume();
-
+void NavigationSimulator::FinishStart(
+    NavigationThrottle::ThrottleCheckResult result) {
   CHECK_EQ(1, num_did_start_navigation_called_);
-  CHECK_EQ(1, num_will_start_request_called_);
+  if (result == NavigationThrottle::PROCEED) {
+    CHECK_EQ(1, num_will_start_request_called_);
+  } else {
+    state_ = FAILED;
+  }
 }
 
-void NavigationSimulator::Redirect(const GURL& new_url) {
+NavigationThrottle::ThrottleCheckResult NavigationSimulator::Redirect(
+    const GURL& new_url) {
   CHECK(state_ <= STARTED) << "NavigationSimulator::Redirect should be "
                               "called before Fail or Commit";
   CHECK_EQ(0, num_did_finish_navigation_called_)
@@ -159,7 +170,7 @@ void NavigationSimulator::Redirect(const GURL& new_url) {
          "navigation has finished";
 
   if (state_ == INITIALIZATION)
-    Start();
+    StartAndExpectToProceed();
 
   navigation_url_ = new_url;
 
@@ -168,6 +179,7 @@ void NavigationSimulator::Redirect(const GURL& new_url) {
   int previous_did_redirect_navigation_called =
       num_did_redirect_navigation_called_;
 
+  PrepareCompleteCallbackOnHandle();
   if (IsBrowserSideNavigationEnabled()) {
     NavigationRequest* request =
         render_frame_host_->frame_tree_node()->navigation_request();
@@ -187,27 +199,33 @@ void NavigationSimulator::Redirect(const GURL& new_url) {
     url_loader->CallOnRequestRedirected(
         redirect_info, scoped_refptr<ResourceResponse>(new ResourceResponse));
   } else {
-    handle_->WillRedirectRequest(
-        new_url, "GET", referrer_.url, false /* is_external_protocol */,
-        scoped_refptr<net::HttpResponseHeaders>(),
-        net::HttpResponseInfo::ConnectionInfo(),
-        base::Callback<void(NavigationThrottle::ThrottleCheckResult)>());
+    handle_->WillRedirectRequest(new_url, "GET", referrer_.url,
+                                 false /* is_external_protocol */,
+                                 scoped_refptr<net::HttpResponseHeaders>(),
+                                 net::HttpResponseInfo::ConnectionInfo(),
+                                 ThrottleChecksFinishedCallback());
   }
+  return RunOrStoreThrottleCompleteCallback(base::Bind(
+      &NavigationSimulator::FinishRedirect, weak_factory_.GetWeakPtr(),
+      previous_num_will_redirect_request_called,
+      previous_did_redirect_navigation_called));
+}
 
-  // Make sure all NavigationThrottles have run.
-  // TODO(clamy): provide a non auto-advance mode if needed.
-  while (handle_->state_for_testing() ==
-         NavigationHandleImpl::DEFERRING_REDIRECT) {
-    handle_->Resume();
+void NavigationSimulator::FinishRedirect(
+    int previous_num_will_redirect_request_called,
+    int previous_did_redirect_navigation_called,
+    NavigationThrottle::ThrottleCheckResult result) {
+  if (result == NavigationThrottle::PROCEED) {
+    CHECK_EQ(previous_num_will_redirect_request_called + 1,
+             num_will_redirect_request_called_);
+    CHECK_EQ(previous_did_redirect_navigation_called + 1,
+             num_did_redirect_navigation_called_);
+  } else {
+    state_ = FAILED;
   }
-
-  CHECK_EQ(previous_num_will_redirect_request_called + 1,
-           num_will_redirect_request_called_);
-  CHECK_EQ(previous_did_redirect_navigation_called + 1,
-           num_did_redirect_navigation_called_);
 }
 
-void NavigationSimulator::Commit() {
+NavigationThrottle::ThrottleCheckResult NavigationSimulator::Commit() {
   CHECK_LE(state_, STARTED) << "NavigationSimulator::Commit can only be "
                                "called once, and cannot be called after "
                                "NavigationSimulator::Fail";
@@ -216,31 +234,35 @@ void NavigationSimulator::Commit() {
          "navigation has finished";
 
   if (state_ == INITIALIZATION)
-    Start();
+    StartAndExpectToProceed();
 
+  PrepareCompleteCallbackOnHandle();
   if (IsBrowserSideNavigationEnabled() &&
       render_frame_host_->frame_tree_node()->navigation_request()) {
     render_frame_host_->PrepareForCommit();
   }
 
   // Call NavigationHandle::WillProcessResponse if needed.
-  if (handle_->state_for_testing() <
-      NavigationHandleImpl::WILL_PROCESS_RESPONSE) {
+  // Note that the handle's state can be CANCELING if a throttle cancelled it
+  // synchronously in PrepareForCommit.
+  if (handle_->state_for_testing() < NavigationHandleImpl::CANCELING) {
     handle_->WillProcessResponse(
         render_frame_host_, scoped_refptr<net::HttpResponseHeaders>(),
         net::HttpResponseInfo::ConnectionInfo(), SSLStatus(), GlobalRequestID(),
         false /* should_replace_current_entry */, false /* is_download */,
         false /* is_stream */, base::Closure(),
-        base::Callback<void(NavigationThrottle::ThrottleCheckResult)>());
+        ThrottleChecksFinishedCallback());
   }
+  return RunOrStoreThrottleCompleteCallback(base::Bind(
+      &NavigationSimulator::FinishCommit, weak_factory_.GetWeakPtr()));
+}
 
-  // Make sure all NavigationThrottles have run.
-  // TODO(clamy): provide a non auto-advance mode if needed.
-  while (handle_->state_for_testing() ==
-         NavigationHandleImpl::DEFERRING_RESPONSE) {
-    handle_->Resume();
+void NavigationSimulator::FinishCommit(
+    NavigationThrottle::ThrottleCheckResult result) {
+  if (result != NavigationThrottle::PROCEED) {
+    state_ = FAILED;
+    return;
   }
-
   CHECK_EQ(1, num_will_process_response_called_);
   CHECK_EQ(1, num_ready_to_commit_called_);
 
@@ -310,7 +332,7 @@ void NavigationSimulator::Fail(int error_code) {
          "navigation has finished";
 
   if (state_ == INITIALIZATION)
-    Start();
+    StartAndExpectToProceed();
 
   state_ = FAILED;
 
@@ -446,6 +468,17 @@ void NavigationSimulator::SetReferrer(const Referrer& referrer) {
   referrer_ = referrer;
 }
 
+NavigationThrottle::ThrottleCheckResult
+NavigationSimulator::WaitForThrottleChecksComplete() {
+  if (!last_throttle_check_result_) {
+    base::RunLoop run_loop;
+    throttle_checks_wait_closure_ = run_loop.QuitClosure();
+    run_loop.Run();
+    throttle_checks_wait_closure_.Reset();
+  }
+  return last_throttle_check_result_.value();
+}
+
 void NavigationSimulator::DidStartNavigation(
     NavigationHandle* navigation_handle) {
   // Check if this navigation is the one we're simulating.
@@ -475,6 +508,8 @@ void NavigationSimulator::DidStartNavigation(
                      weak_factory_.GetWeakPtr()),
           base::Bind(&NavigationSimulator::OnWillProcessResponse,
                      weak_factory_.GetWeakPtr())));
+
+  PrepareCompleteCallbackOnHandle();
 }
 
 void NavigationSimulator::DidRedirectNavigation(
@@ -507,4 +542,57 @@ void NavigationSimulator::OnWillProcessResponse() {
   num_will_process_response_called_++;
 }
 
+// These methods are a little subtle. Overall, the design constraints are:
+// 1. Synchronous throttle checks should not cause subsequent
+// WaitForThrottleChecksComplete to block.
+//
+// 2. |complete_callback_| is called after all synchronous code is executed in
+// response to the throttle checks. This point is important because in some
+// cases (e.g. cancellation), the OnThrottleChecksComplete callback  that the
+// NavigationRequest/NavigationHandle holds on to will be called before things
+// like navigation tear down and DidFinishNavigation. This code ensures that
+// |complete_callback_| is called after that occurs, either synchronously in
+// RunOrStoreThrottleCompleteCallback, or asynchronously in
+// OnThrottleChecksComplete.
+NavigationThrottle::ThrottleCheckResult
+NavigationSimulator::RunOrStoreThrottleCompleteCallback(
+    ThrottleChecksFinishedCallback finish_callback) {
+  // If completed the checks synchronously, call Finish* here. Else, save the
+  // callback to call asynchronously when the checks finish.
+  if (last_throttle_check_result_) {
+    finish_callback.Run(last_throttle_check_result_.value());
+    return last_throttle_check_result_.value();
+  }
+  complete_callback_ = std::move(finish_callback);
+  return NavigationThrottle::DEFER;
+}
+
+void NavigationSimulator::OnThrottleChecksComplete(
+    NavigationThrottle::ThrottleCheckResult result) {
+  // For synchronous throttle checks, set this here, and read from it downstream
+  // in RunOrStoreThrottleCompleteCallback.
+  DCHECK(!last_throttle_check_result_);
+  last_throttle_check_result_ = result;
+  if (!complete_callback_)
+    return;
+  BrowserThread::PostTaskAndReply(
+      BrowserThread::UI, FROM_HERE,
+      base::Bind(std::move(complete_callback_), result),
+      throttle_checks_wait_closure_);
+}
+
+void NavigationSimulator::PrepareCompleteCallbackOnHandle() {
+  last_throttle_check_result_.reset();
+  handle_->set_complete_callback_for_testing(
+      base::Bind(&NavigationSimulator::OnThrottleChecksComplete,
+                 weak_factory_.GetWeakPtr()));
+}
+
+void NavigationSimulator::StartAndExpectToProceed() {
+  NavigationThrottle::ThrottleCheckResult result = Start();
+  if (result == NavigationThrottle::DEFER)
+    result = WaitForThrottleChecksComplete();
+  CHECK_EQ(NavigationThrottle::PROCEED, result);
+}
+
 }  // namespace content