1 files changed, 106 insertions, 0 deletions
diff --git a/chrome_frame/test/http_negotiate_unittest.cc b/chrome_frame/test/http_negotiate_unittest.cc
index 857197e..c56034b 100644
--- a/chrome_frame/test/http_negotiate_unittest.cc
+++ b/chrome_frame/test/http_negotiate_unittest.cc
@@ -16,6 +16,112 @@
 #include "gtest/gtest.h"
 #include "gmock/gmock.h"
 
+class HttpNegotiateTest : public testing::Test {
+ protected:
+  HttpNegotiateTest() {
+  }
+};
+
+class TestHttpNegotiate
+    : public CComObjectRootEx<CComMultiThreadModel>,
+      public IHttpNegotiate {
+ public:
+  TestHttpNegotiate()
+      : beginning_transaction_ret_(S_OK), additional_headers_(NULL) {
+  }
+
+BEGIN_COM_MAP(TestHttpNegotiate)
+  COM_INTERFACE_ENTRY(IHttpNegotiate)
+END_COM_MAP()
+  STDMETHOD(BeginningTransaction)(LPCWSTR url, LPCWSTR headers,  // NOLINT
+                                  DWORD reserved,  // NOLINT
+                                  LPWSTR* additional_headers) {  // NOLINT
+    if (additional_headers_) {
+      int len = lstrlenW(additional_headers_);
+      len++;
+      *additional_headers = reinterpret_cast<wchar_t*>(
+          ::CoTaskMemAlloc(len * sizeof(wchar_t)));
+      lstrcpyW(*additional_headers, additional_headers_);
+    }
+    return beginning_transaction_ret_;
+  }
+
+  STDMETHOD(OnResponse)(DWORD response_code, LPCWSTR response_header,
+                        LPCWSTR request_header,
+                        LPWSTR* additional_request_headers) {
+    return S_OK;
+  }
+
+  HRESULT beginning_transaction_ret_;
+  const wchar_t* additional_headers_;
+};
+
+TEST_F(HttpNegotiateTest, BeginningTransaction) {
+  static const int kBeginningTransactionIndex = 3;
+  CComObjectStackEx<TestHttpNegotiate> test_http;
+  IHttpNegotiate_BeginningTransaction_Fn original =
+      reinterpret_cast<IHttpNegotiate_BeginningTransaction_Fn>(
+          (*reinterpret_cast<void***>(
+              static_cast<IHttpNegotiate*>(
+                  &test_http)))[kBeginningTransactionIndex]);
+
+  std::wstring cf_ua(
+      ASCIIToWide(http_utils::GetDefaultUserAgentHeaderWithCFTag()));
+  std::wstring cf_tag(
+      ASCIIToWide(http_utils::GetChromeFrameUserAgent()));
+
+  EXPECT_NE(std::wstring::npos, cf_ua.find(cf_tag));
+
+  struct TestCase {
+    const std::wstring original_headers_;
+    const std::wstring delegate_additional_;
+    const std::wstring expected_additional_;
+    HRESULT delegate_return_value_;
+  } test_cases[] = {
+    { L"Accept: */*\r\n",
+      L"",
+      cf_ua + L"\r\n",
+      S_OK },
+    { L"Accept: */*\r\n",
+      L"",
+      L"",
+      E_OUTOFMEMORY },
+    { L"",
+      L"Accept: */*\r\n",
+      L"Accept: */*\r\n" + cf_ua + L"\r\n",
+      S_OK },
+    { L"User-Agent: Bingo/1.0\r\n",
+      L"",
+      L"User-Agent: Bingo/1.0 " + cf_tag + L"\r\n",
+      S_OK },
+    { L"User-Agent: NotMe/1.0\r\n",
+      L"User-Agent: MeMeMe/1.0\r\n",
+      L"User-Agent: MeMeMe/1.0 " + cf_tag + L"\r\n",
+      S_OK },
+    { L"",
+      L"User-Agent: MeMeMe/1.0\r\n",
+      L"User-Agent: MeMeMe/1.0 " + cf_tag + L"\r\n",
+      S_OK },
+  };
+
+  for (int i = 0; i < arraysize(test_cases); ++i) {
+    TestCase& test = test_cases[i];
+    wchar_t* additional = NULL;
+    test_http.beginning_transaction_ret_ = test.delegate_return_value_;
+    test_http.additional_headers_ = test.delegate_additional_.c_str();
+    HttpNegotiatePatch::BeginningTransaction(original, &test_http,
+        L"http://www.google.com", test.original_headers_.c_str(), 0,
+        &additional);
+    EXPECT_TRUE(additional != NULL);
+
+    if (additional) {
+      // Check against the expected additional headers.
+      EXPECT_EQ(test.expected_additional_, std::wstring(additional));
+      ::CoTaskMemFree(additional);
+    }
+  }
+}
+
 class TestInternetProtocolSink
     : public CComObjectRootEx<CComMultiThreadModel>,
       public IInternetProtocolSink {