@@ -18,6 +18,7 @@ import (
18
18
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator"
19
19
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/allocator/storepool"
20
20
"github.com/cockroachdb/cockroach/pkg/kv/kvserver/constraint"
21
+ "github.com/cockroachdb/cockroach/pkg/kv/kvserver/mmaintegration"
21
22
"github.com/cockroachdb/cockroach/pkg/roachpb"
22
23
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
23
24
"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -436,10 +437,18 @@ func TestBestRebalanceTarget(t *testing.T) {
436
437
expectedTargets := []roachpb.StoreID {13 , 13 , 11 , 12 }
437
438
expectedExistingRepls := []roachpb.StoreID {3 , 2 , 1 , 1 }
438
439
allocRand := makeAllocatorRand (rand .NewSource (0 ))
440
+ ctx := context .Background ()
441
+ stopper , _ , _ , a , _ := CreateTestAllocatorWithKnobs (ctx , 10 , false , /* deterministic */
442
+ nil /* allocator.TestingKnobs */ , & mmaintegration.TestingKnobs {
443
+ OverrideIsInConflictWithMMA : func (cand roachpb.StoreID ) bool {
444
+ return false
445
+ },
446
+ } /* mmaintegration.TestingKnobs */ )
447
+ defer stopper .Stop (ctx )
439
448
var i int
440
449
for {
441
450
i ++
442
- target , existing , _ := bestRebalanceTarget (allocRand , candidates , nil , nil )
451
+ target , existing , _ := bestRebalanceTarget (allocRand , candidates , a . as )
443
452
if len (expectedTargets ) == 0 {
444
453
if target == nil {
445
454
break
0 commit comments