From 60d744bbf97c628c3efd986afbbce1a2c03583f0 Mon Sep 17 00:00:00 2001
From: Vinicius Guedes <viniciusgued@gmail.com>
Date: Thu, 10 Mar 2022 18:44:03 -0300
Subject: [PATCH] feat: late set response data/body

---
 handler.go       |  1 +
 response.go      | 23 +++++++++++------------
 response_test.go |  6 +++---
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/handler.go b/handler.go
index 2d02410..0debfb8 100644
--- a/handler.go
+++ b/handler.go
@@ -18,5 +18,6 @@ func Handler(handler EasyHandler) MuxHandler {
 		for key, value := range receivedResponse.GetHeaders() {
 			responseWriter.Header().Set(key, value)
 		}
+		responseWriter.Write([]byte(receivedResponse.GetData()))
 	}
 }
diff --git a/response.go b/response.go
index c1f1c26..55dfed1 100644
--- a/response.go
+++ b/response.go
@@ -9,6 +9,7 @@ type Response struct {
 	headers        map[string]string
 	responseWriter http.ResponseWriter
 	status         int
+	data           string
 }
 
 func NewResponse(responseWriter http.ResponseWriter) Response {
@@ -60,25 +61,23 @@ func (response *Response) GetStatus() int {
 	return response.status
 }
 
-func (response *Response) Json(data interface{}) *Response {
-	response.SetHeader("Content-Type", "application/json")
+func (response *Response) GetData() string {
+	return response.data
+}
 
+func (response *Response) Json(data interface{}) *Response {
 	responseData, _ := json.Marshal(data)
-	response.responseWriter.Write(responseData)
+	response.data = string(responseData)
 
-	return response
+	return response.SetHeader("Content-Type", "application/json")
 }
 
 func (response *Response) Html(data string) *Response {
-	response.SetHeader("Content-Type", "text/html")
-	response.responseWriter.Write([]byte(data))
-
-	return response
+	response.data = data
+	return response.SetHeader("Content-Type", "text/html")
 }
 
 func (response *Response) Text(data string) *Response {
-	response.SetHeader("Content-Type", "plain/text")
-	response.responseWriter.Write([]byte(data))
-
-	return response
+	response.data = data
+	return response.SetHeader("Content-Type", "plain/text")
 }
diff --git a/response_test.go b/response_test.go
index 71e8c51..c514dac 100644
--- a/response_test.go
+++ b/response_test.go
@@ -43,7 +43,7 @@ func TestResponse_Json(t *testing.T) {
 			"hello": "world",
 		})
 
-	assert.Equal(t, "{\"hello\":\"world\"}", responseWriter.Body.String())
+	assert.Equal(t, "{\"hello\":\"world\"}", response.GetData())
 	assert.Equal(t, "application/json", response.GetHeader("Content-Type"))
 }
 
@@ -54,7 +54,7 @@ func TestResponse_Html(t *testing.T) {
 	response.SetResponseWriter(responseWriter).
 		Html("<h1>Hello</h1>")
 
-	assert.Equal(t, "<h1>Hello</h1>", responseWriter.Body.String())
+	assert.Equal(t, "<h1>Hello</h1>", response.GetData())
 	assert.Equal(t, "text/html", response.GetHeader("Content-Type"))
 }
 
@@ -65,6 +65,6 @@ func TestResponse_Text(t *testing.T) {
 	response.SetResponseWriter(responseWriter).
 		Text("oi ne")
 
-	assert.Equal(t, "oi ne", responseWriter.Body.String())
+	assert.Equal(t, "oi ne", response.GetData())
 	assert.Equal(t, "plain/text", response.GetHeader("Content-Type"))
 }