remote-write: refactor TestClientRetryAfter

The new version features a set of test cases that simplify the addition
of new HTTP status codes.

Signed-off-by: William Dumont <william.dumont@grafana.com>

storage/remote/client_test.go

@@ -85,12 +85,21 @@ func TestStoreHTTPErrorHandling(t *testing.T) {
 }
 
 func TestClientRetryAfter(t *testing.T) {
-	server := httptest.NewServer(
-		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			http.Error(w, longErrMessage, http.StatusTooManyRequests)
-		}),
-	)
-	defer server.Close()
+	setupServer := func(statusCode int) *httptest.Server {
+		return httptest.NewServer(
+			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				http.Error(w, longErrMessage, statusCode)
+			}),
+		)
+	}
+
+	getClientConfig := func(serverURL *url.URL, retryOnRateLimit bool) *ClientConfig {
+		return &ClientConfig{
+			URL:              &config_util.URL{URL: serverURL},
+			Timeout:          model.Duration(time.Second),
+			RetryOnRateLimit: retryOnRateLimit,
+		}
+	}
 
 	getClient := func(conf *ClientConfig) WriteClient {
 		hash, err := toHash(conf)
@@ -100,30 +109,31 @@ func TestClientRetryAfter(t *testing.T) {
 		return c
 	}
 
-	serverURL, err := url.Parse(server.URL)
-	require.NoError(t, err)
-
-	conf := &ClientConfig{
-		URL:              &config_util.URL{URL: serverURL},
-		Timeout:          model.Duration(time.Second),
-		RetryOnRateLimit: false,
+	testCases := []struct {
+		name                string
+		statusCode          int
+		retryOnRateLimit    bool
+		expectedRecoverable bool
+	}{
+		{"TooManyRequests - No Retry", http.StatusTooManyRequests, false, false},
+		{"TooManyRequests - With Retry", http.StatusTooManyRequests, true, true},
 	}
 
-	var recErr RecoverableError
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			server := setupServer(tc.statusCode)
+			defer server.Close()
 
-	c := getClient(conf)
-	err = c.Store(context.Background(), []byte{}, 0)
-	require.False(t, errors.As(err, &recErr), "Recoverable error not expected.")
+			serverURL, err := url.Parse(server.URL)
+			require.NoError(t, err)
 
-	conf = &ClientConfig{
-		URL:              &config_util.URL{URL: serverURL},
-		Timeout:          model.Duration(time.Second),
-		RetryOnRateLimit: true,
-	}
+			c := getClient(getClientConfig(serverURL, tc.retryOnRateLimit))
 
-	c = getClient(conf)
-	err = c.Store(context.Background(), []byte{}, 0)
-	require.True(t, errors.As(err, &recErr), "Recoverable error was expected.")
+			var recErr RecoverableError
+			err = c.Store(context.Background(), []byte{}, 0)
+			require.Equal(t, tc.expectedRecoverable, errors.As(err, &recErr), "Mismatch in expected recoverable error status.")
+		})
+	}
 }
 
 func TestRetryAfterDuration(t *testing.T) {