Allow to AddOrReplace headers from THttpResponse

2 files changed, 36 insertions, 0 deletions

diff --git a/library/cpp/http/server/response.h b/library/cpp/http/server/response.h
index a75cb85605..8b7819de1e 100644
--- a/library/cpp/http/server/response.h
+++ b/library/cpp/http/server/response.h
@@ -32,6 +32,17 @@ public:
         return *this;
     }
 
+    template <typename ValueType>
+    THttpResponse& AddOrReplaceHeader(const TString& name, const ValueType& value) {
+        return AddOrReplaceHeader(THttpInputHeader(name, ToString(value)));
+    }
+
+    THttpResponse& AddOrReplaceHeader(const THttpInputHeader& header) {
+        Headers.AddOrReplaceHeader(header);
+
+        return *this;
+    }
+
     THttpResponse& AddMultipleHeaders(const THttpHeaders& headers);
 
     const THttpHeaders& GetHeaders() const {
diff --git a/library/cpp/http/server/response_ut.cpp b/library/cpp/http/server/response_ut.cpp
index 73e2112ad3..bdf2fc1c7d 100644
--- a/library/cpp/http/server/response_ut.cpp
+++ b/library/cpp/http/server/response_ut.cpp
@@ -139,4 +139,29 @@ Y_UNIT_TEST_SUITE(TestHttpResponse) {
             .SetContentType("text/xml");
         UNIT_ASSERT_STRINGS_EQUAL(ToString(resp), EXPECTED);
     }
+
+    Y_UNIT_TEST(TestAddOrReplaceHeader) {
+        THttpResponse resp(HTTP_FORBIDDEN);
+        resp.AddOrReplaceHeader(THttpInputHeader("X-Header-1", "ValueOne"));
+        resp.AddOrReplaceHeader("X-Header-2", 10);
+        resp.AddOrReplaceHeader("X-Header-3", true);
+
+        const char* EXPECTED_ADDED = "HTTP/1.1 403 Forbidden\r\n"
+                               "X-Header-1: ValueOne\r\n"
+                               "X-Header-2: 10\r\n"
+                               "X-Header-3: 1\r\n"
+                               "\r\n";
+        UNIT_ASSERT_STRINGS_EQUAL(ToString(resp), EXPECTED_ADDED);
+
+        resp.AddOrReplaceHeader(THttpInputHeader("X-Header-1", "ValueTwo"));
+        resp.AddOrReplaceHeader("X-Header-2", 20);
+        resp.AddOrReplaceHeader("X-Header-3", false);
+
+        const char* EXPECTED_REPLACED = "HTTP/1.1 403 Forbidden\r\n"
+                               "X-Header-1: ValueTwo\r\n"
+                               "X-Header-2: 20\r\n"
+                               "X-Header-3: 0\r\n"
+                               "\r\n";
+        UNIT_ASSERT_STRINGS_EQUAL(ToString(resp), EXPECTED_REPLACED);
+    }
 }