@@ -17,6 +17,8 @@ limitations under the License.
17
17
package quota_test
18
18
19
19
import (
20
+ "reflect"
21
+ "sort"
20
22
"strings"
21
23
"testing"
22
24
"time"
@@ -381,6 +383,211 @@ func TestQuotaManagerQuotaUsedLongRunningConsumers(t *testing.T) {
381
383
382
384
}
383
385
386
+ var (
387
+ tree1String string = `{
388
+ "kind": "QuotaTree",
389
+ "metadata": {
390
+ "name": "tree-1"
391
+ },
392
+ "spec": {
393
+ "resourceNames": [
394
+ "cpu",
395
+ "memory"
396
+ ],
397
+ "nodes": {
398
+ "A": {
399
+ "parent": "nil",
400
+ "hard": "true",
401
+ "quota": {
402
+ "cpu": "10",
403
+ "memory": "256"
404
+ }
405
+ },
406
+ "B": {
407
+ "parent": "A",
408
+ "hard": "true",
409
+ "quota": {
410
+ "cpu": "2",
411
+ "memory": "64"
412
+ }
413
+ },
414
+ "C": {
415
+ "parent": "A",
416
+ "quota": {
417
+ "cpu": "6",
418
+ "memory": "64"
419
+ }
420
+ }
421
+ }
422
+ }
423
+ }`
424
+
425
+ tree2String string = `{
426
+ "kind": "QuotaTree",
427
+ "metadata": {
428
+ "name": "tree-2"
429
+ },
430
+ "spec": {
431
+ "resourceNames": [
432
+ "gpu"
433
+ ],
434
+ "nodes": {
435
+ "X": {
436
+ "parent": "nil",
437
+ "hard": "true",
438
+ "quota": {
439
+ "gpu": "32"
440
+ }
441
+ },
442
+ "Y": {
443
+ "parent": "X",
444
+ "hard": "true",
445
+ "quota": {
446
+ "gpu": "16"
447
+ }
448
+ },
449
+ "Z": {
450
+ "parent": "X",
451
+ "quota": {
452
+ "gpu": "8"
453
+ }
454
+ }
455
+ }
456
+ }
457
+ }`
458
+
459
+ tree3String string = `{
460
+ "kind": "QuotaTree",
461
+ "metadata": {
462
+ "name": "tree-3"
463
+ },
464
+ "spec": {
465
+ "resourceNames": [
466
+ "count"
467
+ ],
468
+ "nodes": {
469
+ "zone": {
470
+ "parent": "nil",
471
+ "hard": "true",
472
+ "quota": {
473
+ "count": "100"
474
+ }
475
+ },
476
+ "rack": {
477
+ "parent": "zone",
478
+ "hard": "true",
479
+ "quota": {
480
+ "count": "100"
481
+ }
482
+ },
483
+ "server": {
484
+ "parent": "rack",
485
+ "quota": {
486
+ "count": "100"
487
+ }
488
+ }
489
+ }
490
+ }
491
+ }`
492
+ )
493
+
494
+ // TestAddDeleteTrees : test adding, retrieving, deleting trees
495
+ func TestAddDeleteTrees (t * testing.T ) {
496
+ qmTest := quota .NewManager ()
497
+ assert .NotNil (t , qmTest , "Expecting no error creating a quota manager" )
498
+ modeSet := qmTest .SetMode (quota .Normal )
499
+ assert .True (t , modeSet , "Expecting no error setting mode to normal" )
500
+ mode := qmTest .GetMode ()
501
+ assert .True (t , mode == quota .Normal , "Expecting mode set to normal" )
502
+ modeString := qmTest .GetModeString ()
503
+ match := strings .Contains (modeString , "Normal" )
504
+ assert .True (t , match , "Expecting mode set to normal" )
505
+ // add a few trees by name
506
+ treeNames := []string {"tree-1" , "tree-2" , "tree-3" }
507
+ treeStrings := []string {tree1String , tree2String , tree3String }
508
+ for i , treeString := range treeStrings {
509
+ name , err := qmTest .AddTreeFromString (treeString )
510
+ assert .NoError (t , err , "No error expected when adding a tree" )
511
+ assert .Equal (t , treeNames [i ], name , "Returned name should match" )
512
+ }
513
+ // get list of names
514
+ names := qmTest .GetTreeNames ()
515
+ assert .ElementsMatch (t , treeNames , names , "Expecting retrieved tree names same as added names" )
516
+ // delete a name
517
+ deletedTreeName := treeNames [0 ]
518
+ remainingTreeNames := treeNames [1 :]
519
+ err := qmTest .DeleteTree (deletedTreeName )
520
+ assert .NoError (t , err , "No error expected when deleting an existing tree" )
521
+ // get list of names after deletion
522
+ names = qmTest .GetTreeNames ()
523
+ assert .ElementsMatch (t , remainingTreeNames , names , "Expecting retrieved tree names to reflect additions/deletions" )
524
+ // delete a non-existing name
525
+ err = qmTest .DeleteTree (deletedTreeName )
526
+ assert .Error (t , err , "Error expected when deleting a non-existing tree" )
527
+ }
528
+
529
+ // TestAddDeleteForests : test adding, retrieving, deleting forests
530
+ func TestAddDeleteForests (t * testing.T ) {
531
+ var err error
532
+ qmTest := quota .NewManager ()
533
+ assert .NotNil (t , qmTest , "Expecting no error creating a quota manager" )
534
+ modeSet := qmTest .SetMode (quota .Normal )
535
+ assert .True (t , modeSet , "Expecting no error setting mode to normal" )
536
+
537
+ // add a few trees by name
538
+ treeNames := []string {"tree-1" , "tree-2" , "tree-3" }
539
+ treeStrings := []string {tree1String , tree2String , tree3String }
540
+ for i , treeString := range treeStrings {
541
+ name , err := qmTest .AddTreeFromString (treeString )
542
+ assert .NoError (t , err , "No error expected when adding a tree" )
543
+ assert .Equal (t , treeNames [i ], name , "Returned name should match" )
544
+ }
545
+ // create two forests
546
+ forestNames := []string {"forest-1" , "forest-2" }
547
+ for _ , forestName := range forestNames {
548
+ err = qmTest .AddForest (forestName )
549
+ assert .NoError (t , err , "No error expected when adding a forest" )
550
+ }
551
+ // assign trees to forests
552
+ err = qmTest .AddTreeToForest ("forest-1" , "tree-1" )
553
+ assert .NoError (t , err , "No error expected when adding a tree to a forest" )
554
+ err = qmTest .AddTreeToForest ("forest-2" , "tree-2" )
555
+ assert .NoError (t , err , "No error expected when adding a tree to a forest" )
556
+ err = qmTest .AddTreeToForest ("forest-2" , "tree-3" )
557
+ assert .NoError (t , err , "No error expected when adding a tree to a forest" )
558
+ // get list of forest names
559
+ fNames := qmTest .GetForestNames ()
560
+ assert .ElementsMatch (t , forestNames , fNames , "Expecting retrieved forest names same as added names" )
561
+ // get forests map
562
+ forestTreeMap := qmTest .GetForestTreeNames ()
563
+ for _ , v := range forestTreeMap {
564
+ sort .Strings (v )
565
+ }
566
+ inputForestTreeMap := map [string ][]string {"forest-1" : {"tree-1" }, "forest-2" : {"tree-2" , "tree-3" }}
567
+ assert .True (t , reflect .DeepEqual (forestTreeMap , inputForestTreeMap ),
568
+ "Expecting retrieved forest tree map same as input, got %v, want %v" , forestTreeMap , inputForestTreeMap )
569
+ // delete a forest
570
+ deletedForestName := forestNames [0 ]
571
+ remainingForestNames := forestNames [1 :]
572
+ err = qmTest .DeleteForest (deletedForestName )
573
+ assert .NoError (t , err , "No error expected when deleting an existing forest" )
574
+ // get list of forest names after deletion
575
+ fNames = qmTest .GetForestNames ()
576
+ assert .ElementsMatch (t , remainingForestNames , fNames , "Expecting retrieved forest names to reflect additions/deletions" )
577
+ // delete a non-existing forest name
578
+ err = qmTest .DeleteForest (deletedForestName )
579
+ assert .Error (t , err , "Error expected when deleting a non-existing forest" )
580
+ // delete a tree from a forest
581
+ err = qmTest .DeleteTreeFromForest ("forest-2" , "tree-2" )
582
+ assert .NoError (t , err , "No error expected when deleting an existing tree from an existing forest" )
583
+ err = qmTest .DeleteTreeFromForest ("forest-2" , "tree-2" )
584
+ assert .Error (t , err , "Error expected when deleting an non-existing tree from an existing forest" )
585
+ // check remaining trees after deletions
586
+ names := qmTest .GetTreeNames ()
587
+ assert .True (t , reflect .DeepEqual (treeNames , names ),
588
+ "Expecting all trees after forest deletions as trees are not deleted, got %v, want %v" , names , treeNames )
589
+ }
590
+
384
591
type AllocationClassifier struct {
385
592
}
386
593
0 commit comments