From 25084ff92d69570ad2bf1034c7a6d8489fbdd14e Mon Sep 17 00:00:00 2001
From: Nicolas Azrak <azrak.nicolas@gmail.com>
Date: Mon, 24 Oct 2016 22:23:58 -0300
Subject: [PATCH] Fixes #2 Change net.http/Headers.set in favor of Add()

---
 cache.go      |  4 +---
 cache_test.go | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/cache.go b/cache.go
index f83e56c..29dd486 100644
--- a/cache.go
+++ b/cache.go
@@ -22,7 +22,7 @@ type CacheHandler struct {
 func respond(response * storage.CachedResponse, w http.ResponseWriter) {
 	for k, values := range response.HeaderMap {
 		for _, v := range values {
-			w.Header().Set(k, v)
+			w.Header().Add(k, v)
 		}
 	}
 	w.WriteHeader(response.Code)
@@ -32,7 +32,6 @@ func respond(response * storage.CachedResponse, w http.ResponseWriter) {
 func shouldUseCache(r *http.Request, config *Config) bool {
 	// TODO Add more logic like get params, ?nocache=true
 
-
 	if r.Method != "GET" && r.Method != "HEAD" {
 		// Only cache Get and head request
 		return false
@@ -79,7 +78,6 @@ func getCacheableStatus(req *http.Request, res *httptest.ResponseRecorder, confi
 	cacheobject.ExpirationObject(&obj, &rv)
 
 	isCacheable := len(rv.OutReasons) == 0
-
 	expiration := rv.OutExpirationTime
 	if expiration.Before(time.Now().UTC().Add(time.Duration(1) * time.Second)) {
 		// If expiration is before now use default MaxAge
diff --git a/cache_test.go b/cache_test.go
index 882f720..0516200 100644
--- a/cache_test.go
+++ b/cache_test.go
@@ -19,20 +19,22 @@ type TestHandler struct {
 
 func (h *TestHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) (int, error) {
 	h.timesCalled = h.timesCalled + 1
-	w.Write(h.ResponseBody)
-	w.WriteHeader(h.ResponseCode)
 	for k, values := range h.ResponseHeaders {
 		for _, v := range values {
-			w.Header().Set(k, v)
+			w.Header().Add(k, v)
 		}
 	}
+	w.WriteHeader(h.ResponseCode)
+	w.Write(h.ResponseBody)
 	return h.ResponseCode, h.ResponseError
 }
 
 func buildBasicHandler(cacheablePaths string) (*CacheHandler, *TestHandler) {
 	memory := storage.MemoryStorage{}
 	memory.Setup()
-	backend := TestHandler{}
+	backend := TestHandler{
+		ResponseCode: 200,
+	}
 
 	return &CacheHandler{
 		Config: &Config {
@@ -53,6 +55,8 @@ func buildGetRequest(path string) *http.Request {
 
 }
 
+
+// TODO avoid code duplication, use r.Run
 func TestBasicCache(t *testing.T) {
 	handler, backend := buildBasicHandler("/assets")
 	rec := httptest.NewRecorder()
@@ -106,7 +110,7 @@ func TestNotCacheableCacheControl(t *testing.T) {
 	handler, backend := buildBasicHandler("/assets")
 	rec := httptest.NewRecorder()
 
-	responseHeaders := make(map[string][]string)
+	responseHeaders := make(http.Header)
 	responseHeaders["Cache-control"] = []string { "private" }
 	backend.ResponseHeaders = responseHeaders
 
@@ -121,6 +125,27 @@ func TestNotCacheableCacheControl(t *testing.T) {
 	assert.Equal(t, 2, backend.timesCalled, "Backend should have been called 2 but it was called", backend.timesCalled)
 }
 
+func TestAddHeaders(t *testing.T) {
+	handler, backend := buildBasicHandler("/assets")
+
+	responseHeaders := make(http.Header)
+	responseHeaders["Content-Type"] = []string { "text/plain; charset=utf-8" }
+	responseHeaders["X-Custom-2"] = []string { "bar", "baz" }
+	responseHeaders["X-Custom"] = []string { "foo", "bar", "baz" }
+	backend.ResponseHeaders = responseHeaders
+
+	req := buildGetRequest("http://somehost.com/assets/1")
+
+	rec := httptest.NewRecorder()
+	_, err := handler.ServeHTTP(rec, req)
+
+	if err != nil {
+		assert.Fail(t, "Error processing request", err)
+	}
+
+	assert.Equal(t, responseHeaders, rec.HeaderMap, "Cache didn't send same headers that backend originally sent")
+}
+
 func TestDefaultCacheTime(t *testing.T) {
 	// TODO test this
 	// isCacheable, expiration := getCacheableStatus(req, res, config)