@@ -944,3 +944,210 @@ func TestUpdateKusciaDeploymentParty(t *testing.T) {
944944 })
945945 }
946946}
947+
948+ func TestGetAffinityMode (t * testing.T ) {
949+ s := servingService {}
950+ tests := []struct {
951+ name string
952+ mode string
953+ expected string
954+ }{
955+ {
956+ name : "empty mode should default to anti-affinity" ,
957+ mode : "" ,
958+ expected : "anti-affinity" ,
959+ },
960+ {
961+ name : "none mode" ,
962+ mode : "none" ,
963+ expected : "none" ,
964+ },
965+ {
966+ name : "affinity mode" ,
967+ mode : "affinity" ,
968+ expected : "affinity" ,
969+ },
970+ {
971+ name : "anti-affinity mode" ,
972+ mode : "anti-affinity" ,
973+ expected : "anti-affinity" ,
974+ },
975+ }
976+
977+ for _ , tt := range tests {
978+ t .Run (tt .name , func (t * testing.T ) {
979+ got := s .getAffinityMode (tt .mode )
980+ assert .Equal (t , tt .expected , got )
981+ })
982+ }
983+ }
984+
985+ func TestBuildAffinityForMode (t * testing.T ) {
986+ s := servingService {}
987+ servingID := "test-serving"
988+
989+ tests := []struct {
990+ name string
991+ mode string
992+ expected * corev1.Affinity
993+ }{
994+ {
995+ name : "none mode should return nil" ,
996+ mode : "none" ,
997+ expected : nil ,
998+ },
999+ {
1000+ name : "affinity mode should return PodAffinity" ,
1001+ mode : "affinity" ,
1002+ expected : & corev1.Affinity {
1003+ PodAffinity : & corev1.PodAffinity {
1004+ PreferredDuringSchedulingIgnoredDuringExecution : []corev1.WeightedPodAffinityTerm {
1005+ {
1006+ Weight : 100 ,
1007+ PodAffinityTerm : corev1.PodAffinityTerm {
1008+ LabelSelector : & metav1.LabelSelector {
1009+ MatchLabels : map [string ]string {
1010+ "kuscia.secretflow/kd-name" : servingID ,
1011+ },
1012+ },
1013+ TopologyKey : "kubernetes.io/hostname" ,
1014+ },
1015+ },
1016+ },
1017+ },
1018+ },
1019+ },
1020+ {
1021+ name : "anti-affinity mode should return PodAntiAffinity" ,
1022+ mode : "anti-affinity" ,
1023+ expected : & corev1.Affinity {
1024+ PodAntiAffinity : & corev1.PodAntiAffinity {
1025+ PreferredDuringSchedulingIgnoredDuringExecution : []corev1.WeightedPodAffinityTerm {
1026+ {
1027+ Weight : 100 ,
1028+ PodAffinityTerm : corev1.PodAffinityTerm {
1029+ LabelSelector : & metav1.LabelSelector {
1030+ MatchLabels : map [string ]string {
1031+ "kuscia.secretflow/kd-name" : servingID ,
1032+ },
1033+ },
1034+ TopologyKey : "kubernetes.io/hostname" ,
1035+ },
1036+ },
1037+ },
1038+ },
1039+ },
1040+ },
1041+ {
1042+ name : "unknown mode should default to anti-affinity" ,
1043+ mode : "unknown" ,
1044+ expected : & corev1.Affinity {
1045+ PodAntiAffinity : & corev1.PodAntiAffinity {
1046+ PreferredDuringSchedulingIgnoredDuringExecution : []corev1.WeightedPodAffinityTerm {
1047+ {
1048+ Weight : 100 ,
1049+ PodAffinityTerm : corev1.PodAffinityTerm {
1050+ LabelSelector : & metav1.LabelSelector {
1051+ MatchLabels : map [string ]string {
1052+ "kuscia.secretflow/kd-name" : servingID ,
1053+ },
1054+ },
1055+ TopologyKey : "kubernetes.io/hostname" ,
1056+ },
1057+ },
1058+ },
1059+ },
1060+ },
1061+ },
1062+ }
1063+
1064+ for _ , tt := range tests {
1065+ t .Run (tt .name , func (t * testing.T ) {
1066+ got := s .buildAffinityForMode (tt .mode , servingID )
1067+ if tt .expected == nil {
1068+ assert .Nil (t , got )
1069+ } else {
1070+ assert .NotNil (t , got )
1071+ if tt .mode == "affinity" {
1072+ assert .NotNil (t , got .PodAffinity )
1073+ assert .Nil (t , got .PodAntiAffinity )
1074+ assert .Equal (t , tt .expected .PodAffinity .PreferredDuringSchedulingIgnoredDuringExecution [0 ].Weight , got .PodAffinity .PreferredDuringSchedulingIgnoredDuringExecution [0 ].Weight )
1075+ } else {
1076+ assert .NotNil (t , got .PodAntiAffinity )
1077+ assert .Nil (t , got .PodAffinity )
1078+ assert .Equal (t , tt .expected .PodAntiAffinity .PreferredDuringSchedulingIgnoredDuringExecution [0 ].Weight , got .PodAntiAffinity .PreferredDuringSchedulingIgnoredDuringExecution [0 ].Weight )
1079+ }
1080+ }
1081+ })
1082+ }
1083+ }
1084+
1085+ func TestBuildKusciaDeploymentWithAffinityMode (t * testing.T ) {
1086+ replicas := int32 (2 )
1087+ parties := []* kusciaapi.ServingParty {{
1088+ DomainId : "alice" ,
1089+ AppImage : "mockImageName" ,
1090+ Replicas : & replicas ,
1091+ }}
1092+
1093+ kusciaClient := kusciafake .NewSimpleClientset (makeMockAppImage ("mockImageName" ))
1094+ s := servingService {kusciaClient : kusciaClient }
1095+
1096+ tests := []struct {
1097+ name string
1098+ affinityMode string
1099+ expectAffinity bool
1100+ }{
1101+ {
1102+ name : "affinity_mode is none, should not set affinity in party template" ,
1103+ affinityMode : "none" ,
1104+ expectAffinity : false ,
1105+ },
1106+ {
1107+ name : "affinity_mode is affinity, should set PodAffinity" ,
1108+ affinityMode : "affinity" ,
1109+ expectAffinity : true ,
1110+ },
1111+ {
1112+ name : "affinity_mode is anti-affinity, should set PodAntiAffinity" ,
1113+ affinityMode : "anti-affinity" ,
1114+ expectAffinity : true ,
1115+ },
1116+ {
1117+ name : "affinity_mode is empty, should default to anti-affinity" ,
1118+ affinityMode : "" ,
1119+ expectAffinity : true ,
1120+ },
1121+ }
1122+
1123+ for _ , tt := range tests {
1124+ t .Run (tt .name , func (t * testing.T ) {
1125+ req := & kusciaapi.CreateServingRequest {
1126+ ServingId : "test-serving" ,
1127+ Initiator : "alice" ,
1128+ AffinityMode : tt .affinityMode ,
1129+ Parties : parties ,
1130+ }
1131+
1132+ kd , err := s .buildKusciaDeployment (context .Background (), req )
1133+ assert .NoError (t , err )
1134+ assert .NotNil (t , kd )
1135+
1136+ if tt .expectAffinity {
1137+ assert .NotNil (t , kd .Spec .Parties [0 ].Template .Spec .Affinity )
1138+ if tt .affinityMode == "affinity" || tt .affinityMode == "" {
1139+ // Empty mode defaults to anti-affinity
1140+ if tt .affinityMode == "affinity" {
1141+ assert .NotNil (t , kd .Spec .Parties [0 ].Template .Spec .Affinity .PodAffinity )
1142+ assert .Nil (t , kd .Spec .Parties [0 ].Template .Spec .Affinity .PodAntiAffinity )
1143+ } else {
1144+ assert .NotNil (t , kd .Spec .Parties [0 ].Template .Spec .Affinity .PodAntiAffinity )
1145+ }
1146+ }
1147+ } else {
1148+ // When affinity_mode is "none", affinity should be nil
1149+ assert .Nil (t , kd .Spec .Parties [0 ].Template .Spec .Affinity )
1150+ }
1151+ })
1152+ }
1153+ }
0 commit comments