Fix test

Signed-off-by: Justin Jung <jungjust@amazon.com>
cortexproject · Jul 10, 2024 · e58b288 · e58b288
1 parent 68e47e8
commit e58b288
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/pkg/scheduler/queue/user_queues.go b/pkg/scheduler/queue/user_queues.go
@@ -433,13 +433,17 @@ type MockLimits struct {
 	MaxOutstanding        int
 	MaxQueriersPerUserVal float64
 	QueryPriorityVal      validation.QueryPriority
+	RandomMaxOutstanding  bool
 }
 
 func (l MockLimits) MaxQueriersPerUser(_ string) float64 {
 	return l.MaxQueriersPerUserVal
 }
 
 func (l MockLimits) MaxOutstandingPerTenant(_ string) int {
+	if l.RandomMaxOutstanding {
+		return rand.Intn(100)
+	}
 	return l.MaxOutstanding
 }
 

diff --git a/pkg/scheduler/queue/user_queues_test.go b/pkg/scheduler/queue/user_queues_test.go
@@ -460,7 +460,10 @@ func TestGetOrAddQueueShouldUpdateProperties(t *testing.T) {
 
 func TestGetOrAddQueueConcurrency(t *testing.T) {
 	const numGoRoutines = 100
-	q := newUserQueues(0, 0, MockLimits{}, nil)
+	limits := MockLimits{
+		RandomMaxOutstanding: true,
+	}
+	q := newUserQueues(0, 0, limits, nil)
 	q.addQuerierConnection("q-1")
 	q.addQuerierConnection("q-2")
 	q.addQuerierConnection("q-3")
@@ -473,9 +476,6 @@ func TestGetOrAddQueueConcurrency(t *testing.T) {
 	for i := 0; i < numGoRoutines; i++ {
 		go func(cnt int) {
 			defer wg.Done()
-			q.limits = MockLimits{
-				MaxOutstanding: cnt + 50,
-			}
 			queue := q.getOrAddQueue("userID", 2)
 			if cnt%2 == 0 {
 				queue.enqueueRequest(MockRequest{})