From e58b2886d2d551d5ef9bc00c6927f5b7cacdb1ab Mon Sep 17 00:00:00 2001
From: Justin Jung <jungjust@amazon.com>
Date: Wed, 10 Jul 2024 10:50:15 -0700
Subject: [PATCH] Fix test

Signed-off-by: Justin Jung <jungjust@amazon.com>
---
 pkg/scheduler/queue/user_queues.go      | 4 ++++
 pkg/scheduler/queue/user_queues_test.go | 8 ++++----
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/pkg/scheduler/queue/user_queues.go b/pkg/scheduler/queue/user_queues.go
index 159df7810bf..d5c6d9302b0 100644
--- a/pkg/scheduler/queue/user_queues.go
+++ b/pkg/scheduler/queue/user_queues.go
@@ -433,6 +433,7 @@ type MockLimits struct {
 	MaxOutstanding        int
 	MaxQueriersPerUserVal float64
 	QueryPriorityVal      validation.QueryPriority
+	RandomMaxOutstanding  bool
 }
 
 func (l MockLimits) MaxQueriersPerUser(_ string) float64 {
@@ -440,6 +441,9 @@ func (l MockLimits) MaxQueriersPerUser(_ string) float64 {
 }
 
 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
index 192cd17b523..ab9aa2f71ac 100644
--- 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{})