Add a function to add extra headers from NavigationThrottle

content/browser/frame_host/navigation_handle_impl_unittest.cc

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/macros.h"
 #include "content/browser/frame_host/navigation_handle_impl.h"
 #include "content/public/browser/navigation_throttle.h"
 #include "content/test/test_render_frame_host.h"
 
 namespace content {
@@ skipping 27 matching lines @@
 
  private:
   // The result returned by the TestNavigationThrottle.
   NavigationThrottle::ThrottleCheckResult result_;
 
   // The number of times WillStartRequest and WillRedirectRequest were called.
   int will_start_calls_;
   int will_redirect_calls_;
 };
 
+// Test version of a NavigationThrottle that will add a header when called.
+class TestHeaderAddingNavigationThrottle : public NavigationThrottle {
+ public:
+  TestHeaderAddingNavigationThrottle(NavigationHandle* handle,
+                                     std::string key,
+                                     std::string value)
+      : NavigationThrottle(handle), key_(key), value_(value) {}
+
+  ~TestHeaderAddingNavigationThrottle() override {}
+
+  NavigationThrottle::ThrottleCheckResult WillStartRequest() override {
+    static_cast<NavigationHandleImpl*>(navigation_handle())
+        ->AddExtraHeader(key_, value_);
+    return PROCEED;
+  }
+
+  NavigationThrottle::ThrottleCheckResult WillRedirectRequest() override {
+    static_cast<NavigationHandleImpl*>(navigation_handle())
+        ->AddExtraHeader(key_, value_);
+    return PROCEED;
+  }
+
+ private:
+  const std::string key_;
+  const std::string value_;
+};
+
 class NavigationHandleImplTest : public RenderViewHostImplTestHarness {
  public:
   NavigationHandleImplTest()
       : was_callback_called_(false),
         callback_result_(NavigationThrottle::DEFER) {}
 
   void SetUp() override {
     RenderViewHostImplTestHarness::SetUp();
     test_handle_ = NavigationHandleImpl::Create(
         GURL(), main_test_rfh()->frame_tree_node());
@@ skipping 16 matching lines @@
 
   bool IsCanceling() {
     return test_handle_->state() == NavigationHandleImpl::CANCELING;
   }
 
   // Helper function to call WillStartRequest on |handle|. If this function
   // returns DEFER, |callback_result_| will be set to the actual result of
   // the throttle checks when they are finished.
   void SimulateWillStartRequest() {
     was_callback_called_ = false;
+    extra_headers_.clear();
     callback_result_ = NavigationThrottle::DEFER;
 
     // It's safe to use base::Unretained since the NavigationHandle is owned by
     // the NavigationHandleImplTest.
     test_handle_->WillStartRequest(
         false, Referrer(), false, ui::PAGE_TRANSITION_LINK, false,
         base::Bind(&NavigationHandleImplTest::UpdateThrottleCheckResult,
                    base::Unretained(this)));
   }
 
   // Helper function to call WillRedirectRequest on |handle|. If this function
   // returns DEFER, |callback_result_| will be set to the actual result of the
   // throttle checks when they are finished.
   // TODO(clamy): this should also simulate that WillStartRequest was called if
   // it has not been called before.
   void SimulateWillRedirectRequest() {
     was_callback_called_ = false;
+    extra_headers_.clear();
     callback_result_ = NavigationThrottle::DEFER;
 
     // It's safe to use base::Unretained since the NavigationHandle is owned by
     // the NavigationHandleImplTest.
     test_handle_->WillRedirectRequest(
         GURL(), false, GURL(), false,
         base::Bind(&NavigationHandleImplTest::UpdateThrottleCheckResult,
                    base::Unretained(this)));
   }
 
@@ skipping 12 matching lines @@
   // |result| on checks.
   TestNavigationThrottle* CreateTestNavigationThrottle(
       NavigationThrottle::ThrottleCheckResult result) {
     TestNavigationThrottle* test_throttle =
         new TestNavigationThrottle(test_handle(), result);
     test_handle()->RegisterThrottleForTesting(
         scoped_ptr<TestNavigationThrottle>(test_throttle));
     return test_throttle;
   }
 
+  // Creates and registers a TestHeaderAddingNavigationThrottle that will add a
+  // header with key |key| and value |value| on checks.
+  void AddHeaderAddingThrottle(std::string key, std::string value) {
+    test_handle()->RegisterThrottleForTesting(scoped_ptr<NavigationThrottle>(
+        new TestHeaderAddingNavigationThrottle(test_handle(), key, value)));
+  }
+
+  const NavigationHandleImpl::ExtraHeadersList& extra_headers() const {
+    return extra_headers_;
+  }
+
  private:
   // The callback provided to NavigationHandleImpl::WillStartRequest and
   // NavigationHandleImpl::WillRedirectRequest during the tests.
   void UpdateThrottleCheckResult(
-      NavigationThrottle::ThrottleCheckResult result) {
+      NavigationThrottle::ThrottleCheckResult result,
+      const NavigationHandleImpl::ExtraHeadersList& extra_headers) {
     callback_result_ = result;
     was_callback_called_ = true;
+    extra_headers_ = extra_headers;
   }
 
   scoped_ptr<NavigationHandleImpl> test_handle_;
   bool was_callback_called_;
   NavigationThrottle::ThrottleCheckResult callback_result_;
+  NavigationHandleImpl::ExtraHeadersList extra_headers_;
 };
 
 // Checks that a deferred navigation can be properly resumed.
 TEST_F(NavigationHandleImplTest, ResumeDeferred) {
   TestNavigationThrottle* test_throttle =
       CreateTestNavigationThrottle(NavigationThrottle::DEFER);
   EXPECT_FALSE(IsDeferringStart());
   EXPECT_FALSE(IsDeferringRedirect());
   EXPECT_EQ(0, test_throttle->will_start_calls());
   EXPECT_EQ(0, test_throttle->will_redirect_calls());
@@ skipping 325 matching lines @@
   EXPECT_FALSE(IsDeferringStart());
   EXPECT_FALSE(IsDeferringRedirect());
   EXPECT_TRUE(was_callback_called());
   EXPECT_EQ(NavigationThrottle::CANCEL_AND_IGNORE, callback_result());
   EXPECT_EQ(0, cancel_throttle->will_start_calls());
   EXPECT_EQ(1, cancel_throttle->will_redirect_calls());
   EXPECT_EQ(0, proceed_throttle->will_start_calls());
   EXPECT_EQ(0, proceed_throttle->will_redirect_calls());
 }
 
+// Checks that extra headers are properly returned to the caller of
+// NavigationHandleImpl::WillStartRequest and
+// NavigationHandleImpl::WillSendRequest. In particular, they should be
+// returned in the order they were registered by the throttles.
+TEST_F(NavigationHandleImplTest, AddExtraHeaders) {
+  const std::string kKey1 = "foo";
+  const std::string kValue1 = "foobar";
+  const std::string kKey2 = "bar";
+  const std::string kValue2 = "barfoo";
+  AddHeaderAddingThrottle(kKey1, kValue1);
+  AddHeaderAddingThrottle(kKey2, kValue2);
+  EXPECT_EQ(0u, extra_headers().size());
+  EXPECT_FALSE(was_callback_called());
+
+  // Simulate WillStartRequest. The headers should have been returned in the
+  // right order.
+  SimulateWillStartRequest();
+  EXPECT_TRUE(was_callback_called());
+  NavigationHandleImpl::ExtraHeadersList headers = extra_headers();
+  ASSERT_TRUE(headers.size() == 2);
+  EXPECT_EQ(kKey1, headers[0].first);
+  EXPECT_EQ(kValue1, headers[0].second);
+  EXPECT_EQ(kKey2, headers[1].first);
+  EXPECT_EQ(kValue2, headers[1].second);
+
+  // Simulate WillRedirectRequest. The headers should have been returned in the
+  // right order.
+  SimulateWillRedirectRequest();
+  EXPECT_TRUE(was_callback_called());
+  headers = extra_headers();
+  ASSERT_TRUE(headers.size() == 2);
+  EXPECT_EQ(kKey1, headers[0].first);
+  EXPECT_EQ(kValue1, headers[0].second);
+  EXPECT_EQ(kKey2, headers[1].first);
+  EXPECT_EQ(kValue2, headers[1].second);
+}
+
 }  // namespace content