From 1fc5c92023defea7984315acf67feadccf785ec4 Mon Sep 17 00:00:00 2001
From: Vinicius Guedes <viniciusgued@gmail.com>
Date: Thu, 10 Mar 2022 17:54:55 -0300
Subject: [PATCH] fix: setting request content after setting header

---
 handler.go       |  5 +----
 response.go      | 30 ++++++++++--------------------
 response_test.go | 24 ++++--------------------
 3 files changed, 15 insertions(+), 44 deletions(-)

diff --git a/handler.go b/handler.go
index 6285440..cde0c20 100644
--- a/handler.go
+++ b/handler.go
@@ -13,9 +13,6 @@ func Handler(handler EasyHandler) MuxHandler {
 		request := NewRequest(httpRequest)
 		response := NewResponse(responseWriter)
 
-		receivedResponse := handler(request, response)
-		for key, value := range receivedResponse.GetHeaders() {
-			responseWriter.Header().Set(key, value)
-		}
+		handler(request, response)
 	}
 }
diff --git a/response.go b/response.go
index f08ae76..7f19562 100644
--- a/response.go
+++ b/response.go
@@ -22,44 +22,34 @@ func (response *Response) SetResponseWriter(responseWriter http.ResponseWriter)
 }
 
 func (response *Response) SetHeader(key string, value string) *Response {
-	if len(response.headers) == 0 {
-		response.headers = map[string]string{}
-	}
-
-	response.headers[key] = value
+	response.responseWriter.Header().Set(key, value)
 	return response
 }
 
-func (response *Response) GetHeader(key string) string {
-	if value, ok := response.headers[key]; ok {
-		return value
-	}
-
-	return ""
-}
-
-func (response *Response) GetHeaders() map[string]string {
-	return response.headers
-}
-
 func (response *Response) SetStatus(status int) *Response {
 	response.responseWriter.WriteHeader(status)
 	return response
 }
 
 func (response *Response) Json(data interface{}) *Response {
+	response.SetHeader("Content-Type", "application/json")
+
 	responseData, _ := json.Marshal(data)
 	response.responseWriter.Write(responseData)
 
-	return response.SetHeader("Content-Type", "application/json")
+	return response
 }
 
 func (response *Response) Html(data string) *Response {
+	response.SetHeader("Content-Type", "text/html")
 	response.responseWriter.Write([]byte(data))
-	return response.SetHeader("Content-Type", "text/html")
+
+	return response
 }
 
 func (response *Response) Text(data string) *Response {
+	response.SetHeader("Content-Type", "plain/text")
 	response.responseWriter.Write([]byte(data))
-	return response.SetHeader("Content-Type", "plain/text")
+
+	return response
 }
diff --git a/response_test.go b/response_test.go
index 6f07959..e2b0c45 100644
--- a/response_test.go
+++ b/response_test.go
@@ -15,23 +15,7 @@ func TestResponse_SetHeader(t *testing.T) {
 	response.SetResponseWriter(responseWriter).
 		SetHeader("Hello", "world")
 
-	assert.Equal(t, "world", response.GetHeader("Hello"))
-}
-
-func TestResponse_GetHeaders(t *testing.T) {
-	responseWriter := httptest.NewRecorder()
-
-	response := Response{}
-	response.SetResponseWriter(responseWriter).
-		SetHeader("Hello", "world")
-
-	assert.Equal(
-		t,
-		map[string]string{
-			"Hello": "world",
-		},
-		response.GetHeaders(),
-	)
+	assert.Equal(t, "world", responseWriter.Header().Get("Hello"))
 }
 
 func TestResponse_SetStatus(t *testing.T) {
@@ -54,7 +38,7 @@ func TestResponse_Json(t *testing.T) {
 		})
 
 	assert.Equal(t, "{\"hello\":\"world\"}", responseWriter.Body.String())
-	assert.Equal(t, "application/json", response.GetHeader("Content-Type"))
+	assert.Equal(t, "application/json", responseWriter.Header().Get("Content-Type"))
 }
 
 func TestResponse_Html(t *testing.T) {
@@ -65,7 +49,7 @@ func TestResponse_Html(t *testing.T) {
 		Html("<h1>Hello</h1>")
 
 	assert.Equal(t, "<h1>Hello</h1>", responseWriter.Body.String())
-	assert.Equal(t, "text/html", response.GetHeader("Content-Type"))
+	assert.Equal(t, "text/html", responseWriter.Header().Get("Content-Type"))
 }
 
 func TestResponse_Text(t *testing.T) {
@@ -76,5 +60,5 @@ func TestResponse_Text(t *testing.T) {
 		Text("oi ne")
 
 	assert.Equal(t, "oi ne", responseWriter.Body.String())
-	assert.Equal(t, "plain/text", response.GetHeader("Content-Type"))
+	assert.Equal(t, "plain/text", responseWriter.Header().Get("Content-Type"))
 }