@@ -723,3 +723,127 @@ func TestCountNodeTaints(t *testing.T) {
723723 got := CountNodeTaints ([]* apiv1.Node {node , node2 }, taintConfig )
724724 assert .Equal (t , want , got )
725725}
726+
727+ func TestAddTaints (t * testing.T ) {
728+ testCases := []struct {
729+ name string
730+ existingTaints []string
731+ newTaints []string
732+ wantTaints []string
733+ }{
734+ {
735+ name : "no existing taints" ,
736+ newTaints : []string {"t1" , "t2" },
737+ wantTaints : []string {"t1" , "t2" },
738+ },
739+ {
740+ name : "existing taints - no overlap" ,
741+ existingTaints : []string {"t1" },
742+ newTaints : []string {"t2" , "t3" },
743+ wantTaints : []string {"t1" , "t2" , "t3" },
744+ },
745+ {
746+ name : "existing taints - duplicates" ,
747+ existingTaints : []string {"t1" },
748+ newTaints : []string {"t1" , "t2" },
749+ wantTaints : []string {"t1" , "t2" },
750+ },
751+ }
752+ for _ , tc := range testCases {
753+ t .Run (tc .name , func (t * testing.T ) {
754+ n := BuildTestNode ("node" , 1000 , 1000 )
755+ existingTaints := make ([]apiv1.Taint , len (tc .existingTaints ))
756+ for i , t := range tc .existingTaints {
757+ existingTaints [i ] = apiv1.Taint {
758+ Key : t ,
759+ Effect : apiv1 .TaintEffectNoSchedule ,
760+ }
761+ }
762+ n .Spec .Taints = append ([]apiv1.Taint {}, existingTaints ... )
763+ fakeClient := buildFakeClient (t , n )
764+ newTaints := make ([]apiv1.Taint , len (tc .newTaints ))
765+ for i , t := range tc .newTaints {
766+ newTaints [i ] = apiv1.Taint {
767+ Key : t ,
768+ Effect : apiv1 .TaintEffectNoSchedule ,
769+ }
770+ }
771+ err := AddTaints (n , fakeClient , newTaints , false )
772+ assert .NoError (t , err )
773+ apiNode := getNode (t , fakeClient , "node" )
774+ for _ , want := range tc .wantTaints {
775+ assert .True (t , HasTaint (n , want ))
776+ assert .True (t , HasTaint (apiNode , want ))
777+ }
778+ })
779+ }
780+ }
781+
782+ func TestCleanTaints (t * testing.T ) {
783+ testCases := []struct {
784+ name string
785+ existingTaints []string
786+ taintsToRemove []string
787+ wantTaints []string
788+ wantModified bool
789+ }{
790+ {
791+ name : "no existing taints" ,
792+ taintsToRemove : []string {"t1" , "t2" },
793+ wantTaints : []string {},
794+ wantModified : false ,
795+ },
796+ {
797+ name : "existing taints - no overlap" ,
798+ existingTaints : []string {"t1" },
799+ taintsToRemove : []string {"t2" , "t3" },
800+ wantTaints : []string {"t1" },
801+ wantModified : false ,
802+ },
803+ {
804+ name : "existing taints - remove one" ,
805+ existingTaints : []string {"t1" , "t2" },
806+ taintsToRemove : []string {"t1" },
807+ wantTaints : []string {"t2" },
808+ wantModified : true ,
809+ },
810+ {
811+ name : "existing taints - remove all" ,
812+ existingTaints : []string {"t1" , "t2" },
813+ taintsToRemove : []string {"t1" , "t2" },
814+ wantTaints : []string {},
815+ wantModified : true ,
816+ },
817+ }
818+
819+ for _ , tc := range testCases {
820+ t .Run (tc .name , func (t * testing.T ) {
821+ n := BuildTestNode ("node" , 1000 , 1000 )
822+ existingTaints := make ([]apiv1.Taint , len (tc .existingTaints ))
823+ for i , taintKey := range tc .existingTaints {
824+ existingTaints [i ] = apiv1.Taint {
825+ Key : taintKey ,
826+ Effect : apiv1 .TaintEffectNoSchedule ,
827+ }
828+ }
829+ n .Spec .Taints = append ([]apiv1.Taint {}, existingTaints ... )
830+ fakeClient := buildFakeClient (t , n )
831+
832+ modified , err := CleanTaints (n , fakeClient , tc .taintsToRemove , false )
833+ assert .NoError (t , err )
834+ assert .Equal (t , tc .wantModified , modified )
835+
836+ apiNode := getNode (t , fakeClient , "node" )
837+
838+ for _ , want := range tc .wantTaints {
839+ assert .True (t , HasTaint (apiNode , want ))
840+ assert .True (t , HasTaint (n , want ))
841+ }
842+
843+ for _ , removed := range tc .taintsToRemove {
844+ assert .False (t , HasTaint (apiNode , removed ))
845+ assert .False (t , HasTaint (n , removed ), "Taint %s should have been removed from local node object" , removed )
846+ }
847+ })
848+ }
849+ }
0 commit comments