Small refactor + tests

Author: jonahgeorge

force-ssl-heroku.go renamed to forcessl.go

@@ -5,10 +5,17 @@ import (
 	"os"
 )
 
+var getenv = os.Getenv
+
+const (
+	xForwardedProtoHeader = "x-forwarded-proto"
+	goEnviron             = "GO_ENV"
+)
+
 func ForceSsl(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if os.Getenv("GO_ENV") == "production" {
-			if r.Header.Get("x-forwarded-proto") != "https" {
+		if getenv(goEnviron) == "production" {
+			if r.Header.Get(xForwardedProtoHeader) != "https" {
 				sslUrl := "https://" + r.Host + r.RequestURI
 				http.Redirect(w, r, sslUrl, http.StatusTemporaryRedirect)
 				return
@@ -17,5 +24,4 @@ func ForceSsl(next http.Handler) http.Handler {
 
 		next.ServeHTTP(w, r)
 	})
-
 }

forcessl_test.go

@@ -0,0 +1,49 @@
+package forcesslheroku
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+var testCases = []struct {
+	goEnv     string
+	proto     string
+	expectLoc string
+}{
+	{goEnv: "production", proto: "http",
+		expectLoc: "https://example.com/test"},
+	{goEnv: "production", proto: "https"},
+	{goEnv: "test", proto: "http"},
+	{goEnv: "test", proto: "https"},
+}
+
+func TestForceSsl(t *testing.T) {
+	noopHandler := func(w http.ResponseWriter, r *http.Request) {}
+	forceSsl := ForceSsl(http.HandlerFunc(noopHandler))
+
+	for _, tt := range testCases {
+		getenv = func(key string) string {
+			switch key {
+			case goEnviron:
+				return tt.goEnv
+			default:
+				return ""
+			}
+		}
+
+		t.Run(tt.goEnv+"_"+tt.proto, func(t *testing.T) {
+			req := httptest.NewRequest("", "/test", nil)
+			req.Header.Add(xForwardedProtoHeader, tt.proto)
+
+			res := httptest.NewRecorder()
+			forceSsl.ServeHTTP(res, req)
+
+			if location := res.Header().Get("Location"); location != tt.expectLoc {
+				t.Errorf("expected Location header '%s', got '%s'",
+					tt.expectLoc, location)
+			}
+		})
+	}
+}
+