Add a function to add extra headers from NavigationThrottle

content/browser/frame_host/navigation_handle_impl_unittest.cc

Index: content/browser/frame_host/navigation_handle_impl_unittest.cc
diff --git a/content/browser/frame_host/navigation_handle_impl_unittest.cc b/content/browser/frame_host/navigation_handle_impl_unittest.cc
index 4e6ade700856bbf3ae76448c97ad3ea554afb9f1..5a6bb5abb60311801e0ddcef13cfbd16bd3f2986 100644
--- a/content/browser/frame_host/navigation_handle_impl_unittest.cc
+++ b/content/browser/frame_host/navigation_handle_impl_unittest.cc
@@ -45,6 +45,33 @@ class TestNavigationThrottle : public NavigationThrottle {
   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()
@@ -81,6 +108,7 @@ class NavigationHandleImplTest : public RenderViewHostImplTestHarness {
   // 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
@@ -98,6 +126,7 @@ class NavigationHandleImplTest : public RenderViewHostImplTestHarness {
   // 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
@@ -130,18 +159,32 @@ class NavigationHandleImplTest : public RenderViewHostImplTestHarness {
     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.
@@ -487,4 +530,41 @@ TEST_F(NavigationHandleImplTest, CancelThenProceedWillRedirectRequest) {
   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