Skip to content

Commit

Permalink
Improve tests
Browse files Browse the repository at this point in the history 
Signed-off-by: Justin Jung <[email protected]>
  • Loading branch information
@justinjung04
justinjung04 committed Nov 30, 2023
1 parent 3ee0b0d commit c0252f1
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 67 deletions.
2 changes: 1 addition & 1 deletion pkg/querier/tripperware/queryrange/limits_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package queryrange

import (
"context"
"github.com/cortexproject/cortex/pkg/util/validation"
"testing"
"time"

Expand All @@ -13,6 +12,7 @@ import (

"github.com/cortexproject/cortex/pkg/querier/tripperware"
"github.com/cortexproject/cortex/pkg/util"
"github.com/cortexproject/cortex/pkg/util/validation"
)

func TestLimitsMiddleware_MaxQueryLookback(t *testing.T) {
Expand Down
17 changes: 14 additions & 3 deletions pkg/scheduler/queue/user_queues.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package queue

import (
"math/rand"
"reflect"
"sort"
"time"

Expand Down Expand Up @@ -136,7 +135,7 @@ func (q *queues) getOrAddQueue(userID string, maxQueriers int) userRequestQueue
uq := q.userQueues[userID]
priorityEnabled := q.limits.QueryPriority(userID).Enabled
maxOutstanding := q.limits.MaxOutstandingPerTenant(userID)
priorityList := getPriorityList(q.limits.QueryPriority(userID), len(q.queriers))
priorityList := getPriorityList(q.limits.QueryPriority(userID), maxQueriers)

if uq == nil {
uq = &userQueue{
Expand Down Expand Up @@ -180,7 +179,7 @@ func (q *queues) getOrAddQueue(userID string, maxQueriers int) userRequestQueue
uq.queriers = shuffleQueriersForUser(uq.seed, maxQueriers, q.sortedQueriers, nil)
}

if priorityEnabled && !reflect.DeepEqual(uq.priorityList, priorityList) {
if priorityEnabled && hasPriorityListChanged(uq.priorityList, priorityList) {
reservedQueriers := make(map[string]int64)

i := 0
Expand Down Expand Up @@ -406,6 +405,18 @@ func getPriorityList(queryPriority validation.QueryPriority, totalQuerierCount i
return priorityList
}

func hasPriorityListChanged(old, new []int64) bool {
if len(old) != len(new) {
return true
}
for i := range old {
if old[i] != new[i] {
return true
}
}
return false
}

// MockLimits implements the Limits interface. Used in tests only.
type MockLimits struct {
MaxOutstanding int
Expand Down
101 changes: 38 additions & 63 deletions pkg/scheduler/queue/user_queues_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -399,13 +399,46 @@ func TestGetOrAddQueueShouldUpdateProperties(t *testing.T) {
assert.Equal(t, 3, q.userQueues["userID"].maxQueriers)
assert.Equal(t, 5, len(q.queriers))
assert.Equal(t, 3, len(q.userQueues["userID"].queriers))
assert.Equal(t, 2, len(q.userQueues["userID"].reservedQueriers))
assert.IsType(t, &PriorityRequestQueue{}, queue)
assert.Equal(t, 1, queue.length())
assert.ElementsMatch(t, []int64{1, 1}, q.userQueues["userID"].priorityList)
assert.True(t, q.userQueues["userID"].priorityEnabled)
assert.Subset(t, getKeys(q.queriers), getKeys(q.userQueues["userID"].queriers))
assert.Subset(t, getKeys(q.userQueues["userID"].queriers), getKeys(q.userQueues["userID"].reservedQueriers))

limits.QueryPriorityVal = validation.QueryPriority{Enabled: true, Priorities: []validation.PriorityDef{
{
Priority: 1,
ReservedQueriers: 0.5,
},
}}
q.limits = limits
_ = q.getOrAddQueue("userID", 3)

assert.Equal(t, 10, q.userQueues["userID"].maxOutstanding)
assert.Equal(t, 3, q.userQueues["userID"].maxQueriers)
assert.Equal(t, 5, len(q.queriers))
assert.Equal(t, 3, len(q.userQueues["userID"].queriers))
assert.Equal(t, 2, len(q.userQueues["userID"].reservedQueriers))
assert.ElementsMatch(t, []int64{1, 1}, q.userQueues["userID"].priorityList)

limits.QueryPriorityVal = validation.QueryPriority{Enabled: true, Priorities: []validation.PriorityDef{
{
Priority: 1,
ReservedQueriers: 10,
},
}}
q.limits = limits
_ = q.getOrAddQueue("userID", 3)

assert.Equal(t, 10, q.userQueues["userID"].maxOutstanding)
assert.Equal(t, 3, q.userQueues["userID"].maxQueriers)
assert.Equal(t, 5, len(q.queriers))
assert.Equal(t, 3, len(q.userQueues["userID"].queriers))
assert.Equal(t, 0, len(q.userQueues["userID"].reservedQueriers))
assert.ElementsMatch(t, []int64{}, q.userQueues["userID"].priorityList)

limits.QueryPriorityVal.Enabled = false
q.limits = limits
queue = q.getOrAddQueue("userID", 3)
Expand Down Expand Up @@ -582,69 +615,11 @@ func TestShuffleQueriersCorrectness(t *testing.T) {
}
}

func TestShuffleQueriers_WithReservedQueriers(t *testing.T) {
//allQueriers := []string{"a", "b", "c", "d", "e"}
//
//queriers, reservedQueriers := shuffleQueriersForUser(12345, 0, allQueriers, 0, nil)
//require.Nil(t, queriers)
//require.Equal(t, 0, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 0, allQueriers, 0.5, nil)
//require.Nil(t, queriers)
//require.Equal(t, 3, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 0, allQueriers, 1, nil)
//require.Nil(t, queriers)
//require.Equal(t, 1, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 0, allQueriers, 100, nil)
//require.Nil(t, queriers)
//require.Equal(t, 5, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 3, allQueriers, 0, nil)
//require.Equal(t, 3, len(queriers))
//require.Equal(t, 0, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 3, allQueriers, 0.5, nil)
//require.Equal(t, 3, len(queriers))
//require.Equal(t, 2, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 3, allQueriers, 1, nil)
//require.Equal(t, 3, len(queriers))
//require.Equal(t, 1, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 3, allQueriers, 100, nil)
//require.Equal(t, 3, len(queriers))
//require.Equal(t, 3, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 100, allQueriers, 0, nil)
//require.Nil(t, queriers)
//require.Equal(t, 0, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 100, allQueriers, 0.5, nil)
//require.Nil(t, queriers)
//require.Equal(t, 3, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 100, allQueriers, 1, nil)
//require.Nil(t, queriers)
//require.Equal(t, 1, len(reservedQueriers))
//
//queriers, reservedQueriers = shuffleQueriersForUser(12345, 100, allQueriers, 100, nil)
//require.Nil(t, queriers)
//require.Equal(t, 5, len(reservedQueriers))
}

func TestShuffleQueriers_WithReservedQueriers_Correctness(t *testing.T) {
//allQueriers := []string{"a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n"}
//
//prevQueriers, prevReservedQueriers := shuffleQueriersForUser(12345, 10, allQueriers, 5, nil)
//for i := 0; i < 100; i++ {
// queriers, reservedQueriers := shuffleQueriersForUser(12345, 10, allQueriers, 5, nil)
// require.Equal(t, prevQueriers, queriers)
// require.Equal(t, prevReservedQueriers, reservedQueriers)
// prevQueriers = queriers
// prevReservedQueriers = reservedQueriers
//}
func TestHasPriorityListChanged(t *testing.T) {
require.True(t, hasPriorityListChanged([]int64{1, 2}, []int64{1, 3}))
require.False(t, hasPriorityListChanged([]int64{1, 2}, []int64{1, 2}))
require.True(t, hasPriorityListChanged([]int64{1, 2}, []int64{1}))
require.False(t, hasPriorityListChanged([]int64{}, []int64{}))
}

func TestGetPriorityList(t *testing.T) {
Expand Down

0 comments on commit c0252f1

Please sign in to comment.