2 files changed, 28 insertions, 3 deletions
diff --git a/net/spdy/spdy_websocket_test_util.cc b/net/spdy/spdy_websocket_test_util.cc
index 2327ee2..0872e3f 100644
--- a/net/spdy/spdy_websocket_test_util.cc
+++ b/net/spdy/spdy_websocket_test_util.cc
@@ -19,11 +19,17 @@ static const int kDefaultExtraHeaderCount = 0;
 SpdyWebSocketTestUtil::SpdyWebSocketTestUtil(
     NextProto protocol) : spdy_util_(protocol) {}
 
+std::string SpdyWebSocketTestUtil::GetHeader(const SpdyHeaderBlock& headers,
+                                             const std::string& key) const {
+  SpdyHeaderBlock::const_iterator it = headers.find(GetHeaderKey(key));
+  return (it == headers.end()) ? "" : it->second;
+}
+
 void SpdyWebSocketTestUtil::SetHeader(
     const std::string& key,
     const std::string& value,
     SpdyHeaderBlock* headers) const {
-  (*headers)[(spdy_util_.is_spdy2() ? "" : ":") + key] = value;
+  (*headers)[GetHeaderKey(key)] = value;
 }
 
 SpdyFrame* SpdyWebSocketTestUtil::ConstructSpdyWebSocketSynStream(
@@ -146,4 +152,13 @@ SpdyFrame* SpdyWebSocketTestUtil::ConstructSpdySettings(
   return spdy_util_.ConstructSpdySettings(settings);
 }
 
+SpdyMajorVersion SpdyWebSocketTestUtil::spdy_version() const {
+  return spdy_util_.spdy_version();
+}
+
+std::string SpdyWebSocketTestUtil::GetHeaderKey(
+    const std::string& key) const {
+  return (spdy_util_.is_spdy2() ? "" : ":") + key;
+}
+
 }  // namespace net
diff --git a/net/spdy/spdy_websocket_test_util.h b/net/spdy/spdy_websocket_test_util.h
index f334919..7a9a59e 100644
--- a/net/spdy/spdy_websocket_test_util.h
+++ b/net/spdy/spdy_websocket_test_util.h
@@ -6,6 +6,7 @@
 #define NET_SPDY_SPDY_WEBSOCKET_TEST_UTIL_H_
 
 #include "net/base/request_priority.h"
+#include "net/spdy/spdy_header_block.h"
 #include "net/spdy/spdy_protocol.h"
 #include "net/spdy/spdy_test_util_common.h"
 
@@ -15,8 +16,13 @@ class SpdyWebSocketTestUtil {
  public:
   explicit SpdyWebSocketTestUtil(NextProto protocol);
 
-  // Adds the given key/value pair to |headers|, tweaking it depending
-  // on SPDY version.
+  // Returns the value corresponding to the given key (passed through
+  // GetHeaderKey()), or the empty string if none exists.
+  std::string GetHeader(const SpdyHeaderBlock& headers,
+                        const std::string& key) const;
+
+  // Adds the given key/value pair to |headers|, passing the key
+  // through GetHeaderKey().
   void SetHeader(const std::string& key,
                  const std::string& value,
                  SpdyHeaderBlock* headers) const;
@@ -57,8 +63,12 @@ class SpdyWebSocketTestUtil {
 
   // Forwards to |spdy_util_|.
   SpdyFrame* ConstructSpdySettings(const SettingsMap& settings) const;
+  SpdyMajorVersion spdy_version() const;
 
  private:
+  // Modify the header key based on the SPDY version and return it.
+  std::string GetHeaderKey(const std::string& key) const;
+
   SpdyTestUtil spdy_util_;
 };