feat: add response wrapper (#27)

Author: mojixcoder

context.go

@@ -10,7 +10,7 @@ import (
 // It holds data related to current HTTP request.
 type Context struct {
 	request  *http.Request
-	response http.ResponseWriter
+	response ResponseWriter
 	params   Params
 	storage  Map
 	kid      *Kid
@@ -26,7 +26,7 @@ func newContext(k *Kid) *Context {
 // reset resets the context.
 func (c *Context) reset(request *http.Request, response http.ResponseWriter) {
 	c.request = request
-	c.response = response
+	c.response = newResponse(response)
 	c.storage = make(Map)
 	c.params = make(Params)
 }
@@ -42,7 +42,7 @@ func (c *Context) Request() *http.Request {
 }
 
 // Response returns plain response of current HTTP request.
-func (c *Context) Response() http.ResponseWriter {
+func (c *Context) Response() ResponseWriter {
 	return c.response
 }
 
@@ -171,6 +171,7 @@ func (c *Context) HTMLString(code int, tpl string) error {
 // NoContent returns an empty response with the given status code.
 func (c *Context) NoContent(code int) {
 	c.response.WriteHeader(code)
+	c.response.WriteHeaderNow()
 }
 
 // writeContentType sets content type header for response.

context_test.go

@@ -36,11 +36,12 @@ func TestContextReset(t *testing.T) {
 
 	req := httptest.NewRequest(http.MethodGet, "/", nil)
 	res := httptest.NewRecorder()
+	expectedRes := newResponse(res)
 
 	ctx.reset(req, res)
 
 	assert.Equal(t, req, ctx.request)
-	assert.Equal(t, res, ctx.response)
+	assert.Equal(t, expectedRes, ctx.response)
 	assert.Equal(t, make(Map), ctx.storage)
 	assert.Equal(t, make(Params), ctx.params)
 }
@@ -59,10 +60,11 @@ func TestContextResponse(t *testing.T) {
 	ctx := newContext(New())
 
 	res := httptest.NewRecorder()
+	expectedRes := newResponse(res)
 
 	ctx.reset(nil, res)
 
-	assert.Equal(t, res, ctx.Response())
+	assert.Equal(t, expectedRes, ctx.Response())
 }
 
 func TestContextSetParams(t *testing.T) {
@@ -205,10 +207,6 @@ func TestNoContent(t *testing.T) {
 
 	ctx.NoContent(http.StatusNoContent)
 	assert.Equal(t, http.StatusNoContent, res.Code)
-
-	// Once status code is written, it can't be rewritten again.
-	ctx.NoContent(http.StatusOK)
-	assert.Equal(t, http.StatusNoContent, res.Code)
 }
 
 func TestContextReadJSON(t *testing.T) {

html_renderer/html.go

@@ -128,8 +128,6 @@ func (r *defaultHTMLRenderer) loadTemplates() error {
 		return newInternalServerHTTPError(err, err.Error())
 	}
 
-	fmt.Println(r.funcMap)
-
 	for _, templateFile := range templateFiles {
 		name := r.getTemplateName(templateFile)
 		files := getFilesToParse(templateFile, layoutFiles)

kid_test.go

@@ -474,13 +474,13 @@ func TestKidRun(t *testing.T) {
 	})
 
 	go func() {
-		assert.NoError(t, k.Run())
+		assert.NoError(t, k.Run(":8080"))
 	}()
 
 	// Wait for the server to start
 	time.Sleep(5 * time.Millisecond)
 
-	resp, err := http.Get("http://localhost:2376")
+	resp, err := http.Get("http://localhost:8080")
 	assert.NoError(t, err)
 
 	defer resp.Body.Close()

response.go

@@ -0,0 +1,98 @@
+package kid
+
+import (
+	"bufio"
+	"net"
+	"net/http"
+)
+
+type (
+	// ResponseWriter is a wrapper for http.ResponseWriter
+	// to make using http.ResponseWriter's methods easier.
+	ResponseWriter interface {
+		http.ResponseWriter
+		http.Hijacker
+		http.Flusher
+
+		// WriteHeaderNow writes status code.
+		WriteHeaderNow()
+
+		// Size returns number of bytes written to response.
+		Size() int
+
+		// Written returns true if response has already been written otherwise returns false.
+		Written() bool
+	}
+
+	// response implements ResponseWriter.
+	response struct {
+		http.ResponseWriter
+		written bool
+		status  int
+		size    int
+	}
+)
+
+// Verifying interface compliance.
+var _ ResponseWriter = (*response)(nil)
+
+// newResponse returns a new response writer.
+func newResponse(w http.ResponseWriter) ResponseWriter {
+	response := response{
+		ResponseWriter: w,
+		status:         http.StatusOK,
+	}
+	return &response
+}
+
+// WriteHeader sets status code.
+func (r *response) WriteHeader(code int) {
+	if r.Written() {
+		return
+	}
+
+	r.status = code
+}
+
+// WriteHeaderNow writes status code.
+// Status code should already be specified using response.WriteHeader method.
+func (r *response) WriteHeaderNow() {
+	if r.Written() {
+		return
+	}
+
+	r.written = true
+	r.ResponseWriter.WriteHeader(r.status)
+}
+
+// Write writes byte data to response.
+func (r *response) Write(b []byte) (int, error) {
+	r.WriteHeaderNow()
+
+	n, err := r.ResponseWriter.Write(b)
+	r.size += n
+
+	return n, err
+}
+
+// Size returns number of bytes written.
+func (r *response) Size() int {
+	return r.size
+}
+
+// Written returns true if response has already been written otherwise returns false.
+func (r *response) Written() bool {
+	return r.written
+}
+
+// Flush implements the http.Flusher interface.
+func (r *response) Flush() {
+	r.WriteHeaderNow()
+	r.ResponseWriter.(http.Flusher).Flush()
+}
+
+// Hijack implements the http.Hijacker interface.
+func (r *response) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	r.written = true
+	return r.ResponseWriter.(http.Hijacker).Hijack()
+}

response_test.go

@@ -0,0 +1,98 @@
+package kid
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewResponse(t *testing.T) {
+	w := httptest.NewRecorder()
+	res := newResponse(w).(*response)
+
+	assert.Equal(t, w, res.ResponseWriter)
+	assert.Equal(t, http.StatusOK, res.status)
+	assert.Zero(t, res.Size())
+	assert.False(t, res.Written())
+}
+
+func TestResponseWriterWriteHeader(t *testing.T) {
+	w := httptest.NewRecorder()
+	res := newResponse(w).(*response)
+
+	res.WriteHeader(http.StatusAccepted)
+
+	assert.Equal(t, http.StatusAccepted, res.status)
+	assert.False(t, res.Written())
+
+	res.WriteHeaderNow()
+
+	// Won't write again because header is already written.
+	res.WriteHeader(http.StatusBadRequest)
+
+	assert.Equal(t, http.StatusAccepted, res.status)
+}
+
+func TestResponseWriterWriteHeaderNow(t *testing.T) {
+	w := httptest.NewRecorder()
+	res := newResponse(w).(*response)
+
+	res.WriteHeader(http.StatusAccepted)
+	res.WriteHeaderNow()
+
+	assert.True(t, res.Written())
+}
+
+func TestResponseWriterSize(t *testing.T) {
+	w := httptest.NewRecorder()
+	res := newResponse(w)
+
+	n1, err := res.Write([]byte("Hello"))
+	assert.NoError(t, err)
+
+	n2, err := res.Write([]byte("Bye"))
+	assert.NoError(t, err)
+
+	assert.Equal(t, 8, n1+n2)
+	assert.Equal(t, n1+n2, res.Size())
+}
+
+func TestResponseWriterWritten(t *testing.T) {
+	w := httptest.NewRecorder()
+	res := newResponse(w)
+
+	assert.False(t, res.Written())
+
+	res.WriteHeaderNow()
+
+	assert.True(t, res.Written())
+}
+
+func TestResponseWriterFlush(t *testing.T) {
+	k := New()
+
+	k.GET("/", func(c *Context) error {
+		c.Response().WriteHeader(http.StatusBadGateway)
+		c.Response().Flush()
+		return nil
+	})
+
+	srv := httptest.NewServer(k)
+	defer srv.Close()
+
+	resp, err := http.Get(srv.URL)
+	assert.NoError(t, err)
+	assert.Equal(t, http.StatusBadGateway, resp.StatusCode)
+}
+
+func TestResponseWriterHijack(t *testing.T) {
+	w := httptest.NewRecorder()
+	res := newResponse(w)
+
+	assert.Panics(t, func() {
+		_, _, _ = res.Hijack()
+	})
+	assert.True(t, res.Written())
+}