if no codes set for repeater, treat any >400 as error

go-pkgz · Nov 5, 2023 · 69afd00 · 69afd00
1 parent 5921c97
commit 69afd00
Show file tree

Hide file tree

Showing 8 changed files with 137 additions and 55 deletions.
diff --git a/middleware/middleware.go b/middleware/middleware.go
@@ -5,10 +5,10 @@ import (
 	"net/http"
 )
 
-//go:generate moq -out mocks/repeater.go -pkg mocks -skip-ensure -fmt goimports . RepeaterSvc
-//go:generate moq -out mocks/circuit_breaker.go -pkg mocks -skip-ensure -fmt goimports . CircuitBreakerSvc
-//go:generate moq -out mocks/logger.go -pkg mocks -skip-ensure -fmt goimports logger Service:LoggerSvc
-//go:generate moq -out mocks/cache.go -pkg mocks -skip-ensure -fmt goimports cache Service:CacheSvc
+//go:generate moq -out mocks/repeater.go -pkg mocks -skip-ensure -with-resets -fmt goimports . RepeaterSvc
+//go:generate moq -out mocks/circuit_breaker.go -pkg mocks -skip-ensure -with-resets -fmt goimports . CircuitBreakerSvc
+//go:generate moq -out mocks/logger.go -pkg mocks -skip-ensure -with-resets -fmt goimports logger Service:LoggerSvc
+//go:generate moq -out mocks/cache.go -pkg mocks -skip-ensure -with-resets -fmt goimports cache Service:CacheSvc
 
 // RoundTripperHandler is a type for middleware handler
 type RoundTripperHandler func(http.RoundTripper) http.RoundTripper

diff --git a/middleware/mocks/cache.go b/middleware/mocks/cache.go
diff --git a/middleware/mocks/circuit_breaker.go b/middleware/mocks/circuit_breaker.go
diff --git a/middleware/mocks/logger.go b/middleware/mocks/logger.go
diff --git a/middleware/mocks/repeater.go b/middleware/mocks/repeater.go
diff --git a/middleware/mocks/roundtripper.go b/middleware/mocks/roundtripper.go
@@ -21,3 +21,8 @@ func (r *RoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
 func (r *RoundTripper) Calls() int {
 	return int(atomic.LoadInt32(&r.calls))
 }
+
+// ResetCalls resets calls counter
+func (r *RoundTripper) ResetCalls() {
+	atomic.StoreInt32(&r.calls, 0)
+}
diff --git a/middleware/repeater.go b/middleware/repeater.go
@@ -29,6 +29,11 @@ func Repeater(repeater RepeaterSvc, failOnCodes ...int) RoundTripperHandler {
 				if err != nil {
 					return err
 				}
+				// no explicit codes provided, fail on any 4xx or 5xx
+				if len(failOnCodes) == 0 && resp.StatusCode >= 400 {
+					return errors.New(resp.Status)
+				}
+				// fail on provided codes only
 				for _, fc := range failOnCodes {
 					if resp.StatusCode == fc {
 						return errors.New(resp.Status)

diff --git a/middleware/repeater_test.go b/middleware/repeater_test.go
@@ -89,19 +89,30 @@ func TestRepeater_FailedStatus(t *testing.T) {
 		return err
 	}}
 
-	{
-		h := Repeater(repeater, 300, 400, 401)
-
+	t.Run("no codes", func(t *testing.T) {
+		rmock.ResetCalls()
+		h := Repeater(repeater)
 		req, err := http.NewRequest("GET", "http://example.com/blah", http.NoBody)
 		require.NoError(t, err)
 
 		_, err = h(rmock).RoundTrip(req)
 		require.EqualError(t, err, "repeater: 400 Bad Request")
-	}
+		assert.Equal(t, 5, rmock.Calls())
+	})
 
-	assert.Equal(t, 5, rmock.Calls())
+	t.Run("with codes", func(t *testing.T) {
+		rmock.ResetCalls()
+		h := Repeater(repeater, 300, 400, 401)
+		req, err := http.NewRequest("GET", "http://example.com/blah", http.NoBody)
+		require.NoError(t, err)
+
+		_, err = h(rmock).RoundTrip(req)
+		require.EqualError(t, err, "repeater: 400 Bad Request")
+		assert.Equal(t, 5, rmock.Calls())
+	})
 
-	{
+	t.Run("no codes, no match", func(t *testing.T) {
+		rmock.ResetCalls()
 		h := Repeater(repeater, 300, 401)
 
 		req, err := http.NewRequest("GET", "http://example.com/blah", http.NoBody)
@@ -110,6 +121,7 @@ func TestRepeater_FailedStatus(t *testing.T) {
 		resp, err := h(rmock).RoundTrip(req)
 		require.NoError(t, err)
 		assert.Equal(t, 400, resp.StatusCode)
-	}
+		assert.Equal(t, 1, rmock.Calls())
+	})
 
 }