Fix test

Signed-off-by: 🌲 Harry 🌊 John 🏔 <[email protected]>
cortexproject · Sep 14, 2023 · 62657e5 · 62657e5
1 parent 26c024e
commit 62657e5
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/pkg/frontend/v2/frontend_test.go b/pkg/frontend/v2/frontend_test.go
@@ -27,7 +27,7 @@ import (
 
 const testFrontendWorkerConcurrency = 5
 
-func setupFrontend(t *testing.T, schedulerReplyFunc func(f *Frontend, msg *schedulerpb.FrontendToScheduler) *schedulerpb.SchedulerToFrontend) (*Frontend, *mockScheduler) {
+func setupFrontend(t *testing.T, schedulerReplyFunc func(f *Frontend, msg *schedulerpb.FrontendToScheduler) *schedulerpb.SchedulerToFrontend, maxRetries int) (*Frontend, *mockScheduler) {
 	l, err := net.Listen("tcp", "")
 	require.NoError(t, err)
 
@@ -48,7 +48,7 @@ func setupFrontend(t *testing.T, schedulerReplyFunc func(f *Frontend, msg *sched
 
 	//logger := log.NewLogfmtLogger(os.Stdout)
 	logger := log.NewNopLogger()
-	f, err := NewFrontend(cfg, logger, nil, transport.NewRetry(3, nil))
+	f, err := NewFrontend(cfg, logger, nil, transport.NewRetry(maxRetries, nil))
 	require.NoError(t, err)
 
 	frontendv2pb.RegisterFrontendForQuerierServer(server, f)
@@ -108,7 +108,7 @@ func TestFrontendBasicWorkflow(t *testing.T) {
 		})
 
 		return &schedulerpb.SchedulerToFrontend{Status: schedulerpb.OK}
-	})
+	}, 0)
 
 	resp, err := f.RoundTripGRPC(user.InjectOrgID(context.Background(), userID), &httpgrpc.HTTPRequest{})
 	require.NoError(t, err)
@@ -138,7 +138,7 @@ func TestFrontendRetryRequest(t *testing.T) {
 		}
 
 		return &schedulerpb.SchedulerToFrontend{Status: schedulerpb.OK}
-	})
+	}, 3)
 
 	res, err := f.RoundTripGRPC(user.InjectOrgID(context.Background(), userID), &httpgrpc.HTTPRequest{})
 	require.NoError(t, err)
@@ -165,7 +165,7 @@ func TestFrontendRetryEnqueue(t *testing.T) {
 		})
 
 		return &schedulerpb.SchedulerToFrontend{Status: schedulerpb.OK}
-	})
+	}, 0)
 
 	_, err := f.RoundTripGRPC(user.InjectOrgID(context.Background(), userID), &httpgrpc.HTTPRequest{})
 	require.NoError(t, err)
@@ -174,15 +174,15 @@ func TestFrontendRetryEnqueue(t *testing.T) {
 func TestFrontendEnqueueFailure(t *testing.T) {
 	f, _ := setupFrontend(t, func(f *Frontend, msg *schedulerpb.FrontendToScheduler) *schedulerpb.SchedulerToFrontend {
 		return &schedulerpb.SchedulerToFrontend{Status: schedulerpb.SHUTTING_DOWN}
-	})
+	}, 0)
 
 	_, err := f.RoundTripGRPC(user.InjectOrgID(context.Background(), "test"), &httpgrpc.HTTPRequest{})
 	require.Error(t, err)
 	require.True(t, strings.Contains(err.Error(), "failed to enqueue request"))
 }
 
 func TestFrontendCancellation(t *testing.T) {
-	f, ms := setupFrontend(t, nil)
+	f, ms := setupFrontend(t, nil, 0)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
 	defer cancel()
@@ -208,7 +208,7 @@ func TestFrontendCancellation(t *testing.T) {
 }
 
 func TestFrontendFailedCancellation(t *testing.T) {
-	f, ms := setupFrontend(t, nil)
+	f, ms := setupFrontend(t, nil, 0)
 
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()