Refactor redirect handling logic and add unit tests

Author: ShocOne

redirecthandler/redirecthandler.go

@@ -47,7 +47,9 @@ func (r *RedirectHandler) WithRedirectHandling(client *http.Client) {
 
 // checkRedirect implements the redirect handling logic.
 func (r *RedirectHandler) checkRedirect(req *http.Request, via []*http.Request) error {
-	defer r.clearRedirectHistory(req) // Ensure redirect history is always cleared to prevent memory leaks
+
+	// Ensure redirect history is always cleared to prevent memory leaks
+	defer r.clearRedirectHistory(req)
 
 	// Non-idempotent methods handling
 	if req.Method == http.MethodPost || req.Method == http.MethodPatch {

redirecthandler/redirecthandler_test.go

@@ -14,70 +14,79 @@ import (
 // It covers various scenarios including redirect loop detection, maximum redirects limit,
 // resolving relative redirects, cross-domain security measures, and handling of 303 See Other response.
 func TestRedirectHandler_CheckRedirect(t *testing.T) {
-	mockLogger := mocklogger.NewMockLogger()
-
-	// Set the mock logger to capture logs at all levels
-	mockLogger.SetLevel(logger.LogLevelDebug)
-
-	redirectHandler := NewRedirectHandler(mockLogger, 10)
+	redirectHandler := NewRedirectHandler(nil, 10) // Logger is not needed for these tests
 
 	reqURL, _ := url.Parse("http://example.com")
 	req := &http.Request{URL: reqURL, Method: http.MethodPost}
-	resp := &http.Response{
-		Status:     "303 See Other",
-		StatusCode: http.StatusSeeOther,
-		Header:     http.Header{"Location": []string{"http://example.com/new"}},
-	}
-
-	t.Run("Redirect Loop Detection", func(t *testing.T) {
-		redirectHandler.VisitedURLs = map[string]int{"http://example.com": 1}
-		err := redirectHandler.checkRedirect(req, []*http.Request{{}, {}})
-		assert.Equal(t, http.ErrUseLastResponse, err)
-		// Verify that a warning log for redirect loop was recorded
-		assert.Contains(t, mockLogger.Calls[0].Arguments.String(0), "Detected redirect loop")
-	})
-
-	t.Run("Maximum Redirects Reached", func(t *testing.T) {
-		redirectHandler.VisitedURLs = map[string]int{}
-		redirectHandler.MaxRedirects = 1
-		err := redirectHandler.checkRedirect(req, []*http.Request{{}, {}})
-		assert.Equal(t, http.ErrUseLastResponse, err)
-		// Verify that a warning log for max redirects was recorded
-		assert.Contains(t, mockLogger.Calls[1].Arguments.String(0), "Stopped after maximum redirects")
-	})
 
-	t.Run("Resolve Relative Redirects", func(t *testing.T) {
-		redirectHandler.MaxRedirects = 10
-		err := redirectHandler.checkRedirect(req, []*http.Request{{}, {}})
-		assert.Nil(t, err)
-		assert.Equal(t, "http://example.com/new", req.URL.String())
-	})
+	// Test cases
+	tests := []struct {
+		name        string
+		prepare     func() *http.Response // Function to prepare the response for each test case
+		expectedErr error
+		expectedURL string
+	}{
+		{
+			name: "Redirect Loop Detection",
+			prepare: func() *http.Response {
+				redirectHandler.VisitedURLs = map[string]int{"http://example.com": 1}
+				return nil
+			},
+			expectedErr: http.ErrUseLastResponse,
+		},
+		{
+			name: "Maximum Redirects Reached",
+			prepare: func() *http.Response {
+				redirectHandler.VisitedURLs = map[string]int{}
+				redirectHandler.MaxRedirects = 1
+				return nil
+			},
+			expectedErr: http.ErrUseLastResponse,
+		},
+		{
+			name: "Resolve Relative Redirects",
+			prepare: func() *http.Response {
+				redirectHandler.MaxRedirects = 10
+				return &http.Response{
+					StatusCode: http.StatusSeeOther,
+					Header:     http.Header{"Location": []string{"http://example.com/new"}},
+				}
+			},
+			expectedURL: "http://example.com/new",
+		},
+		{
+			name: "Cross-Domain Security Measures",
+			prepare: func() *http.Response {
+				return &http.Response{
+					Header: http.Header{"Location": []string{"http://anotherdomain.com/new"}},
+				}
+			},
+			expectedErr: nil,
+		},
+		{
+			name: "Handling 303 See Other",
+			prepare: func() *http.Response {
+				return &http.Response{
+					StatusCode: http.StatusSeeOther,
+					Header:     http.Header{"Location": []string{"http://example.com/new"}},
+				}
+			},
+			expectedErr: nil,
+			expectedURL: "http://example.com/new",
+		},
+	}
 
-	t.Run("Cross-Domain Security Measures", func(t *testing.T) {
-		reqURL, _ = url.Parse("http://example.com")
-		req = &http.Request{URL: reqURL, Method: http.MethodPost}
-		resp.Header.Set("Location", "http://anotherdomain.com/new")
-		err := redirectHandler.checkRedirect(req, []*http.Request{{}, {}})
-		assert.Nil(t, err)
-		// Ensure sensitive headers are removed and corresponding log is recorded
-		assert.Empty(t, req.Header.Get("Authorization"))
-		assert.Contains(t, mockLogger.Calls[2].Arguments.String(0), "Removed sensitive header")
-	})
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			resp := tc.prepare()
+			err := redirectHandler.checkRedirect(req, []*http.Request{{Response: resp}, {}})
 
-	t.Run("Handling 303 See Other", func(t *testing.T) {
-		reqURL, _ = url.Parse("http://example.com")
-		req = &http.Request{URL: reqURL, Method: http.MethodPost}
-		resp.Header.Set("Location", "http://example.com/new")
-		err := redirectHandler.checkRedirect(req, []*http.Request{{}, {}})
-		assert.Nil(t, err)
-		assert.Equal(t, http.MethodGet, req.Method)
-		// Ensure no body, no GetBody, correct ContentLength, no Content-Type header, and a log is recorded
-		assert.Nil(t, req.Body)
-		assert.Nil(t, req.GetBody)
-		assert.Equal(t, int64(0), req.ContentLength)
-		assert.Empty(t, req.Header.Get("Content-Type"))
-		assert.Contains(t, mockLogger.Calls[3].Arguments.String(0), "Changed request method to GET")
-	})
+			assert.Equal(t, tc.expectedErr, err)
+			if tc.expectedURL != "" {
+				assert.Equal(t, tc.expectedURL, req.URL.String())
+			}
+		})
+	}
 }
 
 // TestRedirectHandler_ResolveRedirectURL tests the resolveRedirectURL method of the RedirectHandler.