From c34132114558bd7e4ff18b07b76bda401cc06d10 Mon Sep 17 00:00:00 2001
From: Kelvin Ly <kelvin.ly1618@gmail.com>
Date: Wed, 20 Nov 2019 09:06:04 -0500
Subject: [PATCH] Fix bugs and refactor response collector out into its own
 file TODO write tests for everything

---
 cache.go     | 48 ++----------------------------------------------
 collector.go | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 resize.go    |  4 ++--
 3 files changed, 53 insertions(+), 48 deletions(-)
 create mode 100644 collector.go

diff --git a/cache.go b/cache.go
index c764c6d..3ef0c7f 100644
--- a/cache.go
+++ b/cache.go
@@ -4,50 +4,6 @@ import (
 	"net/http"
 )
 
-type Response struct {
-	Code    int
-	Headers map[string][]string
-	Body    []byte
-}
-
-func (r Response) WriteResponse(rw http.ResponseWriter) {
-	for k, vs := range r.Headers {
-		for _, v := range vs {
-			rw.Header().Add(k, v)
-		}
-	}
-	rw.WriteHeader(r.Code)
-	rw.Write(r.Body)
-}
-
-// implements ResponseWriter to collect HTTP responses
-type ResponseCollector struct {
-	Response
-}
-
-func NewResponseCollector() *ResponseCollector {
-	return &ResponseCollector{
-		Response{Code: 200, Headers: make(map[string][]string)},
-	}
-}
-
-func (rc *ResponseCollector) Header() http.Header {
-	return rc.Headers
-}
-
-func (rc *ResponseCollector) Write(bs []byte) (int, error) {
-	rc.Body = append(rc.Body, bs...)
-	return len(bs), nil
-}
-
-func (rc *ResponseCollector) WriteHeader(code int) {
-	rc.Code = code
-}
-
-func (rc *ResponseCollector) CollectResponse() Response {
-	return rc.Response
-}
-
 type cacheEntry struct {
 	r Response
 }
@@ -66,10 +22,10 @@ func Cache(h http.Handler) http.Handler {
 		if exists {
 			entry.r.WriteResponse(rw)
 		} else {
-			rc := NewResponseCollector()
+			rc := ResponseCollector{}
 			// copy request in case they modify it
 			req := *r
-			h.ServeHTTP(rc, &req)
+			h.ServeHTTP(&rc, &req)
 			resp := rc.CollectResponse()
 			if resp.Code == 200 {
 				c[r.URL.String()] = cacheEntry{resp}
diff --git a/collector.go b/collector.go
new file mode 100644
index 0000000..316d2ca
--- /dev/null
+++ b/collector.go
@@ -0,0 +1,49 @@
+package main
+
+import (
+	"net/http"
+)
+
+type Response struct {
+	Code    int
+	Headers map[string][]string
+	Body    []byte
+}
+
+func (r Response) WriteResponse(rw http.ResponseWriter) {
+	for k, vs := range r.Headers {
+		for _, v := range vs {
+			rw.Header().Add(k, v)
+		}
+	}
+	rw.WriteHeader(r.Code)
+	rw.Write(r.Body)
+}
+
+// implements ResponseWriter to collect HTTP responses
+type ResponseCollector struct {
+	Response
+}
+
+func (rc *ResponseCollector) Header() http.Header {
+	if rc.Headers == nil {
+		rc.Headers = make(map[string][]string)
+	}
+	return rc.Headers
+}
+
+func (rc *ResponseCollector) Write(bs []byte) (int, error) {
+	if rc.Code == 0 {
+		rc.Code = 200
+	}
+	rc.Body = append(rc.Body, bs...)
+	return len(bs), nil
+}
+
+func (rc *ResponseCollector) WriteHeader(code int) {
+	rc.Code = code
+}
+
+func (rc *ResponseCollector) CollectResponse() Response {
+	return rc.Response
+}
diff --git a/resize.go b/resize.go
index baaab9a..4fcdfd4 100644
--- a/resize.go
+++ b/resize.go
@@ -12,9 +12,9 @@ import (
 
 func Resize(maxWidth uint, h http.Handler) http.Handler {
 	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
-		rc := NewResponseCollector()
+		rc := ResponseCollector{}
 		req := *r
-		h.ServeHTTP(rc, &req)
+		h.ServeHTTP(&rc, &req)
 		imageResp := rc.CollectResponse()
 
 		if imageResp.Code != 200 {