Rework HTTPS tests and assert location header

1 files changed, 28 insertions, 19 deletions

diff --git a/internal/handlers/https_test.go b/internal/handlers/https_test.go
index df78ecbd..c5ec812c 100644
--- a/internal/handlers/https_test.go
+++ b/internal/handlers/https_test.go
@@ -2,6 +2,7 @@ package handlers_test
 
 import (
 	"net/http"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -9,7 +10,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-pages/internal/handlers"
 )
 
-func TestAcmeMiddleware(t *testing.T) {
+func TestHTTPSMiddleware(t *testing.T) {
 	h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusAccepted)
 	})
@@ -17,42 +18,50 @@ func TestAcmeMiddleware(t *testing.T) {
 	httpsURL := "https://example.com"
 	httpURL := "http://example.com"
 
-	testCases := []struct {
-		name           string
-		redirect       bool
-		path           string
-		expectedStatus int
+	testCases := map[string]struct {
+		redirect         bool
+		path             string
+		expectedStatus   int
+		expectedLocation string
 	}{
-		{
-			name:           "http redirects to https with redirect enabled",
-			redirect:       true,
-			path:           httpURL,
-			expectedStatus: http.StatusTemporaryRedirect,
+		"http redirects to https with redirect enabled": {
+			redirect:         true,
+			path:             httpURL,
+			expectedStatus:   http.StatusTemporaryRedirect,
+			expectedLocation: httpsURL,
 		},
-		{
-			name:           "https handled successfully with redirect enabled",
+		"https handled successfully with redirect enabled": {
 			redirect:       true,
 			path:           httpsURL,
 			expectedStatus: http.StatusAccepted,
 		},
-		{
-			name:           "http does not redirect to https with redirect disabled",
+		"http does not redirect to https with redirect disabled": {
 			redirect:       false,
 			path:           httpURL,
 			expectedStatus: http.StatusAccepted,
 		},
-		{
-			name:           "https handled successfully with redirect disabled",
+		"https handled successfully with redirect disabled": {
 			redirect:       false,
 			path:           httpsURL,
 			expectedStatus: http.StatusAccepted,
 		},
 	}
 
-	for _, tc := range testCases {
-		t.Run(tc.name, func(t *testing.T) {
+	for name, tc := range testCases {
+		t.Run(name, func(t *testing.T) {
 			m := handlers.HTTPSRedirectMiddleware(h, tc.redirect)
 			require.HTTPStatusCode(t, m.ServeHTTP, http.MethodGet, tc.path, nil, tc.expectedStatus)
+
+			// if we expected a redirect make sure the location header is correct
+			if tc.expectedStatus == http.StatusTemporaryRedirect {
+				w := httptest.NewRecorder()
+				req, err := http.NewRequest(http.MethodGet, tc.path, nil)
+				require.NoError(t, err)
+
+				m.ServeHTTP(w, req)
+
+				require.Equal(t, []string{httpsURL}, w.Result().Header["Location"])
+			}
 		})
 	}
 }