Merge 3bb4d02f67 into 420bf6dff8

api/src/main/java/com/cloud/kubernetes/cluster/KubernetesServiceHelper.java

@@ -18,6 +18,7 @@ package com.cloud.kubernetes.cluster;
 
 import org.apache.cloudstack.acl.ControlledEntity;
 
+import java.util.List;
 import java.util.Map;
 
 import com.cloud.user.Account;
@@ -36,5 +37,6 @@ public interface KubernetesServiceHelper extends Adapter {
     boolean isValidNodeType(String nodeType);
     Map<String, Long> getServiceOfferingNodeTypeMap(Map<String, Map<String, String>> serviceOfferingNodeTypeMap);
     Map<String, Long> getTemplateNodeTypeMap(Map<String, Map<String, String>> templateNodeTypeMap);
+    Map<String, List<Long>> getAffinityGroupNodeTypeMap(Map<String, Map<String, String>> affinityGroupNodeTypeMap);
     void cleanupForAccount(Account account);
 }

api/src/main/java/com/cloud/vm/VmDetailConstants.java

@@ -96,6 +96,7 @@ public interface VmDetailConstants {
     String CKS_NODE_TYPE = "node";
     String OFFERING = "offering";
     String TEMPLATE = "template";
+    String AFFINITY_GROUP = "affinitygroup";
 
     // VMware to KVM VM migrations specific
     String VMWARE_TO_KVM_PREFIX = "vmware-to-kvm";

api/src/main/java/org/apache/cloudstack/api/ApiConstants.java

@@ -1237,6 +1237,13 @@ public class ApiConstants {
     public static final String MAX_SIZE = "maxsize";
     public static final String NODE_TYPE_OFFERING_MAP = "nodeofferings";
     public static final String NODE_TYPE_TEMPLATE_MAP = "nodetemplates";
+    public static final String NODE_TYPE_AFFINITY_GROUP_MAP = "nodeaffinitygroups";
+    public static final String CONTROL_AFFINITY_GROUP_IDS = "controlaffinitygroupids";
+    public static final String CONTROL_AFFINITY_GROUP_NAMES = "controlaffinitygroupnames";
+    public static final String WORKER_AFFINITY_GROUP_IDS = "workeraffinitygroupids";
+    public static final String WORKER_AFFINITY_GROUP_NAMES = "workeraffinitygroupnames";
+    public static final String ETCD_AFFINITY_GROUP_IDS = "etcdaffinitygroupids";
+    public static final String ETCD_AFFINITY_GROUP_NAMES = "etcdaffinitygroupnames";
 
     public static final String BOOT_TYPE = "boottype";
     public static final String BOOT_MODE = "bootmode";

engine/schema/src/main/resources/META-INF/db/schema-42210to42300.sql

@@ -34,6 +34,19 @@ CREATE TABLE `cloud`.`backup_offering_details` (
 UPDATE `cloud`.`configuration` SET value='random' WHERE name IN ('vm.allocation.algorithm', 'volume.allocation.algorithm') AND value='userconcentratedpod_random';
 UPDATE `cloud`.`configuration` SET value='firstfit' WHERE name IN ('vm.allocation.algorithm', 'volume.allocation.algorithm') AND value='userconcentratedpod_firstfit';
 
+-- Create kubernetes_cluster_affinity_group_map table for CKS per-node-type affinity groups
+CREATE TABLE IF NOT EXISTS `cloud`.`kubernetes_cluster_affinity_group_map` (
+    `id` bigint unsigned NOT NULL AUTO_INCREMENT,
+    `cluster_id` bigint unsigned NOT NULL COMMENT 'kubernetes cluster id',
+    `node_type` varchar(32) NOT NULL COMMENT 'CONTROL, WORKER, or ETCD',
+    `affinity_group_id` bigint unsigned NOT NULL COMMENT 'affinity group id',
+    PRIMARY KEY (`id`),
+    CONSTRAINT `fk_kubernetes_cluster_ag_map__cluster_id` FOREIGN KEY (`cluster_id`) REFERENCES `kubernetes_cluster`(`id`) ON DELETE CASCADE,
+    CONSTRAINT `fk_kubernetes_cluster_ag_map__ag_id` FOREIGN KEY (`affinity_group_id`) REFERENCES `affinity_group`(`id`) ON DELETE CASCADE,
+    INDEX `i_kubernetes_cluster_ag_map__cluster_id`(`cluster_id`),
+    INDEX `i_kubernetes_cluster_ag_map__ag_id`(`affinity_group_id`)
+) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
+
 -- Create webhook_filter table
 DROP TABLE IF EXISTS `cloud`.`webhook_filter`;
 CREATE TABLE IF NOT EXISTS `cloud`.`webhook_filter` (

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/KubernetesClusterAffinityGroupMapVO.java

@@ -0,0 +1,83 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+package com.cloud.kubernetes.cluster;
+
+import javax.persistence.Column;
+import javax.persistence.Entity;
+import javax.persistence.GeneratedValue;
+import javax.persistence.GenerationType;
+import javax.persistence.Id;
+import javax.persistence.Table;
+
+import org.apache.cloudstack.api.InternalIdentity;
+
+@Entity
+@Table(name = "kubernetes_cluster_affinity_group_map")
+public class KubernetesClusterAffinityGroupMapVO implements InternalIdentity {
+
+    @Id
+    @GeneratedValue(strategy = GenerationType.IDENTITY)
+    @Column(name = "id")
+    private Long id;
+
+    @Column(name = "cluster_id")
+    private long clusterId;
+
+    @Column(name = "node_type")
+    private String nodeType;
+
+    @Column(name = "affinity_group_id")
+    private long affinityGroupId;
+
+    public KubernetesClusterAffinityGroupMapVO() {
+    }
+
+    public KubernetesClusterAffinityGroupMapVO(long clusterId, String nodeType, long affinityGroupId) {
+        this.clusterId = clusterId;
+        this.nodeType = nodeType;
+        this.affinityGroupId = affinityGroupId;
+    }
+
+    @Override
+    public long getId() {
+        return id;
+    }
+
+    public long getClusterId() {
+        return clusterId;
+    }
+
+    public void setClusterId(long clusterId) {
+        this.clusterId = clusterId;
+    }
+
+    public String getNodeType() {
+        return nodeType;
+    }
+
+    public void setNodeType(String nodeType) {
+        this.nodeType = nodeType;
+    }
+
+    public long getAffinityGroupId() {
+        return affinityGroupId;
+    }
+
+    public void setAffinityGroupId(long affinityGroupId) {
+        this.affinityGroupId = affinityGroupId;
+    }
+}

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/KubernetesClusterManagerImpl.java

@@ -169,6 +169,7 @@ import com.cloud.kubernetes.cluster.actionworkers.KubernetesClusterScaleWorker;
 import com.cloud.kubernetes.cluster.actionworkers.KubernetesClusterStartWorker;
 import com.cloud.kubernetes.cluster.actionworkers.KubernetesClusterStopWorker;
 import com.cloud.kubernetes.cluster.actionworkers.KubernetesClusterUpgradeWorker;
+import com.cloud.kubernetes.cluster.dao.KubernetesClusterAffinityGroupMapDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDetailsDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterVmMapDao;
@@ -315,6 +316,8 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
     @Inject
     public KubernetesClusterDetailsDao kubernetesClusterDetailsDao;
     @Inject
+    public KubernetesClusterAffinityGroupMapDao kubernetesClusterAffinityGroupMapDao;
+    @Inject
     public KubernetesSupportedVersionDao kubernetesSupportedVersionDao;
     @Inject
     protected SSHKeyPairDao sshKeyPairDao;
@@ -858,24 +861,38 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
 
         List<KubernetesUserVmResponse> vmResponses = new ArrayList<>();
         List<KubernetesClusterVmMapVO> vmList = kubernetesClusterVmMapDao.listByClusterId(kubernetesCluster.getId());
-        ResponseView respView = ResponseView.Restricted;
+        ResponseView userVmResponseView = ResponseView.Restricted;
         Account caller = CallContext.current().getCallingAccount();
         if (accountService.isRootAdmin(caller.getId())) {
-            respView = ResponseView.Full;
+            userVmResponseView = ResponseView.Full;
         }
         final String responseName = "virtualmachine";
         if (vmList != null && !vmList.isEmpty()) {
-            for (KubernetesClusterVmMapVO vmMapVO : vmList) {
-                UserVmJoinVO userVM = userVmJoinDao.findById(vmMapVO.getVmId());
-                if (userVM != null) {
-                    UserVmResponse vmResponse = ApiDBUtils.newUserVmResponse(respView, responseName, userVM,
-                        EnumSet.of(VMDetails.nics), caller);
+            Map<Long, KubernetesClusterVmMapVO> vmMapById = vmList.stream()
+                    .collect(Collectors.toMap(KubernetesClusterVmMapVO::getVmId, vm -> vm));
+            Long[] vmIds = vmMapById.keySet().toArray(new Long[0]);
+            List<UserVmJoinVO> userVmJoinVOs = userVmJoinDao.searchByIds(vmIds);
+            if (userVmJoinVOs != null && !userVmJoinVOs.isEmpty()) {
+                Map<Long, UserVmResponse> vmResponseMap = new HashMap<>();
+                for (UserVmJoinVO userVM : userVmJoinVOs) {
+                    Long vmId = userVM.getId();
+                    UserVmResponse vmResponse = vmResponseMap.get(vmId);
+                    if (vmResponse == null) {
+                        vmResponse = ApiDBUtils.newUserVmResponse(userVmResponseView, responseName, userVM,
+                                EnumSet.of(VMDetails.nics, VMDetails.affgrp), caller);
+                        vmResponseMap.put(vmId, vmResponse);
+                    } else {
+                        ApiDBUtils.fillVmDetails(userVmResponseView, vmResponse, userVM);
+                    }
+                }
+                for (Map.Entry<Long, UserVmResponse> vmIdResponseEntry : vmResponseMap.entrySet()) {
                     KubernetesUserVmResponse kubernetesUserVmResponse = new KubernetesUserVmResponse();
                     try {
-                        BeanUtils.copyProperties(kubernetesUserVmResponse, vmResponse);
+                        BeanUtils.copyProperties(kubernetesUserVmResponse, vmIdResponseEntry.getValue());
                     } catch (IllegalAccessException | InvocationTargetException e) {
                         throw new ServerApiException(ApiErrorCode.INTERNAL_ERROR, "Failed to generate zone metrics response");
                     }
+                    KubernetesClusterVmMapVO vmMapVO = vmMapById.get(vmIdResponseEntry.getKey());
                     kubernetesUserVmResponse.setExternalNode(vmMapVO.isExternalNode());
                     kubernetesUserVmResponse.setEtcdNode(vmMapVO.isEtcdNode());
                     kubernetesUserVmResponse.setNodeVersion(vmMapVO.getNodeVersion());
@@ -905,10 +922,45 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
         response.setClusterType(kubernetesCluster.getClusterType());
         response.setCsiEnabled(kubernetesCluster.isCsiEnabled());
         response.setCreated(kubernetesCluster.getCreated());
+        setNodeTypeAffinityGroupResponse(response, kubernetesCluster.getId());
 
         return response;
     }
 
+    protected void setNodeTypeAffinityGroupResponse(KubernetesClusterResponse response, long clusterId) {
+        setAffinityGroupResponseForNodeType(response, clusterId, CONTROL.name());
+        setAffinityGroupResponseForNodeType(response, clusterId, WORKER.name());
+        setAffinityGroupResponseForNodeType(response, clusterId, ETCD.name());
+    }
+
+    protected void setAffinityGroupResponseForNodeType(KubernetesClusterResponse response, long clusterId, String nodeType) {
+        List<Long> affinityGroupIds = kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(clusterId, nodeType);
+        if (affinityGroupIds == null || affinityGroupIds.isEmpty()) {
+            return;
+        }
+        List<String> affinityGroupUuids = new ArrayList<>();
+        List<String> affinityGroupNames = new ArrayList<>();
+        for (Long affinityGroupId : affinityGroupIds) {
+            AffinityGroupVO affinityGroup = affinityGroupDao.findById(affinityGroupId);
+            if (affinityGroup != null) {
+                affinityGroupUuids.add(affinityGroup.getUuid());
+                affinityGroupNames.add(affinityGroup.getName());
+            }
+        }
+        String affinityGroupUuidsCsv = String.join(",", affinityGroupUuids);
+        String affinityGroupNamesCsv = String.join(",", affinityGroupNames);
+        if (CONTROL.name().equals(nodeType)) {
+            response.setControlAffinityGroupIds(affinityGroupUuidsCsv);
+            response.setControlAffinityGroupNames(affinityGroupNamesCsv);
+        } else if (WORKER.name().equals(nodeType)) {
+            response.setWorkerAffinityGroupIds(affinityGroupUuidsCsv);
+            response.setWorkerAffinityGroupNames(affinityGroupNamesCsv);
+        } else if (ETCD.name().equals(nodeType)) {
+            response.setEtcdAffinityGroupIds(affinityGroupUuidsCsv);
+            response.setEtcdAffinityGroupNames(affinityGroupNamesCsv);
+        }
+    }
+
     private DataCenter validateAndGetZoneForKubernetesCreateParameters(Long zoneId, Long networkId) {
         DataCenter zone = dataCenterDao.findById(zoneId);
         if (zone == null) {
@@ -1187,6 +1239,20 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
         return network;
     }
 
+    private void persistAffinityGroupMappings(long clusterId, Map<String, List<Long>> affinityGroupNodeTypeMap) {
+        if (MapUtils.isEmpty(affinityGroupNodeTypeMap)) {
+            return;
+        }
+        for (Map.Entry<String, List<Long>> nodeTypeAffinityGroupEntry : affinityGroupNodeTypeMap.entrySet()) {
+            String nodeType = nodeTypeAffinityGroupEntry.getKey();
+            List<Long> affinityGroupIds = nodeTypeAffinityGroupEntry.getValue();
+            for (Long affinityGroupId : affinityGroupIds) {
+                kubernetesClusterAffinityGroupMapDao.persist(
+                    new KubernetesClusterAffinityGroupMapVO(clusterId, nodeType, affinityGroupId));
+            }
+        }
+    }
+
     private void addKubernetesClusterDetails(final KubernetesCluster kubernetesCluster, final Network network, final CreateKubernetesClusterCmd cmd) {
         final String externalLoadBalancerIpAddress = cmd.getExternalLoadBalancerIpAddress();
         final String dockerRegistryUserName = cmd.getDockerRegistryUserName();
@@ -1626,6 +1692,7 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
         }
 
         Map<String, Long> templateNodeTypeMap = cmd.getTemplateNodeTypeMap();
+        Map<String, List<Long>> affinityGroupNodeTypeMap = cmd.getAffinityGroupNodeTypeMap();
         final VMTemplateVO finalTemplate = getKubernetesServiceTemplate(zone, hypervisorType, templateNodeTypeMap, DEFAULT, clusterKubernetesVersion);
         final VMTemplateVO controlNodeTemplate = getKubernetesServiceTemplate(zone, hypervisorType, templateNodeTypeMap, CONTROL, clusterKubernetesVersion);
         final VMTemplateVO workerNodeTemplate = getKubernetesServiceTemplate(zone, hypervisorType, templateNodeTypeMap, WORKER, clusterKubernetesVersion);
@@ -1671,6 +1738,7 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
                 }
                 newCluster.setCsiEnabled(cmd.getEnableCsi());
                 kubernetesClusterDao.persist(newCluster);
+                persistAffinityGroupMappings(newCluster.getId(), affinityGroupNodeTypeMap);
                 addKubernetesClusterDetails(newCluster, defaultNetwork, cmd);
                 return newCluster;
             }
@@ -2288,6 +2356,7 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
         if (validNodeIds.isEmpty()) {
             throw new CloudRuntimeException("No valid nodes found to be added to the Kubernetes cluster");
         }
+        validateNodeAffinityGroups(validNodeIds, kubernetesCluster);
         KubernetesClusterAddWorker addWorker = new KubernetesClusterAddWorker(kubernetesCluster, KubernetesClusterManagerImpl.this);
         addWorker = ComponentContext.inject(addWorker);
         return addWorker.addNodesToCluster(validNodeIds, cmd.isMountCksIsoOnVr(), cmd.isManualUpgrade());
@@ -2347,6 +2416,83 @@ public class KubernetesClusterManagerImpl extends ManagerBase implements Kuberne
         return validNodeIds;
     }
 
+    protected void validateNodeAffinityGroups(List<Long> nodeIds, KubernetesCluster cluster) {
+        List<Long> workerAffinityGroupIds = kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(
+                cluster.getId(), WORKER.name());
+        if (CollectionUtils.isEmpty(workerAffinityGroupIds)) {
+            return;
+        }
+
+        List<KubernetesClusterVmMapVO> existingWorkerVms = kubernetesClusterVmMapDao.listByClusterIdAndVmType(cluster.getId(), WORKER);
+        Set<Long> existingWorkerHostIds = new HashSet<>();
+        for (KubernetesClusterVmMapVO workerVmMap : existingWorkerVms) {
+            VMInstanceVO workerVm = vmInstanceDao.findById(workerVmMap.getVmId());
+            if (workerVm != null && workerVm.getHostId() != null) {
+                existingWorkerHostIds.add(workerVm.getHostId());
+            }
+        }
+
+        for (Long affinityGroupId : workerAffinityGroupIds) {
+            AffinityGroupVO affinityGroup = affinityGroupDao.findById(affinityGroupId);
+            if (affinityGroup == null) {
+                continue;
+            }
+            String affinityGroupType = affinityGroup.getType();
+
+            for (Long nodeId : nodeIds) {
+                VMInstanceVO node = vmInstanceDao.findById(nodeId);
+                if (node == null || node.getHostId() == null) {
+                    continue;
+                }
+                Long nodeHostId = node.getHostId();
+                HostVO nodeHost = hostDao.findById(nodeHostId);
+                String nodeHostName = nodeHost != null ? nodeHost.getName() : String.valueOf(nodeHostId);
+
+                if ("host anti-affinity".equalsIgnoreCase(affinityGroupType)) {
+                    if (existingWorkerHostIds.contains(nodeHostId)) {
+                        throw new InvalidParameterValueException(String.format(
+                                "Cannot add VM %s to cluster %s. VM is running on host %s which violates the cluster's " +
+                                "host anti-affinity rule (affinity group: %s). Existing worker VMs are already running on this host.",
+                                node.getInstanceName(), cluster.getName(), nodeHostName, affinityGroup.getName()));
+                    }
+                } else if ("host affinity".equalsIgnoreCase(affinityGroupType)) {
+                    if (!existingWorkerHostIds.isEmpty() && !existingWorkerHostIds.contains(nodeHostId)) {
+                        List<String> existingHostNames = new ArrayList<>();
+                        for (Long hostId : existingWorkerHostIds) {
+                            HostVO host = hostDao.findById(hostId);
+                            existingHostNames.add(host != null ? host.getName() : String.valueOf(hostId));
+                        }
+                        throw new InvalidParameterValueException(String.format(
+                                "Cannot add VM %s to cluster %s. VM is running on host %s which violates the cluster's " +
+                                "host affinity rule (affinity group: %s). All worker VMs must run on the same host. " +
+                                "Existing workers are on host(s): %s.",
+                                node.getInstanceName(), cluster.getName(), nodeHostName, affinityGroup.getName(),
+                                String.join(", ", existingHostNames)));
+                    }
+                }
+            }
+
+            if ("host anti-affinity".equalsIgnoreCase(affinityGroupType)) {
+                Set<Long> newNodeHostIds = new HashSet<>();
+                for (Long nodeId : nodeIds) {
+                    VMInstanceVO node = vmInstanceDao.findById(nodeId);
+                    if (node != null && node.getHostId() != null) {
+                        Long nodeHostId = node.getHostId();
+                        if (newNodeHostIds.contains(nodeHostId)) {
+                            HostVO nodeHost = hostDao.findById(nodeHostId);
+                            String nodeHostName = nodeHost != null ? nodeHost.getName() : String.valueOf(nodeHostId);
+                            throw new InvalidParameterValueException(String.format(
+                                    "Cannot add VM %s to cluster %s. Multiple VMs being added are running on the same host %s, " +
+                                    "which violates the cluster's host anti-affinity rule (affinity group: %s).",
+                                    node.getInstanceName(), cluster.getName(), nodeHostName, affinityGroup.getName()));
+                        }
+                        newNodeHostIds.add(nodeHostId);
+                    }
+                }
+            }
+        }
+    }
+
     @Override
     public List<RemoveVirtualMachinesFromKubernetesClusterResponse> removeVmsFromCluster(RemoveVirtualMachinesFromKubernetesClusterCmd cmd) {
         if (!KubernetesServiceEnabled.value()) {

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/KubernetesServiceHelperImpl.java

@@ -18,12 +18,16 @@ package com.cloud.kubernetes.cluster;
 
 import java.lang.reflect.Field;
 import java.lang.reflect.Modifier;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
 import javax.inject.Inject;
 
+import org.apache.cloudstack.affinity.AffinityGroup;
+import org.apache.cloudstack.affinity.dao.AffinityGroupDao;
 import com.cloud.exception.InvalidParameterValueException;
 import com.cloud.offering.ServiceOffering;
 import com.cloud.service.dao.ServiceOfferingDao;
@@ -66,6 +70,8 @@ public class KubernetesServiceHelperImpl extends AdapterBase implements Kubernet
     @Inject
     protected VMTemplateDao vmTemplateDao;
     @Inject
+    protected AffinityGroupDao affinityGroupDao;
+    @Inject
     KubernetesClusterService kubernetesClusterService;
 
     protected void setEventTypeEntityDetails(Class<?> eventTypeDefinedClass, Class<?> entityClass) {
@@ -244,6 +250,81 @@ public class KubernetesServiceHelperImpl extends AdapterBase implements Kubernet
         return mapping;
     }
 
+    protected void checkNodeTypeAffinityGroupEntryCompleteness(String nodeType, String affinityGroupUuids) {
+        if (StringUtils.isAnyBlank(nodeType, affinityGroupUuids)) {
+            String error = String.format("Any Node Type to Affinity Group entry should have a valid '%s' and '%s' values",
+                    VmDetailConstants.CKS_NODE_TYPE, VmDetailConstants.AFFINITY_GROUP);
+            logger.error(error);
+            throw new InvalidParameterValueException(error);
+        }
+    }
+
+    protected void checkNodeTypeAffinityGroupEntryNodeType(String nodeType) {
+        if (!isValidNodeType(nodeType)) {
+            String error = String.format("The provided value '%s' for Node Type is invalid", nodeType);
+            logger.error(error);
+            throw new InvalidParameterValueException(error);
+        }
+    }
+
+    protected Long validateAffinityGroupUuidAndGetId(String affinityGroupUuid) {
+        if (StringUtils.isBlank(affinityGroupUuid)) {
+            String error = "Empty affinity group UUID provided";
+            logger.error(error);
+            throw new InvalidParameterValueException(error);
+        }
+        AffinityGroup affinityGroup = affinityGroupDao.findByUuid(affinityGroupUuid);
+        if (affinityGroup == null) {
+            String error = String.format("Cannot find an affinity group with ID %s", affinityGroupUuid);
+            logger.error(error);
+            throw new InvalidParameterValueException(error);
+        }
+        return affinityGroup.getId();
+    }
+
+    protected List<Long> validateAndGetAffinityGroupIds(String affinityGroupUuids) {
+        String[] uuids = affinityGroupUuids.split(",");
+        List<Long> affinityGroupIds = new ArrayList<>();
+        for (String uuid : uuids) {
+            String trimmedUuid = uuid.trim();
+            Long affinityGroupId = validateAffinityGroupUuidAndGetId(trimmedUuid);
+            affinityGroupIds.add(affinityGroupId);
+        }
+        return affinityGroupIds;
+    }
+
+    protected void addNodeTypeAffinityGroupEntry(String nodeType, List<Long> affinityGroupIds, Map<String, List<Long>> nodeTypeToAffinityGroupIds) {
+        if (logger.isDebugEnabled()) {
+            logger.debug(String.format("Node Type: '%s' should use affinity group IDs: '%s'", nodeType, affinityGroupIds));
+        }
+        KubernetesClusterNodeType clusterNodeType = KubernetesClusterNodeType.valueOf(nodeType.toUpperCase());
+        nodeTypeToAffinityGroupIds.put(clusterNodeType.name(), affinityGroupIds);
+    }
+
+    protected void processNodeTypeAffinityGroupEntryAndAddToMappingIfValid(Map<String, String> nodeTypeAffinityConfig, Map<String, List<Long>> nodeTypeToAffinityGroupIds) {
+        if (MapUtils.isEmpty(nodeTypeAffinityConfig)) {
+            return;
+        }
+        String nodeType = nodeTypeAffinityConfig.get(VmDetailConstants.CKS_NODE_TYPE);
+        String affinityGroupUuids = nodeTypeAffinityConfig.get(VmDetailConstants.AFFINITY_GROUP);
+        checkNodeTypeAffinityGroupEntryCompleteness(nodeType, affinityGroupUuids);
+        checkNodeTypeAffinityGroupEntryNodeType(nodeType);
+
+        List<Long> affinityGroupIds = validateAndGetAffinityGroupIds(affinityGroupUuids);
+        addNodeTypeAffinityGroupEntry(nodeType, affinityGroupIds, nodeTypeToAffinityGroupIds);
+    }
+
+    @Override
+    public Map<String, List<Long>> getAffinityGroupNodeTypeMap(Map<String, Map<String, String>> affinityGroupNodeTypeMap) {
+        Map<String, List<Long>> nodeTypeToAffinityGroupIds = new HashMap<>();
+        if (MapUtils.isNotEmpty(affinityGroupNodeTypeMap)) {
+            for (Map<String, String> nodeTypeAffinityConfig : affinityGroupNodeTypeMap.values()) {
+                processNodeTypeAffinityGroupEntryAndAddToMappingIfValid(nodeTypeAffinityConfig, nodeTypeToAffinityGroupIds);
+            }
+        }
+        return nodeTypeToAffinityGroupIds;
+    }
+
     public void cleanupForAccount(Account account) {
         kubernetesClusterService.cleanupForAccount(account);
     }

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/actionworkers/KubernetesClusterActionWorker.java

@@ -90,6 +90,7 @@ import com.cloud.kubernetes.cluster.KubernetesClusterDetailsVO;
 import com.cloud.kubernetes.cluster.KubernetesClusterManagerImpl;
 import com.cloud.kubernetes.cluster.KubernetesClusterVO;
 import com.cloud.kubernetes.cluster.KubernetesClusterVmMapVO;
+import com.cloud.kubernetes.cluster.dao.KubernetesClusterAffinityGroupMapDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDetailsDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterVmMapDao;
@@ -217,6 +218,7 @@ public class KubernetesClusterActionWorker {
     protected KubernetesClusterDao kubernetesClusterDao;
     protected KubernetesClusterVmMapDao kubernetesClusterVmMapDao;
     protected KubernetesClusterDetailsDao kubernetesClusterDetailsDao;
+    protected KubernetesClusterAffinityGroupMapDao kubernetesClusterAffinityGroupMapDao;
     protected KubernetesSupportedVersionDao kubernetesSupportedVersionDao;
 
     protected KubernetesCluster kubernetesCluster;
@@ -251,6 +253,7 @@ public class KubernetesClusterActionWorker {
         this.kubernetesClusterDao = clusterManager.kubernetesClusterDao;
         this.kubernetesClusterDetailsDao = clusterManager.kubernetesClusterDetailsDao;
         this.kubernetesClusterVmMapDao = clusterManager.kubernetesClusterVmMapDao;
+        this.kubernetesClusterAffinityGroupMapDao = clusterManager.kubernetesClusterAffinityGroupMapDao;
         this.kubernetesSupportedVersionDao = clusterManager.kubernetesSupportedVersionDao;
         this.manager = clusterManager;
     }
@@ -1112,4 +1115,18 @@ public class KubernetesClusterActionWorker {
         }
         return null;
     }
+
+    protected List<Long> getAffinityGroupIdsForNodeType(KubernetesClusterNodeType nodeType) {
+        return new ArrayList<>(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(
+            kubernetesCluster.getId(), nodeType.name()));
+    }
+
+    protected List<Long> getMergedAffinityGroupIds(KubernetesClusterNodeType nodeType, Long domainId, Long accountId) {
+        List<Long> affinityGroupIds = getAffinityGroupIdsForNodeType(nodeType);
+        Long explicitAffinityGroupId = getExplicitAffinityGroup(domainId, accountId);
+        if (explicitAffinityGroupId != null && !affinityGroupIds.contains(explicitAffinityGroupId)) {
+            affinityGroupIds.add(explicitAffinityGroupId);
+        }
+        return affinityGroupIds.isEmpty() ? null : affinityGroupIds;
+    }
 }

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/actionworkers/KubernetesClusterDestroyWorker.java

@@ -348,6 +348,7 @@ public class KubernetesClusterDestroyWorker extends KubernetesClusterResourceMod
         stateTransitTo(kubernetesCluster.getId(), KubernetesCluster.Event.OperationSucceeded);
         annotationDao.removeByEntityType(AnnotationService.EntityType.KUBERNETES_CLUSTER.name(), kubernetesCluster.getUuid());
         kubernetesClusterDetailsDao.removeDetails(kubernetesCluster.getId());
+        kubernetesClusterAffinityGroupMapDao.removeByClusterId(kubernetesCluster.getId());
         boolean deleted = kubernetesClusterDao.remove(kubernetesCluster.getId());
         if (!deleted) {
             logMessage(Level.WARN, String.format("Failed to delete Kubernetes cluster: %s", kubernetesCluster), null);

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/actionworkers/KubernetesClusterResourceModifierActionWorker.java

@@ -26,7 +26,6 @@ import static com.cloud.utils.db.Transaction.execute;
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -426,21 +425,19 @@ public class KubernetesClusterResourceModifierActionWorker extends KubernetesClu
         if (StringUtils.isNotBlank(kubernetesCluster.getKeyPair())) {
             keypairs.add(kubernetesCluster.getKeyPair());
         }
-        Long affinityGroupId = getExplicitAffinityGroup(domainId, accountId);
+        List<Long> affinityGroupIds = getMergedAffinityGroupIds(WORKER, domainId, accountId);
         if (kubernetesCluster.getSecurityGroupId() != null && networkModel.checkSecurityGroupSupportForNetwork(owner, zone, networkIds, List.of(kubernetesCluster.getSecurityGroupId()))) {
             List<Long> securityGroupIds = new ArrayList<>();
             securityGroupIds.add(kubernetesCluster.getSecurityGroupId());
             nodeVm = userVmService.createAdvancedSecurityGroupVirtualMachine(zone, serviceOffering, workerNodeTemplate, networkIds, securityGroupIds, owner,
                     hostName, hostName, null, null, null, null, Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST,base64UserData, null, null, keypairs,
-                    null, addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null,
+                    null, addrs, null, null, affinityGroupIds, customParameterMap, null, null, null,
                     null, true, null, UserVmManager.CKS_NODE, null, null);
         } else {
             nodeVm = userVmService.createAdvancedVirtualMachine(zone, serviceOffering, workerNodeTemplate, networkIds, owner,
                     hostName, hostName, null, null, null, null,
                     Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST, base64UserData, null, null, keypairs,
-                    null, addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
+                    null, addrs, null, null, affinityGroupIds, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
         }
         if (logger.isInfoEnabled()) {
             logger.info("Created node VM : {}, {} in the Kubernetes cluster : {}", hostName, nodeVm, kubernetesCluster.getName());

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/actionworkers/KubernetesClusterStartWorker.java

@@ -270,7 +270,7 @@ public class KubernetesClusterStartWorker extends KubernetesClusterResourceModif
             keypairs.add(kubernetesCluster.getKeyPair());
         }
 
-        Long affinityGroupId = getExplicitAffinityGroup(domainId, accountId);
+        List<Long> affinityGroupIds = getMergedAffinityGroupIds(CONTROL, domainId, accountId);
         String userDataDetails = kubernetesCluster.getCniConfigDetails();
         if (kubernetesCluster.getSecurityGroupId() != null &&
                 networkModel.checkSecurityGroupSupportForNetwork(owner, zone, networkIds,
@@ -279,15 +279,13 @@ public class KubernetesClusterStartWorker extends KubernetesClusterResourceModif
             securityGroupIds.add(kubernetesCluster.getSecurityGroupId());
             controlVm = userVmService.createAdvancedSecurityGroupVirtualMachine(zone, serviceOffering, controlNodeTemplate, networkIds, securityGroupIds, owner,
             hostName, hostName, null, null, null, null, Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST,base64UserData, userDataId, userDataDetails, keypairs,
-                    requestedIps, addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null,
+                    requestedIps, addrs, null, null, affinityGroupIds, customParameterMap, null, null, null,
                     null, true, null, UserVmManager.CKS_NODE, null, null);
         } else {
             controlVm = userVmService.createAdvancedVirtualMachine(zone, serviceOffering, controlNodeTemplate, networkIds, owner,
                     hostName, hostName, null, null, null, null,
                     Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST, base64UserData, userDataId, userDataDetails, keypairs,
-                    requestedIps, addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
+                    requestedIps, addrs, null, null, affinityGroupIds, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
         }
         if (logger.isInfoEnabled()) {
             logger.info("Created control VM: {}, {} in the Kubernetes cluster: {}", controlVm, hostName, kubernetesCluster);
@@ -439,7 +437,7 @@ public class KubernetesClusterStartWorker extends KubernetesClusterResourceModif
             keypairs.add(kubernetesCluster.getKeyPair());
         }
 
-        Long affinityGroupId = getExplicitAffinityGroup(domainId, accountId);
+        List<Long> affinityGroupIds = getMergedAffinityGroupIds(CONTROL, domainId, accountId);
         if (kubernetesCluster.getSecurityGroupId() != null &&
                 networkModel.checkSecurityGroupSupportForNetwork(owner, zone, networkIds,
                         List.of(kubernetesCluster.getSecurityGroupId()))) {
@@ -447,15 +445,13 @@ public class KubernetesClusterStartWorker extends KubernetesClusterResourceModif
             securityGroupIds.add(kubernetesCluster.getSecurityGroupId());
             additionalControlVm = userVmService.createAdvancedSecurityGroupVirtualMachine(zone, serviceOffering, controlNodeTemplate, networkIds, securityGroupIds, owner,
                     hostName, hostName, null, null, null, null, Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST,base64UserData, null, null, keypairs,
-                    null, addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null,
+                    null, addrs, null, null, affinityGroupIds, customParameterMap, null, null, null,
                     null, true, null, UserVmManager.CKS_NODE, null, null);
         } else {
             additionalControlVm = userVmService.createAdvancedVirtualMachine(zone, serviceOffering, controlNodeTemplate, networkIds, owner,
                     hostName, hostName, null, null, null, null,
                     Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST, base64UserData, null, null, keypairs,
-                    null, addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
+                    null, addrs, null, null, affinityGroupIds, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
         }
 
         if (logger.isInfoEnabled()) {
@@ -483,7 +479,7 @@ public class KubernetesClusterStartWorker extends KubernetesClusterResourceModif
         if (StringUtils.isNotBlank(kubernetesCluster.getKeyPair())) {
             keypairs.add(kubernetesCluster.getKeyPair());
         }
-        Long affinityGroupId = getExplicitAffinityGroup(domainId, accountId);
+        List<Long> affinityGroupIds = getMergedAffinityGroupIds(ETCD, domainId, accountId);
         String hostName = etcdNodeHostnames.get(etcdNodeIndex);
         Map<String, String> customParameterMap = new HashMap<String, String>();
         if (zone.isSecurityGroupEnabled()) {
@@ -491,15 +487,13 @@ public class KubernetesClusterStartWorker extends KubernetesClusterResourceModif
             securityGroupIds.add(kubernetesCluster.getSecurityGroupId());
             etcdNode = userVmService.createAdvancedSecurityGroupVirtualMachine(zone, serviceOffering, etcdTemplate, networkIds, securityGroupIds, owner,
                     hostName, hostName, null, null, null, null, Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST,base64UserData, null, null, keypairs,
-                    Map.of(kubernetesCluster.getNetworkId(), requestedIps.get(etcdNodeIndex)), addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null,
+                    Map.of(kubernetesCluster.getNetworkId(), requestedIps.get(etcdNodeIndex)), addrs, null, null, affinityGroupIds, customParameterMap, null, null, null,
                     null, true, null, null, null, null);
         } else {
             etcdNode = userVmService.createAdvancedVirtualMachine(zone, serviceOffering, etcdTemplate, networkIds, owner,
                     hostName, hostName, null, null, null, null,
                     Hypervisor.HypervisorType.None, BaseCmd.HTTPMethod.POST, base64UserData, null, null, keypairs,
-                    Map.of(kubernetesCluster.getNetworkId(), requestedIps.get(etcdNodeIndex)), addrs, null, null, Objects.nonNull(affinityGroupId) ?
-                            Collections.singletonList(affinityGroupId) : null, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
+                    Map.of(kubernetesCluster.getNetworkId(), requestedIps.get(etcdNodeIndex)), addrs, null, null, affinityGroupIds, customParameterMap, null, null, null, null, true, UserVmManager.CKS_NODE, null, null, null);
         }
 
         if (logger.isInfoEnabled()) {

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/dao/KubernetesClusterAffinityGroupMapDao.java

@@ -0,0 +1,31 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+package com.cloud.kubernetes.cluster.dao;
+
+import java.util.List;
+
+import com.cloud.kubernetes.cluster.KubernetesClusterAffinityGroupMapVO;
+import com.cloud.utils.db.GenericDao;
+
+public interface KubernetesClusterAffinityGroupMapDao extends GenericDao<KubernetesClusterAffinityGroupMapVO, Long> {
+
+    List<KubernetesClusterAffinityGroupMapVO> listByClusterIdAndNodeType(long clusterId, String nodeType);
+
+    List<Long> listAffinityGroupIdsByClusterIdAndNodeType(long clusterId, String nodeType);
+
+    int removeByClusterId(long clusterId);
+}

plugins/integrations/kubernetes-service/src/main/java/com/cloud/kubernetes/cluster/dao/KubernetesClusterAffinityGroupMapDaoImpl.java

@@ -0,0 +1,67 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+package com.cloud.kubernetes.cluster.dao;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+import org.springframework.stereotype.Component;
+
+import com.cloud.kubernetes.cluster.KubernetesClusterAffinityGroupMapVO;
+import com.cloud.utils.db.GenericDaoBase;
+import com.cloud.utils.db.SearchBuilder;
+import com.cloud.utils.db.SearchCriteria;
+
+@Component
+public class KubernetesClusterAffinityGroupMapDaoImpl extends GenericDaoBase<KubernetesClusterAffinityGroupMapVO, Long>
+        implements KubernetesClusterAffinityGroupMapDao {
+
+    private final SearchBuilder<KubernetesClusterAffinityGroupMapVO> clusterIdAndNodeTypeSearch;
+    private final SearchBuilder<KubernetesClusterAffinityGroupMapVO> clusterIdSearch;
+
+    public KubernetesClusterAffinityGroupMapDaoImpl() {
+        clusterIdAndNodeTypeSearch = createSearchBuilder();
+        clusterIdAndNodeTypeSearch.and("clusterId", clusterIdAndNodeTypeSearch.entity().getClusterId(), SearchCriteria.Op.EQ);
+        clusterIdAndNodeTypeSearch.and("nodeType", clusterIdAndNodeTypeSearch.entity().getNodeType(), SearchCriteria.Op.EQ);
+        clusterIdAndNodeTypeSearch.done();
+
+        clusterIdSearch = createSearchBuilder();
+        clusterIdSearch.and("clusterId", clusterIdSearch.entity().getClusterId(), SearchCriteria.Op.EQ);
+        clusterIdSearch.done();
+    }
+
+    @Override
+    public List<KubernetesClusterAffinityGroupMapVO> listByClusterIdAndNodeType(long clusterId, String nodeType) {
+        SearchCriteria<KubernetesClusterAffinityGroupMapVO> sc = clusterIdAndNodeTypeSearch.create();
+        sc.setParameters("clusterId", clusterId);
+        sc.setParameters("nodeType", nodeType);
+        return listBy(sc);
+    }
+
+    @Override
+    public List<Long> listAffinityGroupIdsByClusterIdAndNodeType(long clusterId, String nodeType) {
+        List<KubernetesClusterAffinityGroupMapVO> maps = listByClusterIdAndNodeType(clusterId, nodeType);
+        return maps.stream().map(KubernetesClusterAffinityGroupMapVO::getAffinityGroupId).collect(Collectors.toList());
+    }
+
+    @Override
+    public int removeByClusterId(long clusterId) {
+        SearchCriteria<KubernetesClusterAffinityGroupMapVO> sc = clusterIdSearch.create();
+        sc.setParameters("clusterId", clusterId);
+        return remove(sc);
+    }
+}

plugins/integrations/kubernetes-service/src/main/java/org/apache/cloudstack/api/command/user/kubernetes/cluster/CreateKubernetesClusterCmd.java

@@ -17,6 +17,7 @@
 package org.apache.cloudstack.api.command.user.kubernetes.cluster;
 
 import java.security.InvalidParameterException;
+import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 
@@ -79,7 +80,7 @@ public class CreateKubernetesClusterCmd extends BaseAsyncCreateCmd {
     @Inject
     public KubernetesClusterService kubernetesClusterService;
     @Inject
-    protected KubernetesServiceHelper kubernetesClusterHelper;
+    protected KubernetesServiceHelper kubernetesServiceHelper;
     @Inject
     private ConfigurationDao configurationDao;
     @Inject
@@ -125,6 +126,12 @@ public class CreateKubernetesClusterCmd extends BaseAsyncCreateCmd {
             since = "4.21.0")
     private Map<String, Map<String, String>> templateNodeTypeMap;
 
+    @ACL(accessType = AccessType.UseEntry)
+    @Parameter(name = ApiConstants.NODE_TYPE_AFFINITY_GROUP_MAP, type = CommandType.MAP,
+            description = "(Optional) Node Type to Affinity Group ID mapping. If provided, VMs of each node type will be added to the specified affinity group",
+            since = "4.23.0")
+    private Map<String, Map<String, String>> affinityGroupNodeTypeMap;
+
     @ACL(accessType = AccessType.UseEntry)
     @Parameter(name = ApiConstants.ETCD_NODES, type = CommandType.LONG,
             description = "(Optional) Number of Kubernetes cluster etcd nodes, default is 0." +
@@ -314,11 +321,15 @@ public class CreateKubernetesClusterCmd extends BaseAsyncCreateCmd {
     }
 
     public Map<String, Long> getServiceOfferingNodeTypeMap() {
-        return kubernetesClusterHelper.getServiceOfferingNodeTypeMap(serviceOfferingNodeTypeMap);
+        return kubernetesServiceHelper.getServiceOfferingNodeTypeMap(serviceOfferingNodeTypeMap);
     }
 
     public Map<String, Long> getTemplateNodeTypeMap() {
-        return kubernetesClusterHelper.getTemplateNodeTypeMap(templateNodeTypeMap);
+        return kubernetesServiceHelper.getTemplateNodeTypeMap(templateNodeTypeMap);
+    }
+
+    public Map<String, List<Long>> getAffinityGroupNodeTypeMap() {
+        return kubernetesServiceHelper.getAffinityGroupNodeTypeMap(affinityGroupNodeTypeMap);
     }
 
     public Hypervisor.HypervisorType getHypervisorType() {

plugins/integrations/kubernetes-service/src/main/java/org/apache/cloudstack/api/command/user/kubernetes/cluster/ScaleKubernetesClusterCmd.java

@@ -57,7 +57,7 @@ public class ScaleKubernetesClusterCmd extends BaseAsyncCmd {
     @Inject
     public KubernetesClusterService kubernetesClusterService;
     @Inject
-    protected KubernetesServiceHelper kubernetesClusterHelper;
+    protected KubernetesServiceHelper kubernetesServiceHelper;
 
     /////////////////////////////////////////////////////
     //////////////// API parameters /////////////////////
@@ -114,7 +114,7 @@ public class ScaleKubernetesClusterCmd extends BaseAsyncCmd {
     }
 
     public Map<String, Long> getServiceOfferingNodeTypeMap() {
-        return kubernetesClusterHelper.getServiceOfferingNodeTypeMap(this.serviceOfferingNodeTypeMap);
+        return kubernetesServiceHelper.getServiceOfferingNodeTypeMap(this.serviceOfferingNodeTypeMap);
     }
 
     public Long getClusterSize() {

plugins/integrations/kubernetes-service/src/main/java/org/apache/cloudstack/api/response/KubernetesClusterResponse.java

@@ -220,6 +220,30 @@ public class KubernetesClusterResponse extends BaseResponseWithAnnotations imple
     @Param(description = "The date when this Kubernetes cluster was created")
     private Date created;
 
+    @SerializedName(ApiConstants.CONTROL_AFFINITY_GROUP_IDS)
+    @Param(description = "The IDs of affinity groups associated with control nodes", since = "4.23.0")
+    private String controlAffinityGroupIds;
+
+    @SerializedName(ApiConstants.CONTROL_AFFINITY_GROUP_NAMES)
+    @Param(description = "The names of affinity groups associated with control nodes", since = "4.23.0")
+    private String controlAffinityGroupNames;
+
+    @SerializedName(ApiConstants.WORKER_AFFINITY_GROUP_IDS)
+    @Param(description = "The IDs of affinity groups associated with worker nodes", since = "4.23.0")
+    private String workerAffinityGroupIds;
+
+    @SerializedName(ApiConstants.WORKER_AFFINITY_GROUP_NAMES)
+    @Param(description = "The names of affinity groups associated with worker nodes", since = "4.23.0")
+    private String workerAffinityGroupNames;
+
+    @SerializedName(ApiConstants.ETCD_AFFINITY_GROUP_IDS)
+    @Param(description = "The IDs of affinity groups associated with etcd nodes", since = "4.23.0")
+    private String etcdAffinityGroupIds;
+
+    @SerializedName(ApiConstants.ETCD_AFFINITY_GROUP_NAMES)
+    @Param(description = "The names of affinity groups associated with etcd nodes", since = "4.23.0")
+    private String etcdAffinityGroupNames;
+
     public KubernetesClusterResponse() {
     }
 
@@ -535,4 +559,28 @@ public class KubernetesClusterResponse extends BaseResponseWithAnnotations imple
     public void setCsiEnabled(Boolean csiEnabled) {
         isCsiEnabled = csiEnabled;
     }
+
+    public void setControlAffinityGroupIds(String controlAffinityGroupIds) {
+        this.controlAffinityGroupIds = controlAffinityGroupIds;
+    }
+
+    public void setControlAffinityGroupNames(String controlAffinityGroupNames) {
+        this.controlAffinityGroupNames = controlAffinityGroupNames;
+    }
+
+    public void setWorkerAffinityGroupIds(String workerAffinityGroupIds) {
+        this.workerAffinityGroupIds = workerAffinityGroupIds;
+    }
+
+    public void setWorkerAffinityGroupNames(String workerAffinityGroupNames) {
+        this.workerAffinityGroupNames = workerAffinityGroupNames;
+    }
+
+    public void setEtcdAffinityGroupIds(String etcdAffinityGroupIds) {
+        this.etcdAffinityGroupIds = etcdAffinityGroupIds;
+    }
+
+    public void setEtcdAffinityGroupNames(String etcdAffinityGroupNames) {
+        this.etcdAffinityGroupNames = etcdAffinityGroupNames;
+    }
 }

plugins/integrations/kubernetes-service/src/main/resources/META-INF/cloudstack/kubernetes-service/spring-kubernetes-service-context.xml

@@ -32,6 +32,7 @@
     <bean id="kubernetesClusterDaoImpl" class="com.cloud.kubernetes.cluster.dao.KubernetesClusterDaoImpl" />
     <bean id="kubernetesClusterDetailsDaoImpl" class="com.cloud.kubernetes.cluster.dao.KubernetesClusterDetailsDaoImpl" />
     <bean id="kubernetesClusterVmMapDaoImpl" class="com.cloud.kubernetes.cluster.dao.KubernetesClusterVmMapDaoImpl" />
+    <bean id="kubernetesClusterAffinityGroupMapDaoImpl" class="com.cloud.kubernetes.cluster.dao.KubernetesClusterAffinityGroupMapDaoImpl" />
     <bean id="kubernetesClusterManagerImpl" class="com.cloud.kubernetes.cluster.KubernetesClusterManagerImpl" />
 
     <bean id="kubernetesServiceHelper" class="com.cloud.kubernetes.cluster.KubernetesServiceHelperImpl" >

plugins/integrations/kubernetes-service/src/test/java/com/cloud/kubernetes/cluster/KubernetesClusterAffinityGroupMapVOTest.java

@@ -0,0 +1,87 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+package com.cloud.kubernetes.cluster;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class KubernetesClusterAffinityGroupMapVOTest {
+
+    @Test
+    public void testConstructorAndGetters() {
+        KubernetesClusterAffinityGroupMapVO vo =
+            new KubernetesClusterAffinityGroupMapVO(1L, "CONTROL", 100L);
+
+        Assert.assertEquals(1L, vo.getClusterId());
+        Assert.assertEquals("CONTROL", vo.getNodeType());
+        Assert.assertEquals(100L, vo.getAffinityGroupId());
+    }
+
+    @Test
+    public void testDefaultConstructor() {
+        KubernetesClusterAffinityGroupMapVO vo = new KubernetesClusterAffinityGroupMapVO();
+        Assert.assertNotNull(vo);
+    }
+
+    @Test
+    public void testSetClusterId() {
+        KubernetesClusterAffinityGroupMapVO vo = new KubernetesClusterAffinityGroupMapVO();
+        vo.setClusterId(2L);
+        Assert.assertEquals(2L, vo.getClusterId());
+    }
+
+    @Test
+    public void testSetNodeType() {
+        KubernetesClusterAffinityGroupMapVO vo = new KubernetesClusterAffinityGroupMapVO();
+        vo.setNodeType("WORKER");
+        Assert.assertEquals("WORKER", vo.getNodeType());
+    }
+
+    @Test
+    public void testSetAffinityGroupId() {
+        KubernetesClusterAffinityGroupMapVO vo = new KubernetesClusterAffinityGroupMapVO();
+        vo.setAffinityGroupId(200L);
+        Assert.assertEquals(200L, vo.getAffinityGroupId());
+    }
+
+    @Test
+    public void testAllNodeTypes() {
+        KubernetesClusterAffinityGroupMapVO controlVo =
+            new KubernetesClusterAffinityGroupMapVO(1L, "CONTROL", 10L);
+        KubernetesClusterAffinityGroupMapVO workerVo =
+            new KubernetesClusterAffinityGroupMapVO(1L, "WORKER", 20L);
+        KubernetesClusterAffinityGroupMapVO etcdVo =
+            new KubernetesClusterAffinityGroupMapVO(1L, "ETCD", 30L);
+
+        Assert.assertEquals("CONTROL", controlVo.getNodeType());
+        Assert.assertEquals("WORKER", workerVo.getNodeType());
+        Assert.assertEquals("ETCD", etcdVo.getNodeType());
+    }
+
+    @Test
+    public void testSettersChain() {
+        KubernetesClusterAffinityGroupMapVO vo = new KubernetesClusterAffinityGroupMapVO();
+
+        vo.setClusterId(5L);
+        vo.setNodeType("ETCD");
+        vo.setAffinityGroupId(500L);
+
+        Assert.assertEquals(5L, vo.getClusterId());
+        Assert.assertEquals("ETCD", vo.getNodeType());
+        Assert.assertEquals(500L, vo.getAffinityGroupId());
+    }
+}

plugins/integrations/kubernetes-service/src/test/java/com/cloud/kubernetes/cluster/KubernetesClusterHelperImplTest.java

@@ -1,145 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements.  See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership.  The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License.  You may obtain a copy of the License at
-//
-//   http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied.  See the License for the
-// specific language governing permissions and limitations
-// under the License.
-package com.cloud.kubernetes.cluster;
-
-import com.cloud.exception.InvalidParameterValueException;
-import com.cloud.service.ServiceOfferingVO;
-import com.cloud.service.dao.ServiceOfferingDao;
-import com.cloud.vm.VmDetailConstants;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.Mock;
-import org.mockito.Mockito;
-import org.mockito.junit.MockitoJUnitRunner;
-
-import java.util.HashMap;
-import java.util.Map;
-import java.util.UUID;
-
-import static com.cloud.kubernetes.cluster.KubernetesServiceHelper.KubernetesClusterNodeType.CONTROL;
-import static com.cloud.kubernetes.cluster.KubernetesServiceHelper.KubernetesClusterNodeType.ETCD;
-import static com.cloud.kubernetes.cluster.KubernetesServiceHelper.KubernetesClusterNodeType.WORKER;
-
-@RunWith(MockitoJUnitRunner.class)
-public class KubernetesClusterHelperImplTest {
-
-    @Mock
-    private ServiceOfferingDao serviceOfferingDao;
-    @Mock
-    private ServiceOfferingVO workerServiceOffering;
-    @Mock
-    private ServiceOfferingVO controlServiceOffering;
-    @Mock
-    private ServiceOfferingVO etcdServiceOffering;
-
-    private static final String workerNodesOfferingId = UUID.randomUUID().toString();
-    private static final String controlNodesOfferingId = UUID.randomUUID().toString();
-    private static final String etcdNodesOfferingId = UUID.randomUUID().toString();
-    private static final Long workerOfferingId = 1L;
-    private static final Long controlOfferingId = 2L;
-    private static final Long etcdOfferingId = 3L;
-
-    private final KubernetesServiceHelperImpl helper = new KubernetesServiceHelperImpl();
-
-    @Before
-    public void setUp() {
-        helper.serviceOfferingDao = serviceOfferingDao;
-        Mockito.when(serviceOfferingDao.findByUuid(workerNodesOfferingId)).thenReturn(workerServiceOffering);
-        Mockito.when(serviceOfferingDao.findByUuid(controlNodesOfferingId)).thenReturn(controlServiceOffering);
-        Mockito.when(serviceOfferingDao.findByUuid(etcdNodesOfferingId)).thenReturn(etcdServiceOffering);
-        Mockito.when(workerServiceOffering.getId()).thenReturn(workerOfferingId);
-        Mockito.when(controlServiceOffering.getId()).thenReturn(controlOfferingId);
-        Mockito.when(etcdServiceOffering.getId()).thenReturn(etcdOfferingId);
-    }
-
-    @Test
-    public void testIsValidNodeTypeEmptyNodeType() {
-        Assert.assertFalse(helper.isValidNodeType(null));
-    }
-
-    @Test
-    public void testIsValidNodeTypeInvalidNodeType() {
-        String nodeType = "invalidNodeType";
-        Assert.assertFalse(helper.isValidNodeType(nodeType));
-    }
-
-    @Test
-    public void testIsValidNodeTypeValidNodeTypeLowercase() {
-        String nodeType = KubernetesServiceHelper.KubernetesClusterNodeType.WORKER.name().toLowerCase();
-        Assert.assertTrue(helper.isValidNodeType(nodeType));
-    }
-
-    private Map<String, String> createMapEntry(KubernetesServiceHelper.KubernetesClusterNodeType nodeType,
-                                               String nodeTypeOfferingUuid) {
-        Map<String, String> map = new HashMap<>();
-        map.put(VmDetailConstants.CKS_NODE_TYPE, nodeType.name().toLowerCase());
-        map.put(VmDetailConstants.OFFERING, nodeTypeOfferingUuid);
-        return map;
-    }
-
-    @Test
-    public void testNodeOfferingMap() {
-        Map<String, Map<String, String>> serviceOfferingNodeTypeMap = new HashMap<>();
-        Map<String, String> firstMap = createMapEntry(WORKER, workerNodesOfferingId);
-        Map<String, String> secondMap = createMapEntry(CONTROL, controlNodesOfferingId);
-        serviceOfferingNodeTypeMap.put("map1", firstMap);
-        serviceOfferingNodeTypeMap.put("map2", secondMap);
-        Map<String, Long> map = helper.getServiceOfferingNodeTypeMap(serviceOfferingNodeTypeMap);
-        Assert.assertNotNull(map);
-        Assert.assertEquals(2, map.size());
-        Assert.assertTrue(map.containsKey(WORKER.name()) && map.containsKey(CONTROL.name()));
-        Assert.assertEquals(workerOfferingId, map.get(WORKER.name()));
-        Assert.assertEquals(controlOfferingId, map.get(CONTROL.name()));
-    }
-
-    @Test
-    public void testNodeOfferingMapNullMap() {
-        Map<String, Long> map = helper.getServiceOfferingNodeTypeMap(null);
-        Assert.assertTrue(map.isEmpty());
-    }
-
-    @Test
-    public void testNodeOfferingMapEtcdNodes() {
-        Map<String, Map<String, String>> serviceOfferingNodeTypeMap = new HashMap<>();
-        Map<String, String> firstMap = createMapEntry(ETCD, etcdNodesOfferingId);
-        serviceOfferingNodeTypeMap.put("map1", firstMap);
-        Map<String, Long> map = helper.getServiceOfferingNodeTypeMap(serviceOfferingNodeTypeMap);
-        Assert.assertNotNull(map);
-        Assert.assertEquals(1, map.size());
-        Assert.assertTrue(map.containsKey(ETCD.name()));
-        Assert.assertEquals(etcdOfferingId, map.get(ETCD.name()));
-    }
-
-    @Test(expected = InvalidParameterValueException.class)
-    public void testCheckNodeTypeOfferingEntryCompletenessInvalidParameters() {
-        helper.checkNodeTypeOfferingEntryCompleteness(WORKER.name(), null);
-    }
-
-    @Test(expected = InvalidParameterValueException.class)
-    public void testCheckNodeTypeOfferingEntryValuesInvalidNodeType() {
-        String invalidNodeType = "invalidNodeTypeName";
-        helper.checkNodeTypeOfferingEntryValues(invalidNodeType, workerServiceOffering, workerNodesOfferingId);
-    }
-
-    @Test(expected = InvalidParameterValueException.class)
-    public void testCheckNodeTypeOfferingEntryValuesEmptyOffering() {
-        String nodeType = WORKER.name();
-        helper.checkNodeTypeOfferingEntryValues(nodeType, null, workerNodesOfferingId);
-    }
-}

plugins/integrations/kubernetes-service/src/test/java/com/cloud/kubernetes/cluster/KubernetesClusterManagerImplTest.java

@@ -26,6 +26,7 @@ import com.cloud.dc.DataCenter;
 import com.cloud.exception.InvalidParameterValueException;
 import com.cloud.exception.PermissionDeniedException;
 import com.cloud.kubernetes.cluster.actionworkers.KubernetesClusterActionWorker;
+import com.cloud.kubernetes.cluster.dao.KubernetesClusterAffinityGroupMapDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterVmMapDao;
 import com.cloud.kubernetes.version.KubernetesSupportedVersion;
@@ -46,9 +47,14 @@ import com.cloud.utils.Pair;
 import com.cloud.utils.net.NetUtils;
 import com.cloud.vm.VMInstanceVO;
 import com.cloud.vm.dao.VMInstanceDao;
+import com.cloud.host.HostVO;
+import com.cloud.host.dao.HostDao;
+import org.apache.cloudstack.affinity.AffinityGroupVO;
+import org.apache.cloudstack.affinity.dao.AffinityGroupDao;
 import org.apache.cloudstack.api.BaseCmd;
 import org.apache.cloudstack.api.command.user.kubernetes.cluster.AddVirtualMachinesToKubernetesClusterCmd;
 import org.apache.cloudstack.api.command.user.kubernetes.cluster.RemoveVirtualMachinesFromKubernetesClusterCmd;
+import org.apache.cloudstack.api.response.KubernetesClusterResponse;
 import org.apache.cloudstack.context.CallContext;
 import org.apache.cloudstack.framework.config.ConfigKey;
 import org.apache.commons.collections.MapUtils;
@@ -103,6 +109,15 @@ public class KubernetesClusterManagerImplTest {
     @Mock
     private ServiceOfferingDao serviceOfferingDao;
 
+    @Mock
+    private KubernetesClusterAffinityGroupMapDao kubernetesClusterAffinityGroupMapDao;
+
+    @Mock
+    private AffinityGroupDao affinityGroupDao;
+
+    @Mock
+    private HostDao hostDao;
+
     @Spy
     @InjectMocks
     KubernetesClusterManagerImpl kubernetesClusterManager;
@@ -441,4 +456,462 @@ public class KubernetesClusterManagerImplTest {
         String cksClusterPreferredArch = kubernetesClusterManager.getCksClusterPreferredArch(systemVMArch, cksIso);
         Assert.assertEquals(CPU.CPUArch.amd64.getType(), cksClusterPreferredArch);
     }
+
+    @Test
+    public void testSetAffinityGroupResponseForNodeTypeControl() {
+        KubernetesClusterResponse response = new KubernetesClusterResponse();
+        long clusterId = 1L;
+
+        AffinityGroupVO ag1 = Mockito.mock(AffinityGroupVO.class);
+        AffinityGroupVO ag2 = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(ag1.getUuid()).thenReturn("uuid-1");
+        Mockito.when(ag1.getName()).thenReturn("affinity-group-1");
+        Mockito.when(ag2.getUuid()).thenReturn("uuid-2");
+        Mockito.when(ag2.getName()).thenReturn("affinity-group-2");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(clusterId, CONTROL.name()))
+            .thenReturn(Arrays.asList(1L, 2L));
+        Mockito.when(affinityGroupDao.findById(1L)).thenReturn(ag1);
+        Mockito.when(affinityGroupDao.findById(2L)).thenReturn(ag2);
+
+        kubernetesClusterManager.setAffinityGroupResponseForNodeType(response, clusterId, CONTROL.name());
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(clusterId, CONTROL.name());
+        Mockito.verify(affinityGroupDao).findById(1L);
+        Mockito.verify(affinityGroupDao).findById(2L);
+    }
+
+    @Test
+    public void testSetAffinityGroupResponseForNodeTypeWorker() {
+        KubernetesClusterResponse response = new KubernetesClusterResponse();
+        long clusterId = 1L;
+
+        AffinityGroupVO ag = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(ag.getUuid()).thenReturn("worker-uuid");
+        Mockito.when(ag.getName()).thenReturn("worker-affinity");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(clusterId, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(ag);
+
+        kubernetesClusterManager.setAffinityGroupResponseForNodeType(response, clusterId, WORKER.name());
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(clusterId, WORKER.name());
+        Mockito.verify(affinityGroupDao).findById(10L);
+    }
+
+    @Test
+    public void testSetAffinityGroupResponseForNodeTypeEtcd() {
+        KubernetesClusterResponse response = new KubernetesClusterResponse();
+        long clusterId = 1L;
+
+        AffinityGroupVO ag = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(ag.getUuid()).thenReturn("etcd-uuid");
+        Mockito.when(ag.getName()).thenReturn("etcd-affinity");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(clusterId, ETCD.name()))
+            .thenReturn(Arrays.asList(20L));
+        Mockito.when(affinityGroupDao.findById(20L)).thenReturn(ag);
+
+        kubernetesClusterManager.setAffinityGroupResponseForNodeType(response, clusterId, ETCD.name());
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(clusterId, ETCD.name());
+        Mockito.verify(affinityGroupDao).findById(20L);
+    }
+
+    @Test
+    public void testSetAffinityGroupResponseForNodeTypeEmptyList() {
+        KubernetesClusterResponse response = new KubernetesClusterResponse();
+        long clusterId = 1L;
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(clusterId, CONTROL.name()))
+            .thenReturn(Collections.emptyList());
+
+        kubernetesClusterManager.setAffinityGroupResponseForNodeType(response, clusterId, CONTROL.name());
+
+        Mockito.verify(affinityGroupDao, Mockito.never()).findById(Mockito.anyLong());
+    }
+
+    @Test
+    public void testSetAffinityGroupResponseForNodeTypeNullList() {
+        KubernetesClusterResponse response = new KubernetesClusterResponse();
+        long clusterId = 1L;
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(clusterId, ETCD.name()))
+            .thenReturn(null);
+
+        kubernetesClusterManager.setAffinityGroupResponseForNodeType(response, clusterId, ETCD.name());
+
+        Mockito.verify(affinityGroupDao, Mockito.never()).findById(Mockito.anyLong());
+    }
+
+    @Test
+    public void testSetAffinityGroupResponseForNodeTypeNullAffinityGroup() {
+        KubernetesClusterResponse response = new KubernetesClusterResponse();
+        long clusterId = 1L;
+
+        AffinityGroupVO ag1 = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(ag1.getUuid()).thenReturn("uuid-1");
+        Mockito.when(ag1.getName()).thenReturn("affinity-group-1");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(clusterId, CONTROL.name()))
+            .thenReturn(Arrays.asList(1L, 2L));
+        Mockito.when(affinityGroupDao.findById(1L)).thenReturn(ag1);
+        Mockito.when(affinityGroupDao.findById(2L)).thenReturn(null);
+
+        kubernetesClusterManager.setAffinityGroupResponseForNodeType(response, clusterId, CONTROL.name());
+
+        Mockito.verify(affinityGroupDao).findById(1L);
+        Mockito.verify(affinityGroupDao).findById(2L);
+    }
+
+    @Test
+    public void testSetNodeTypeAffinityGroupResponse() {
+        KubernetesClusterResponse response = new KubernetesClusterResponse();
+        long clusterId = 1L;
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(Mockito.eq(clusterId), Mockito.anyString()))
+            .thenReturn(Collections.emptyList());
+
+        kubernetesClusterManager.setNodeTypeAffinityGroupResponse(response, clusterId);
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(clusterId, CONTROL.name());
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(clusterId, WORKER.name());
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(clusterId, ETCD.name());
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsNoAffinityGroups() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        List<Long> nodeIds = Arrays.asList(100L, 101L);
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Collections.emptyList());
+
+        kubernetesClusterManager.validateNodeAffinityGroups(nodeIds, cluster);
+
+        Mockito.verify(kubernetesClusterVmMapDao, Mockito.never()).listByClusterIdAndVmType(Mockito.anyLong(), Mockito.any());
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsNullAffinityGroups() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        List<Long> nodeIds = Arrays.asList(100L, 101L);
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(null);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(nodeIds, cluster);
+
+        Mockito.verify(kubernetesClusterVmMapDao, Mockito.never()).listByClusterIdAndVmType(Mockito.anyLong(), Mockito.any());
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testValidateNodeAffinityGroupsAntiAffinityNewNodeOnExistingHost() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId = 100L;
+        Long existingWorkerVmId = 200L;
+        Long sharedHostId = 1000L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup.getType()).thenReturn("host anti-affinity");
+        Mockito.when(affinityGroup.getName()).thenReturn("anti-affinity-group");
+
+        VMInstanceVO newNode = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode.getHostId()).thenReturn(sharedHostId);
+        Mockito.when(newNode.getInstanceName()).thenReturn("new-node-vm");
+
+        VMInstanceVO existingWorkerVm = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(existingWorkerVm.getHostId()).thenReturn(sharedHostId);
+
+        KubernetesClusterVmMapVO workerVmMap = Mockito.mock(KubernetesClusterVmMapVO.class);
+        Mockito.when(workerVmMap.getVmId()).thenReturn(existingWorkerVmId);
+
+        HostVO host = Mockito.mock(HostVO.class);
+        Mockito.when(host.getName()).thenReturn("host-1");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Arrays.asList(workerVmMap));
+        Mockito.when(vmInstanceDao.findById(existingWorkerVmId)).thenReturn(existingWorkerVm);
+        Mockito.when(vmInstanceDao.findById(newNodeId)).thenReturn(newNode);
+        Mockito.when(hostDao.findById(sharedHostId)).thenReturn(host);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId), cluster);
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsAntiAffinityNewNodeOnDifferentHost() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.lenient().when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId = 100L;
+        Long existingWorkerVmId = 200L;
+        Long existingHostId = 1000L;
+        Long newNodeHostId = 1001L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup.getType()).thenReturn("host anti-affinity");
+        Mockito.lenient().when(affinityGroup.getName()).thenReturn("anti-affinity-group");
+
+        VMInstanceVO newNode = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode.getHostId()).thenReturn(newNodeHostId);
+
+        VMInstanceVO existingWorkerVm = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(existingWorkerVm.getHostId()).thenReturn(existingHostId);
+
+        KubernetesClusterVmMapVO workerVmMap = Mockito.mock(KubernetesClusterVmMapVO.class);
+        Mockito.when(workerVmMap.getVmId()).thenReturn(existingWorkerVmId);
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Arrays.asList(workerVmMap));
+        Mockito.when(vmInstanceDao.findById(existingWorkerVmId)).thenReturn(existingWorkerVm);
+        Mockito.when(vmInstanceDao.findById(newNodeId)).thenReturn(newNode);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId), cluster);
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name());
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsAffinityNewNodeOnSameHost() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.lenient().when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId = 100L;
+        Long existingWorkerVmId = 200L;
+        Long sharedHostId = 1000L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup.getType()).thenReturn("host affinity");
+        Mockito.lenient().when(affinityGroup.getName()).thenReturn("affinity-group");
+
+        VMInstanceVO newNode = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode.getHostId()).thenReturn(sharedHostId);
+
+        VMInstanceVO existingWorkerVm = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(existingWorkerVm.getHostId()).thenReturn(sharedHostId);
+
+        KubernetesClusterVmMapVO workerVmMap = Mockito.mock(KubernetesClusterVmMapVO.class);
+        Mockito.when(workerVmMap.getVmId()).thenReturn(existingWorkerVmId);
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Arrays.asList(workerVmMap));
+        Mockito.when(vmInstanceDao.findById(existingWorkerVmId)).thenReturn(existingWorkerVm);
+        Mockito.when(vmInstanceDao.findById(newNodeId)).thenReturn(newNode);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId), cluster);
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name());
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testValidateNodeAffinityGroupsAffinityNewNodeOnDifferentHost() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId = 100L;
+        Long existingWorkerVmId = 200L;
+        Long existingHostId = 1000L;
+        Long newNodeHostId = 1001L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup.getType()).thenReturn("host affinity");
+        Mockito.when(affinityGroup.getName()).thenReturn("affinity-group");
+
+        VMInstanceVO newNode = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode.getHostId()).thenReturn(newNodeHostId);
+        Mockito.when(newNode.getInstanceName()).thenReturn("new-node-vm");
+
+        VMInstanceVO existingWorkerVm = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(existingWorkerVm.getHostId()).thenReturn(existingHostId);
+
+        KubernetesClusterVmMapVO workerVmMap = Mockito.mock(KubernetesClusterVmMapVO.class);
+        Mockito.when(workerVmMap.getVmId()).thenReturn(existingWorkerVmId);
+
+        HostVO newHost = Mockito.mock(HostVO.class);
+        Mockito.when(newHost.getName()).thenReturn("host-2");
+
+        HostVO existingHost = Mockito.mock(HostVO.class);
+        Mockito.when(existingHost.getName()).thenReturn("host-1");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Arrays.asList(workerVmMap));
+        Mockito.when(vmInstanceDao.findById(existingWorkerVmId)).thenReturn(existingWorkerVm);
+        Mockito.when(vmInstanceDao.findById(newNodeId)).thenReturn(newNode);
+        Mockito.when(hostDao.findById(newNodeHostId)).thenReturn(newHost);
+        Mockito.when(hostDao.findById(existingHostId)).thenReturn(existingHost);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId), cluster);
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testValidateNodeAffinityGroupsAntiAffinityMultipleNewNodesOnSameHost() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId1 = 100L;
+        Long newNodeId2 = 101L;
+        Long sharedHostId = 1000L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.lenient().when(affinityGroup.getType()).thenReturn("host anti-affinity");
+        Mockito.when(affinityGroup.getName()).thenReturn("anti-affinity-group");
+
+        VMInstanceVO newNode1 = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode1.getHostId()).thenReturn(sharedHostId);
+        Mockito.lenient().when(newNode1.getInstanceName()).thenReturn("new-node-vm-1");
+
+        VMInstanceVO newNode2 = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode2.getHostId()).thenReturn(sharedHostId);
+        Mockito.when(newNode2.getInstanceName()).thenReturn("new-node-vm-2");
+
+        HostVO host = Mockito.mock(HostVO.class);
+        Mockito.when(host.getName()).thenReturn("host-1");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Collections.emptyList());
+        Mockito.when(vmInstanceDao.findById(newNodeId1)).thenReturn(newNode1);
+        Mockito.when(vmInstanceDao.findById(newNodeId2)).thenReturn(newNode2);
+        Mockito.when(hostDao.findById(sharedHostId)).thenReturn(host);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId1, newNodeId2), cluster);
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsAntiAffinityMultipleNewNodesOnDifferentHosts() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.lenient().when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId1 = 100L;
+        Long newNodeId2 = 101L;
+        Long hostId1 = 1000L;
+        Long hostId2 = 1001L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.lenient().when(affinityGroup.getType()).thenReturn("host anti-affinity");
+        Mockito.lenient().when(affinityGroup.getName()).thenReturn("anti-affinity-group");
+
+        VMInstanceVO newNode1 = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode1.getHostId()).thenReturn(hostId1);
+
+        VMInstanceVO newNode2 = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode2.getHostId()).thenReturn(hostId2);
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Collections.emptyList());
+        Mockito.when(vmInstanceDao.findById(newNodeId1)).thenReturn(newNode1);
+        Mockito.when(vmInstanceDao.findById(newNodeId2)).thenReturn(newNode2);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId1, newNodeId2), cluster);
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name());
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsNodeWithNullHost() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.lenient().when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId = 100L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.lenient().when(affinityGroup.getType()).thenReturn("host anti-affinity");
+        Mockito.lenient().when(affinityGroup.getName()).thenReturn("anti-affinity-group");
+
+        VMInstanceVO newNode = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode.getHostId()).thenReturn(null);
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Collections.emptyList());
+        Mockito.when(vmInstanceDao.findById(newNodeId)).thenReturn(newNode);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId), cluster);
+
+        Mockito.verify(vmInstanceDao, Mockito.atLeastOnce()).findById(newNodeId);
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsNullNode() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.lenient().when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId = 100L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.lenient().when(affinityGroup.getType()).thenReturn("host anti-affinity");
+        Mockito.lenient().when(affinityGroup.getName()).thenReturn("anti-affinity-group");
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Collections.emptyList());
+        Mockito.when(vmInstanceDao.findById(newNodeId)).thenReturn(null);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId), cluster);
+
+        Mockito.verify(vmInstanceDao, Mockito.atLeastOnce()).findById(newNodeId);
+    }
+
+    @Test
+    public void testValidateNodeAffinityGroupsAffinityNoExistingWorkers() {
+        KubernetesCluster cluster = Mockito.mock(KubernetesCluster.class);
+        Mockito.when(cluster.getId()).thenReturn(1L);
+        Mockito.lenient().when(cluster.getName()).thenReturn("test-cluster");
+
+        Long newNodeId = 100L;
+        Long newNodeHostId = 1000L;
+
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.lenient().when(affinityGroup.getType()).thenReturn("host affinity");
+        Mockito.lenient().when(affinityGroup.getName()).thenReturn("affinity-group");
+
+        VMInstanceVO newNode = Mockito.mock(VMInstanceVO.class);
+        Mockito.when(newNode.getHostId()).thenReturn(newNodeHostId);
+
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name()))
+            .thenReturn(Arrays.asList(10L));
+        Mockito.when(affinityGroupDao.findById(10L)).thenReturn(affinityGroup);
+        Mockito.when(kubernetesClusterVmMapDao.listByClusterIdAndVmType(1L, WORKER))
+            .thenReturn(Collections.emptyList());
+        Mockito.when(vmInstanceDao.findById(newNodeId)).thenReturn(newNode);
+
+        kubernetesClusterManager.validateNodeAffinityGroups(Arrays.asList(newNodeId), cluster);
+
+        Mockito.verify(kubernetesClusterAffinityGroupMapDao).listAffinityGroupIdsByClusterIdAndNodeType(1L, WORKER.name());
+    }
+
 }

plugins/integrations/kubernetes-service/src/test/java/com/cloud/kubernetes/cluster/KubernetesServiceHelperImplTest.java

@@ -17,6 +17,15 @@
 package com.cloud.kubernetes.cluster;
 
 
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import org.apache.cloudstack.affinity.AffinityGroupVO;
+import org.apache.cloudstack.affinity.dao.AffinityGroupDao;
+import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.InjectMocks;
@@ -24,11 +33,16 @@ import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnitRunner;
 
+import com.cloud.exception.InvalidParameterValueException;
+import com.cloud.kubernetes.cluster.KubernetesServiceHelper.KubernetesClusterNodeType;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterVmMapDao;
+import com.cloud.service.ServiceOfferingVO;
+import com.cloud.service.dao.ServiceOfferingDao;
 import com.cloud.uservm.UserVm;
 import com.cloud.utils.exception.CloudRuntimeException;
 import com.cloud.vm.UserVmManager;
+import com.cloud.vm.VmDetailConstants;
 
 @RunWith(MockitoJUnitRunner.class)
 public class KubernetesServiceHelperImplTest {
@@ -36,6 +50,10 @@ public class KubernetesServiceHelperImplTest {
     KubernetesClusterVmMapDao kubernetesClusterVmMapDao;
     @Mock
     KubernetesClusterDao kubernetesClusterDao;
+    @Mock
+    AffinityGroupDao affinityGroupDao;
+    @Mock
+    ServiceOfferingDao serviceOfferingDao;
 
     @InjectMocks
     KubernetesServiceHelperImpl kubernetesServiceHelper = new KubernetesServiceHelperImpl();
@@ -84,4 +102,302 @@ public class KubernetesServiceHelperImplTest {
         Mockito.when(kubernetesCluster.getClusterType()).thenReturn(KubernetesCluster.ClusterType.ExternalManaged);
         kubernetesServiceHelper.checkVmCanBeDestroyed(vm);
     }
+
+    @Test
+    public void testIsValidNodeTypeEmptyNodeType() {
+        Assert.assertFalse(kubernetesServiceHelper.isValidNodeType(null));
+    }
+
+    @Test
+    public void testIsValidNodeTypeInvalidNodeType() {
+        Assert.assertFalse(kubernetesServiceHelper.isValidNodeType("invalidNodeType"));
+    }
+
+    @Test
+    public void testIsValidNodeTypeValidNodeTypeLowercase() {
+        String nodeType = KubernetesClusterNodeType.WORKER.name().toLowerCase();
+        Assert.assertTrue(kubernetesServiceHelper.isValidNodeType(nodeType));
+    }
+
+    private Map<String, String> createServiceOfferingMapEntry(KubernetesClusterNodeType nodeType, String offeringUuid) {
+        Map<String, String> map = new HashMap<>();
+        map.put(VmDetailConstants.CKS_NODE_TYPE, nodeType.name().toLowerCase());
+        map.put(VmDetailConstants.OFFERING, offeringUuid);
+        return map;
+    }
+
+    @Test
+    public void testGetServiceOfferingNodeTypeMap() {
+        String workerOfferingUuid = UUID.randomUUID().toString();
+        String controlOfferingUuid = UUID.randomUUID().toString();
+
+        ServiceOfferingVO workerOffering = Mockito.mock(ServiceOfferingVO.class);
+        Mockito.when(workerOffering.getId()).thenReturn(1L);
+        Mockito.when(serviceOfferingDao.findByUuid(workerOfferingUuid)).thenReturn(workerOffering);
+
+        ServiceOfferingVO controlOffering = Mockito.mock(ServiceOfferingVO.class);
+        Mockito.when(controlOffering.getId()).thenReturn(2L);
+        Mockito.when(serviceOfferingDao.findByUuid(controlOfferingUuid)).thenReturn(controlOffering);
+
+        Map<String, Map<String, String>> serviceOfferingNodeTypeMap = new HashMap<>();
+        serviceOfferingNodeTypeMap.put("map1", createServiceOfferingMapEntry(KubernetesClusterNodeType.WORKER, workerOfferingUuid));
+        serviceOfferingNodeTypeMap.put("map2", createServiceOfferingMapEntry(KubernetesClusterNodeType.CONTROL, controlOfferingUuid));
+
+        Map<String, Long> result = kubernetesServiceHelper.getServiceOfferingNodeTypeMap(serviceOfferingNodeTypeMap);
+
+        Assert.assertNotNull(result);
+        Assert.assertEquals(2, result.size());
+        Assert.assertTrue(result.containsKey(KubernetesClusterNodeType.WORKER.name()));
+        Assert.assertTrue(result.containsKey(KubernetesClusterNodeType.CONTROL.name()));
+        Assert.assertEquals(Long.valueOf(1L), result.get(KubernetesClusterNodeType.WORKER.name()));
+        Assert.assertEquals(Long.valueOf(2L), result.get(KubernetesClusterNodeType.CONTROL.name()));
+    }
+
+    @Test
+    public void testGetServiceOfferingNodeTypeMapNullMap() {
+        Map<String, Long> result = kubernetesServiceHelper.getServiceOfferingNodeTypeMap(null);
+        Assert.assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void testGetServiceOfferingNodeTypeMapEtcdNodes() {
+        String etcdOfferingUuid = UUID.randomUUID().toString();
+
+        ServiceOfferingVO etcdOffering = Mockito.mock(ServiceOfferingVO.class);
+        Mockito.when(etcdOffering.getId()).thenReturn(3L);
+        Mockito.when(serviceOfferingDao.findByUuid(etcdOfferingUuid)).thenReturn(etcdOffering);
+
+        Map<String, Map<String, String>> serviceOfferingNodeTypeMap = new HashMap<>();
+        serviceOfferingNodeTypeMap.put("map1", createServiceOfferingMapEntry(KubernetesClusterNodeType.ETCD, etcdOfferingUuid));
+
+        Map<String, Long> result = kubernetesServiceHelper.getServiceOfferingNodeTypeMap(serviceOfferingNodeTypeMap);
+
+        Assert.assertNotNull(result);
+        Assert.assertEquals(1, result.size());
+        Assert.assertTrue(result.containsKey(KubernetesClusterNodeType.ETCD.name()));
+        Assert.assertEquals(Long.valueOf(3L), result.get(KubernetesClusterNodeType.ETCD.name()));
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testCheckNodeTypeOfferingEntryCompletenessInvalidParameters() {
+        kubernetesServiceHelper.checkNodeTypeOfferingEntryCompleteness(KubernetesClusterNodeType.WORKER.name(), null);
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testCheckNodeTypeOfferingEntryValuesInvalidNodeType() {
+        ServiceOfferingVO offering = Mockito.mock(ServiceOfferingVO.class);
+        kubernetesServiceHelper.checkNodeTypeOfferingEntryValues("invalidNodeTypeName", offering, "some-uuid");
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testCheckNodeTypeOfferingEntryValuesEmptyOffering() {
+        kubernetesServiceHelper.checkNodeTypeOfferingEntryValues(KubernetesClusterNodeType.WORKER.name(), null, "some-uuid");
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testCheckNodeTypeAffinityGroupEntryCompletenessBlankNodeType() {
+        kubernetesServiceHelper.checkNodeTypeAffinityGroupEntryCompleteness("", "affinity-group-uuid");
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testCheckNodeTypeAffinityGroupEntryCompletenessBlankAffinityGroupUuid() {
+        kubernetesServiceHelper.checkNodeTypeAffinityGroupEntryCompleteness("control", "");
+    }
+
+    @Test
+    public void testCheckNodeTypeAffinityGroupEntryCompletenessValid() {
+        kubernetesServiceHelper.checkNodeTypeAffinityGroupEntryCompleteness("control", "affinity-group-uuid");
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testCheckNodeTypeAffinityGroupEntryNodeTypeInvalid() {
+        kubernetesServiceHelper.checkNodeTypeAffinityGroupEntryNodeType("invalid-node-type");
+    }
+
+    @Test
+    public void testCheckNodeTypeAffinityGroupEntryNodeTypeValid() {
+        kubernetesServiceHelper.checkNodeTypeAffinityGroupEntryNodeType("control");
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testValidateAffinityGroupUuidAndGetIdBlank() {
+        kubernetesServiceHelper.validateAffinityGroupUuidAndGetId("");
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testValidateAffinityGroupUuidAndGetIdNotFound() {
+        Mockito.when(affinityGroupDao.findByUuid("non-existent-uuid")).thenReturn(null);
+        kubernetesServiceHelper.validateAffinityGroupUuidAndGetId("non-existent-uuid");
+    }
+
+    @Test
+    public void testValidateAffinityGroupUuidAndGetIdValid() {
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup.getId()).thenReturn(100L);
+        Mockito.when(affinityGroupDao.findByUuid("valid-uuid")).thenReturn(affinityGroup);
+        Long result = kubernetesServiceHelper.validateAffinityGroupUuidAndGetId("valid-uuid");
+        Assert.assertEquals(Long.valueOf(100L), result);
+    }
+
+    @Test
+    public void testValidateAndGetAffinityGroupIdsSingleUuid() {
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup.getId()).thenReturn(1L);
+        Mockito.when(affinityGroupDao.findByUuid("uuid1")).thenReturn(affinityGroup);
+
+        List<Long> result = kubernetesServiceHelper.validateAndGetAffinityGroupIds("uuid1");
+        Assert.assertEquals(1, result.size());
+        Assert.assertEquals(Long.valueOf(1L), result.get(0));
+    }
+
+    @Test
+    public void testValidateAndGetAffinityGroupIdsMultipleUuids() {
+        AffinityGroupVO affinityGroup1 = Mockito.mock(AffinityGroupVO.class);
+        AffinityGroupVO affinityGroup2 = Mockito.mock(AffinityGroupVO.class);
+        AffinityGroupVO affinityGroup3 = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup1.getId()).thenReturn(1L);
+        Mockito.when(affinityGroup2.getId()).thenReturn(2L);
+        Mockito.when(affinityGroup3.getId()).thenReturn(3L);
+        Mockito.when(affinityGroupDao.findByUuid("uuid1")).thenReturn(affinityGroup1);
+        Mockito.when(affinityGroupDao.findByUuid("uuid2")).thenReturn(affinityGroup2);
+        Mockito.when(affinityGroupDao.findByUuid("uuid3")).thenReturn(affinityGroup3);
+
+        List<Long> result = kubernetesServiceHelper.validateAndGetAffinityGroupIds("uuid1,uuid2,uuid3");
+        Assert.assertEquals(3, result.size());
+        Assert.assertEquals(Arrays.asList(1L, 2L, 3L), result);
+    }
+
+    @Test
+    public void testValidateAndGetAffinityGroupIdsWithSpaces() {
+        AffinityGroupVO affinityGroup1 = Mockito.mock(AffinityGroupVO.class);
+        AffinityGroupVO affinityGroup2 = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup1.getId()).thenReturn(1L);
+        Mockito.when(affinityGroup2.getId()).thenReturn(2L);
+        Mockito.when(affinityGroupDao.findByUuid("uuid1")).thenReturn(affinityGroup1);
+        Mockito.when(affinityGroupDao.findByUuid("uuid2")).thenReturn(affinityGroup2);
+
+        List<Long> result = kubernetesServiceHelper.validateAndGetAffinityGroupIds("  uuid1  ,  uuid2  ");
+        Assert.assertEquals(2, result.size());
+        Assert.assertEquals(Arrays.asList(1L, 2L), result);
+    }
+
+    @Test(expected = InvalidParameterValueException.class)
+    public void testValidateAndGetAffinityGroupIdsOneInvalid() {
+        AffinityGroupVO affinityGroup1 = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroupDao.findByUuid("uuid1")).thenReturn(affinityGroup1);
+        Mockito.when(affinityGroupDao.findByUuid("invalid-uuid")).thenReturn(null);
+
+        kubernetesServiceHelper.validateAndGetAffinityGroupIds("uuid1,invalid-uuid");
+    }
+
+    @Test
+    public void testAddNodeTypeAffinityGroupEntry() {
+        Map<String, List<Long>> mapping = new HashMap<>();
+        kubernetesServiceHelper.addNodeTypeAffinityGroupEntry("control", Arrays.asList(1L, 2L), mapping);
+        Assert.assertEquals(1, mapping.size());
+        Assert.assertEquals(Arrays.asList(1L, 2L), mapping.get("CONTROL"));
+    }
+
+    @Test
+    public void testProcessNodeTypeAffinityGroupEntryAndAddToMappingIfValidEmptyEntry() {
+        Map<String, List<Long>> mapping = new HashMap<>();
+        kubernetesServiceHelper.processNodeTypeAffinityGroupEntryAndAddToMappingIfValid(new HashMap<>(), mapping);
+        Assert.assertTrue(mapping.isEmpty());
+    }
+
+    @Test
+    public void testProcessNodeTypeAffinityGroupEntryAndAddToMappingIfValidValidEntry() {
+        AffinityGroupVO affinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup.getId()).thenReturn(100L);
+        Mockito.when(affinityGroupDao.findByUuid("affinity-group-uuid")).thenReturn(affinityGroup);
+
+        Map<String, String> entry = new HashMap<>();
+        entry.put(VmDetailConstants.CKS_NODE_TYPE, "control");
+        entry.put(VmDetailConstants.AFFINITY_GROUP, "affinity-group-uuid");
+
+        Map<String, List<Long>> mapping = new HashMap<>();
+        kubernetesServiceHelper.processNodeTypeAffinityGroupEntryAndAddToMappingIfValid(entry, mapping);
+        Assert.assertEquals(1, mapping.size());
+        Assert.assertEquals(Arrays.asList(100L), mapping.get("CONTROL"));
+    }
+
+    @Test
+    public void testProcessNodeTypeAffinityGroupEntryAndAddToMappingIfValidMultipleUuids() {
+        AffinityGroupVO affinityGroup1 = Mockito.mock(AffinityGroupVO.class);
+        AffinityGroupVO affinityGroup2 = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(affinityGroup1.getId()).thenReturn(1L);
+        Mockito.when(affinityGroup2.getId()).thenReturn(2L);
+        Mockito.when(affinityGroupDao.findByUuid("uuid1")).thenReturn(affinityGroup1);
+        Mockito.when(affinityGroupDao.findByUuid("uuid2")).thenReturn(affinityGroup2);
+
+        Map<String, String> entry = new HashMap<>();
+        entry.put(VmDetailConstants.CKS_NODE_TYPE, "worker");
+        entry.put(VmDetailConstants.AFFINITY_GROUP, "uuid1,uuid2");
+
+        Map<String, List<Long>> mapping = new HashMap<>();
+        kubernetesServiceHelper.processNodeTypeAffinityGroupEntryAndAddToMappingIfValid(entry, mapping);
+        Assert.assertEquals(1, mapping.size());
+        Assert.assertEquals(Arrays.asList(1L, 2L), mapping.get("WORKER"));
+    }
+
+    @Test
+    public void testGetAffinityGroupNodeTypeMapEmptyMap() {
+        Map<String, List<Long>> result = kubernetesServiceHelper.getAffinityGroupNodeTypeMap(null);
+        Assert.assertTrue(result.isEmpty());
+
+        result = kubernetesServiceHelper.getAffinityGroupNodeTypeMap(new HashMap<>());
+        Assert.assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void testGetAffinityGroupNodeTypeMapValidEntries() {
+        AffinityGroupVO controlAffinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(controlAffinityGroup.getId()).thenReturn(100L);
+        Mockito.when(affinityGroupDao.findByUuid("control-affinity-uuid")).thenReturn(controlAffinityGroup);
+
+        AffinityGroupVO workerAffinityGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(workerAffinityGroup.getId()).thenReturn(200L);
+        Mockito.when(affinityGroupDao.findByUuid("worker-affinity-uuid")).thenReturn(workerAffinityGroup);
+
+        Map<String, Map<String, String>> affinityGroupNodeTypeMap = new HashMap<>();
+
+        Map<String, String> controlEntry = new HashMap<>();
+        controlEntry.put(VmDetailConstants.CKS_NODE_TYPE, "control");
+        controlEntry.put(VmDetailConstants.AFFINITY_GROUP, "control-affinity-uuid");
+        affinityGroupNodeTypeMap.put("0", controlEntry);
+
+        Map<String, String> workerEntry = new HashMap<>();
+        workerEntry.put(VmDetailConstants.CKS_NODE_TYPE, "worker");
+        workerEntry.put(VmDetailConstants.AFFINITY_GROUP, "worker-affinity-uuid");
+        affinityGroupNodeTypeMap.put("1", workerEntry);
+
+        Map<String, List<Long>> result = kubernetesServiceHelper.getAffinityGroupNodeTypeMap(affinityGroupNodeTypeMap);
+        Assert.assertEquals(2, result.size());
+        Assert.assertEquals(Arrays.asList(100L), result.get("CONTROL"));
+        Assert.assertEquals(Arrays.asList(200L), result.get("WORKER"));
+    }
+
+    @Test
+    public void testGetAffinityGroupNodeTypeMapMultipleIdsPerNodeType() {
+        AffinityGroupVO ag1 = Mockito.mock(AffinityGroupVO.class);
+        AffinityGroupVO ag2 = Mockito.mock(AffinityGroupVO.class);
+        AffinityGroupVO ag3 = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(ag1.getId()).thenReturn(1L);
+        Mockito.when(ag2.getId()).thenReturn(2L);
+        Mockito.when(ag3.getId()).thenReturn(3L);
+        Mockito.when(affinityGroupDao.findByUuid("ag1")).thenReturn(ag1);
+        Mockito.when(affinityGroupDao.findByUuid("ag2")).thenReturn(ag2);
+        Mockito.when(affinityGroupDao.findByUuid("ag3")).thenReturn(ag3);
+
+        Map<String, Map<String, String>> affinityGroupNodeTypeMap = new HashMap<>();
+
+        Map<String, String> controlEntry = new HashMap<>();
+        controlEntry.put(VmDetailConstants.CKS_NODE_TYPE, "control");
+        controlEntry.put(VmDetailConstants.AFFINITY_GROUP, "ag1,ag2,ag3");
+        affinityGroupNodeTypeMap.put("0", controlEntry);
+
+        Map<String, List<Long>> result = kubernetesServiceHelper.getAffinityGroupNodeTypeMap(affinityGroupNodeTypeMap);
+        Assert.assertEquals(1, result.size());
+        Assert.assertEquals(Arrays.asList(1L, 2L, 3L), result.get("CONTROL"));
+    }
 }

plugins/integrations/kubernetes-service/src/test/java/com/cloud/kubernetes/cluster/actionworkers/KubernetesClusterActionWorkerTest.java

@@ -16,8 +16,14 @@
 // under the License.
 package com.cloud.kubernetes.cluster.actionworkers;
 
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.UUID;
 
+import org.apache.cloudstack.affinity.AffinityGroupVO;
+import org.apache.cloudstack.affinity.dao.AffinityGroupDao;
 import org.apache.cloudstack.api.ApiConstants;
 import org.junit.Assert;
 import org.junit.Before;
@@ -30,6 +36,8 @@ import org.mockito.junit.MockitoJUnitRunner;
 import com.cloud.kubernetes.cluster.KubernetesCluster;
 import com.cloud.kubernetes.cluster.KubernetesClusterDetailsVO;
 import com.cloud.kubernetes.cluster.KubernetesClusterManagerImpl;
+import com.cloud.kubernetes.cluster.KubernetesServiceHelper.KubernetesClusterNodeType;
+import com.cloud.kubernetes.cluster.dao.KubernetesClusterAffinityGroupMapDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterDetailsDao;
 import com.cloud.kubernetes.cluster.dao.KubernetesClusterVmMapDao;
@@ -60,6 +68,12 @@ public class KubernetesClusterActionWorkerTest {
     @Mock
     IPAddressDao ipAddressDao;
 
+    @Mock
+    KubernetesClusterAffinityGroupMapDao kubernetesClusterAffinityGroupMapDao;
+
+    @Mock
+    AffinityGroupDao affinityGroupDao;
+
     KubernetesClusterActionWorker actionWorker = null;
 
     final static Long DEFAULT_ID = 1L;
@@ -70,10 +84,12 @@ public class KubernetesClusterActionWorkerTest {
         kubernetesClusterManager.kubernetesSupportedVersionDao = kubernetesSupportedVersionDao;
         kubernetesClusterManager.kubernetesClusterDetailsDao = kubernetesClusterDetailsDao;
         kubernetesClusterManager.kubernetesClusterVmMapDao = kubernetesClusterVmMapDao;
+        kubernetesClusterManager.kubernetesClusterAffinityGroupMapDao = kubernetesClusterAffinityGroupMapDao;
         KubernetesCluster kubernetesCluster = Mockito.mock(KubernetesCluster.class);
         Mockito.when(kubernetesCluster.getId()).thenReturn(DEFAULT_ID);
         actionWorker = new KubernetesClusterActionWorker(kubernetesCluster, kubernetesClusterManager);
         actionWorker.ipAddressDao = ipAddressDao;
+        actionWorker.affinityGroupDao = affinityGroupDao;
     }
 
     @Test
@@ -130,4 +146,87 @@ public class KubernetesClusterActionWorkerTest {
         IpAddress result = actionWorker.getVpcTierKubernetesPublicIp(mockNetworkForGetVpcTierKubernetesPublicIpTest());
         Assert.assertNotNull(result);
     }
+
+    @Test
+    public void testGetAffinityGroupIdsForNodeTypeReturnsIds() {
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(DEFAULT_ID, "CONTROL"))
+            .thenReturn(Arrays.asList(1L, 2L));
+
+        List<Long> result = actionWorker.getAffinityGroupIdsForNodeType(KubernetesClusterNodeType.CONTROL);
+
+        Assert.assertEquals(2, result.size());
+        Assert.assertTrue(result.containsAll(Arrays.asList(1L, 2L)));
+    }
+
+    @Test
+    public void testGetAffinityGroupIdsForNodeTypeReturnsEmptyList() {
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(DEFAULT_ID, "WORKER"))
+            .thenReturn(Collections.emptyList());
+
+        List<Long> result = actionWorker.getAffinityGroupIdsForNodeType(KubernetesClusterNodeType.WORKER);
+
+        Assert.assertTrue(result.isEmpty());
+    }
+
+    @Test
+    public void testGetMergedAffinityGroupIdsWithExplicitDedication() {
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(DEFAULT_ID, "CONTROL"))
+            .thenReturn(new ArrayList<>(Arrays.asList(1L)));
+
+        AffinityGroupVO explicitGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(explicitGroup.getId()).thenReturn(99L);
+        Mockito.when(affinityGroupDao.findByAccountAndType(Mockito.anyLong(), Mockito.eq("ExplicitDedication")))
+            .thenReturn(explicitGroup);
+
+        List<Long> result = actionWorker.getMergedAffinityGroupIds(KubernetesClusterNodeType.CONTROL, 1L, 1L);
+
+        Assert.assertEquals(2, result.size());
+        Assert.assertTrue(result.contains(1L));
+        Assert.assertTrue(result.contains(99L));
+    }
+
+    @Test
+    public void testGetMergedAffinityGroupIdsNoExplicitDedication() {
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(DEFAULT_ID, "WORKER"))
+            .thenReturn(new ArrayList<>(Arrays.asList(1L, 2L)));
+        Mockito.when(affinityGroupDao.findByAccountAndType(Mockito.anyLong(), Mockito.eq("ExplicitDedication")))
+            .thenReturn(null);
+        Mockito.when(affinityGroupDao.findDomainLevelGroupByType(Mockito.anyLong(), Mockito.eq("ExplicitDedication")))
+            .thenReturn(null);
+
+        List<Long> result = actionWorker.getMergedAffinityGroupIds(KubernetesClusterNodeType.WORKER, 1L, 1L);
+
+        Assert.assertEquals(2, result.size());
+    }
+
+    @Test
+    public void testGetMergedAffinityGroupIdsReturnsNullWhenEmpty() {
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(DEFAULT_ID, "ETCD"))
+            .thenReturn(new ArrayList<>());
+        Mockito.when(affinityGroupDao.findByAccountAndType(Mockito.anyLong(), Mockito.anyString()))
+            .thenReturn(null);
+        Mockito.when(affinityGroupDao.findDomainLevelGroupByType(Mockito.anyLong(), Mockito.anyString()))
+            .thenReturn(null);
+
+        List<Long> result = actionWorker.getMergedAffinityGroupIds(KubernetesClusterNodeType.ETCD, 1L, 1L);
+
+        Assert.assertNull(result);
+    }
+
+    @Test
+    public void testGetMergedAffinityGroupIdsExplicitDedicationAlreadyInList() {
+        Mockito.when(kubernetesClusterAffinityGroupMapDao.listAffinityGroupIdsByClusterIdAndNodeType(DEFAULT_ID, "CONTROL"))
+            .thenReturn(new ArrayList<>(Arrays.asList(99L, 2L)));
+
+        AffinityGroupVO explicitGroup = Mockito.mock(AffinityGroupVO.class);
+        Mockito.when(explicitGroup.getId()).thenReturn(99L);
+        Mockito.when(affinityGroupDao.findByAccountAndType(Mockito.anyLong(), Mockito.eq("ExplicitDedication")))
+            .thenReturn(explicitGroup);
+
+        List<Long> result = actionWorker.getMergedAffinityGroupIds(KubernetesClusterNodeType.CONTROL, 1L, 1L);
+
+        Assert.assertEquals(2, result.size());
+        Assert.assertTrue(result.contains(99L));
+        Assert.assertTrue(result.contains(2L));
+    }
 }

test/integration/component/test_kubernetes_cluster_affinity_groups.py

@@ -0,0 +1,931 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for Kubernetes cluster affinity groups feature"""
+
+import unittest
+from marvin.cloudstackTestCase import cloudstackTestCase
+from marvin.cloudstackAPI import (listInfrastructure,
+                                  listKubernetesSupportedVersions,
+                                  addKubernetesSupportedVersion,
+                                  deleteKubernetesSupportedVersion,
+                                  listKubernetesClusters,
+                                  createKubernetesCluster,
+                                  stopKubernetesCluster,
+                                  startKubernetesCluster,
+                                  deleteKubernetesCluster,
+                                  scaleKubernetesCluster,
+                                  destroyVirtualMachine,
+                                  deleteNetwork)
+from marvin.cloudstackException import CloudstackAPIException
+from marvin.lib.base import (ServiceOffering,
+                             Account,
+                             AffinityGroup,
+                             Configurations)
+from marvin.lib.utils import (cleanup_resources,
+                              random_gen)
+from marvin.lib.common import (get_zone,
+                               get_domain)
+from marvin.sshClient import SshClient
+from nose.plugins.attrib import attr
+from marvin.lib.decoratorGenerators import skipTestIf
+
+import time
+
+_multiprocess_shared_ = True
+
+RAND_SUFFIX = random_gen()
+
+
+class TestKubernetesClusterAffinityGroups(cloudstackTestCase):
+    """
+    Tests for CKS Affinity Groups feature (since 4.23.0)
+
+    This feature allows specifying different affinity groups for each
+    Kubernetes node type (CONTROL, WORKER, ETCD).
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        testClient = super(TestKubernetesClusterAffinityGroups, cls).getClsTestClient()
+        if testClient is None:
+            raise unittest.SkipTest("Marvin test client not available - check marvin configuration")
+        cls.apiclient = testClient.getApiClient()
+        cls.services = testClient.getParsedTestDataConfig()
+        cls.zone = get_zone(cls.apiclient, testClient.getZoneForTests())
+        cls.hypervisor = testClient.getHypervisorInfo()
+        cls.mgtSvrDetails = cls.config.__dict__["mgtSvr"][0].__dict__
+
+        cls.hypervisorNotSupported = False
+        if cls.hypervisor.lower() not in ["kvm", "vmware", "xenserver"]:
+            cls.hypervisorNotSupported = True
+
+        cls.setup_failed = False
+        cls._cleanup = []
+        cls.kubernetes_version_ids = []
+        cls.initial_configuration_cks_enabled = None
+
+        cls.k8s_version = cls.services.get("cks_kubernetes_version_upgrade_to",
+                                           cls.services.get("cks_kubernetes_version_upgrade_from"))
+
+        if cls.hypervisorNotSupported == False:
+            cls.endpoint_url = Configurations.list(cls.apiclient, name="endpoint.url")[0].value
+            if "localhost" in cls.endpoint_url:
+                endpoint_url = "http://%s:%d/client/api" % (cls.mgtSvrDetails["mgtSvrIp"], cls.mgtSvrDetails["port"])
+                cls.debug("Setting endpoint.url to %s" % endpoint_url)
+                Configurations.update(cls.apiclient, "endpoint.url", endpoint_url)
+
+            cls.initial_configuration_cks_enabled = Configurations.list(
+                cls.apiclient, name="cloud.kubernetes.service.enabled")[0].value
+            if cls.initial_configuration_cks_enabled not in ["true", True]:
+                cls.debug("Enabling CloudStack Kubernetes Service plugin and restarting management server")
+                Configurations.update(cls.apiclient, "cloud.kubernetes.service.enabled", "true")
+                cls.restartServer()
+
+            cls.cks_service_offering = None
+
+            if cls.setup_failed == False:
+                try:
+                    cls.kubernetes_version = cls.addKubernetesSupportedVersion(
+                        cls.services["cks_kubernetes_versions"][cls.k8s_version])
+                    cls.kubernetes_version_ids.append(cls.kubernetes_version.id)
+                except Exception as e:
+                    cls.setup_failed = True
+                    cls.debug("Failed to get Kubernetes version ISO in ready state: %s" % e)
+
+            if cls.setup_failed == False:
+                cks_offering_data = cls.services["cks_service_offering"]
+                cks_offering_data["name"] = 'CKS-Instance-' + random_gen()
+                cls.cks_service_offering = ServiceOffering.create(
+                    cls.apiclient,
+                    cks_offering_data
+                )
+                cls._cleanup.append(cls.cks_service_offering)
+
+                cls.domain = get_domain(cls.apiclient)
+                cls.account = Account.create(
+                    cls.apiclient,
+                    cls.services["account"],
+                    domainid=cls.domain.id
+                )
+                cls._cleanup.append(cls.account)
+
+        cls.default_network = None
+
+        return
+
+    @classmethod
+    def tearDownClass(cls):
+        # Delete added Kubernetes supported version
+        for version_id in cls.kubernetes_version_ids:
+            try:
+                cls.deleteKubernetesSupportedVersion(version_id)
+            except Exception as e:
+                cls.debug("Error during cleanup for Kubernetes versions: %s" % e)
+
+        # Restore CKS enabled
+        if cls.initial_configuration_cks_enabled not in ["true", True]:
+            cls.debug("Restoring Kubernetes Service enabled value")
+            Configurations.update(cls.apiclient, "cloud.kubernetes.service.enabled", "false")
+            cls.restartServer()
+
+        super(TestKubernetesClusterAffinityGroups, cls).tearDownClass()
+
+    @classmethod
+    def restartServer(cls):
+        """Restart management server"""
+        cls.debug("Restarting management server")
+        sshClient = SshClient(
+            cls.mgtSvrDetails["mgtSvrIp"],
+            22,
+            cls.mgtSvrDetails["user"],
+            cls.mgtSvrDetails["passwd"]
+        )
+        command = "service cloudstack-management stop"
+        sshClient.execute(command)
+
+        command = "service cloudstack-management start"
+        sshClient.execute(command)
+
+        # Wait for management to come up in 5 mins
+        timeout = time.time() + 300
+        while time.time() < timeout:
+            if cls.isManagementUp() is True:
+                return
+            time.sleep(5)
+        cls.setup_failed = True
+        cls.debug("Management server did not come up, failing")
+        return
+
+    @classmethod
+    def isManagementUp(cls):
+        try:
+            cls.apiclient.listInfrastructure(listInfrastructure.listInfrastructureCmd())
+            return True
+        except Exception:
+            return False
+
+    @classmethod
+    def waitForKubernetesSupportedVersionIsoReadyState(cls, version_id, retries=30, interval=60):
+        """Check if Kubernetes supported version ISO is in Ready state"""
+        while retries > 0:
+            time.sleep(interval)
+            list_versions_response = cls.listKubernetesSupportedVersion(version_id)
+            if not hasattr(list_versions_response, 'isostate') or not list_versions_response or not list_versions_response.isostate:
+                retries = retries - 1
+                continue
+            if 'Ready' == list_versions_response.isostate:
+                return
+            elif 'Failed' == list_versions_response.isostate:
+                raise Exception("Failed to download template: status - %s" % list_versions_response.isostate)
+            retries = retries - 1
+        raise Exception("Kubernetes supported version Ready state timed out")
+
+    @classmethod
+    def listKubernetesSupportedVersion(cls, version_id):
+        listKubernetesSupportedVersionsCmd = listKubernetesSupportedVersions.listKubernetesSupportedVersionsCmd()
+        listKubernetesSupportedVersionsCmd.id = version_id
+        versionResponse = cls.apiclient.listKubernetesSupportedVersions(listKubernetesSupportedVersionsCmd)
+        return versionResponse[0]
+
+    @classmethod
+    def addKubernetesSupportedVersion(cls, version_service):
+        addKubernetesSupportedVersionCmd = addKubernetesSupportedVersion.addKubernetesSupportedVersionCmd()
+        addKubernetesSupportedVersionCmd.semanticversion = version_service["semanticversion"]
+        addKubernetesSupportedVersionCmd.name = 'v' + version_service["semanticversion"] + '-' + random_gen()
+        addKubernetesSupportedVersionCmd.url = version_service["url"]
+        addKubernetesSupportedVersionCmd.mincpunumber = version_service["mincpunumber"]
+        addKubernetesSupportedVersionCmd.minmemory = version_service["minmemory"]
+        kubernetes_version = cls.apiclient.addKubernetesSupportedVersion(addKubernetesSupportedVersionCmd)
+        cls.debug("Waiting for Kubernetes version with ID %s to be ready" % kubernetes_version.id)
+        cls.waitForKubernetesSupportedVersionIsoReadyState(kubernetes_version.id)
+        kubernetes_version = cls.listKubernetesSupportedVersion(kubernetes_version.id)
+        return kubernetes_version
+
+    @classmethod
+    def deleteKubernetesSupportedVersion(cls, version_id):
+        deleteKubernetesSupportedVersionCmd = deleteKubernetesSupportedVersion.deleteKubernetesSupportedVersionCmd()
+        deleteKubernetesSupportedVersionCmd.id = version_id
+        cls.apiclient.deleteKubernetesSupportedVersion(deleteKubernetesSupportedVersionCmd)
+
+    @classmethod
+    def listKubernetesCluster(cls, cluster_id=None, cluster_name=None):
+        listKubernetesClustersCmd = listKubernetesClusters.listKubernetesClustersCmd()
+        listKubernetesClustersCmd.listall = True
+        if cluster_id is not None:
+            listKubernetesClustersCmd.id = cluster_id
+        if cluster_name is not None:
+            listKubernetesClustersCmd.name = cluster_name
+        clusterResponse = cls.apiclient.listKubernetesClusters(listKubernetesClustersCmd)
+        if (cluster_id is not None or cluster_name is not None) and clusterResponse is not None:
+            return clusterResponse[0]
+        return clusterResponse
+
+    @classmethod
+    def deleteKubernetesCluster(cls, cluster_id):
+        deleteKubernetesClusterCmd = deleteKubernetesCluster.deleteKubernetesClusterCmd()
+        deleteKubernetesClusterCmd.id = cluster_id
+        response = cls.apiclient.deleteKubernetesCluster(deleteKubernetesClusterCmd)
+        return response
+
+    @classmethod
+    def stopKubernetesCluster(cls, cluster_id):
+        stopKubernetesClusterCmd = stopKubernetesCluster.stopKubernetesClusterCmd()
+        stopKubernetesClusterCmd.id = cluster_id
+        response = cls.apiclient.stopKubernetesCluster(stopKubernetesClusterCmd)
+        return response
+
+    def setUp(self):
+        self.services = self.testClient.getParsedTestDataConfig()
+        self.apiclient = self.testClient.getApiClient()
+        self.dbclient = self.testClient.getDbConnection()
+        self.cleanup = []
+        self.aff_grp = []
+        return
+
+    def tearDown(self):
+        super(TestKubernetesClusterAffinityGroups, self).tearDown()
+
+    def deleteKubernetesClusterAndVerify(self, cluster_id, verify=True, forced=False):
+        """Delete Kubernetes cluster and check if it is really deleted"""
+        delete_response = {}
+        forceDeleted = False
+        try:
+            delete_response = self.deleteKubernetesCluster(cluster_id)
+        except Exception as e:
+            if forced:
+                cluster = self.listKubernetesCluster(cluster_id)
+                if cluster is not None:
+                    if cluster.state in ['Starting', 'Running', 'Upgrading', 'Scaling']:
+                        self.stopKubernetesCluster(cluster_id)
+                        self.deleteKubernetesCluster(cluster_id)
+                    else:
+                        forceDeleted = True
+                        for cluster_vm in cluster.virtualmachines:
+                            cmd = destroyVirtualMachine.destroyVirtualMachineCmd()
+                            cmd.id = cluster_vm.id
+                            cmd.expunge = True
+                            self.apiclient.destroyVirtualMachine(cmd)
+                        cmd = deleteNetwork.deleteNetworkCmd()
+                        cmd.id = cluster.networkid
+                        cmd.forced = True
+                        self.apiclient.deleteNetwork(cmd)
+                        self.dbclient.execute(
+                            "update kubernetes_cluster set state='Destroyed', removed=now() where uuid = '%s';" % cluster.id)
+            else:
+                raise Exception("Error: Exception during delete cluster : %s" % e)
+
+        if verify and not forceDeleted:
+            self.assertEqual(
+                delete_response.success,
+                True,
+                "Check KubernetesCluster delete response {}, {}".format(delete_response.success, True)
+            )
+
+            db_cluster_removed = \
+            self.dbclient.execute("select removed from kubernetes_cluster where uuid = '%s';" % cluster_id)[0][0]
+
+            self.assertNotEqual(
+                db_cluster_removed,
+                None,
+                "KubernetesCluster not removed in DB, {}".format(db_cluster_removed)
+            )
+
+    def create_aff_grp(self, aff_grp_name=None, aff_grp_type="host anti-affinity"):
+        """Create an affinity group"""
+        if aff_grp_name is None:
+            aff_grp_name = "aff_grp_" + random_gen(size=6)
+
+        aff_grp_data = {
+            "name": aff_grp_name,
+            "type": aff_grp_type
+        }
+        aff_grp = AffinityGroup.create(
+            self.apiclient,
+            aff_grp_data,
+            self.account.name,
+            self.domain.id
+        )
+        self.aff_grp.append(aff_grp)
+        self.cleanup.append(aff_grp)
+        return aff_grp
+
+    def createKubernetesCluster(self, name, version_id, size=1, control_nodes=1, etcd_nodes=0,
+                                control_aff_grp=None, worker_aff_grp=None, etcd_aff_grp=None):
+        """Create a Kubernetes cluster with optional affinity groups for each node type"""
+        createKubernetesClusterCmd = createKubernetesCluster.createKubernetesClusterCmd()
+        createKubernetesClusterCmd.name = name
+        createKubernetesClusterCmd.description = name + "-description"
+        createKubernetesClusterCmd.kubernetesversionid = version_id
+        createKubernetesClusterCmd.size = size
+        createKubernetesClusterCmd.controlnodes = control_nodes
+        createKubernetesClusterCmd.serviceofferingid = self.cks_service_offering.id
+        createKubernetesClusterCmd.zoneid = self.zone.id
+        createKubernetesClusterCmd.noderootdisksize = 10
+        createKubernetesClusterCmd.account = self.account.name
+        createKubernetesClusterCmd.domainid = self.domain.id
+
+        if etcd_nodes > 0:
+            createKubernetesClusterCmd.etcdnodes = etcd_nodes
+
+        # Set affinity groups for node types using the nodeaffinitygroups parameter
+        # Format: list of {node: "<NODE_TYPE>", affinitygroup: "<UUID>"}
+        if control_aff_grp is not None:
+            if not hasattr(createKubernetesClusterCmd, 'nodeaffinitygroups'):
+                createKubernetesClusterCmd.nodeaffinitygroups = []
+            createKubernetesClusterCmd.nodeaffinitygroups.append({
+                "node": "CONTROL",
+                "affinitygroup": control_aff_grp.id
+            })
+        if worker_aff_grp is not None:
+            if not hasattr(createKubernetesClusterCmd, 'nodeaffinitygroups'):
+                createKubernetesClusterCmd.nodeaffinitygroups = []
+            createKubernetesClusterCmd.nodeaffinitygroups.append({
+                "node": "WORKER",
+                "affinitygroup": worker_aff_grp.id
+            })
+        if etcd_aff_grp is not None:
+            if not hasattr(createKubernetesClusterCmd, 'nodeaffinitygroups'):
+                createKubernetesClusterCmd.nodeaffinitygroups = []
+            createKubernetesClusterCmd.nodeaffinitygroups.append({
+                "node": "ETCD",
+                "affinitygroup": etcd_aff_grp.id
+            })
+
+        if self.default_network:
+            createKubernetesClusterCmd.networkid = self.default_network.id
+
+        clusterResponse = self.apiclient.createKubernetesCluster(createKubernetesClusterCmd)
+        return clusterResponse
+
+    def startKubernetesCluster(self, cluster_id):
+        startKubernetesClusterCmd = startKubernetesCluster.startKubernetesClusterCmd()
+        startKubernetesClusterCmd.id = cluster_id
+        response = self.apiclient.startKubernetesCluster(startKubernetesClusterCmd)
+        return response
+
+    def scaleKubernetesCluster(self, cluster_id, size):
+        scaleKubernetesClusterCmd = scaleKubernetesCluster.scaleKubernetesClusterCmd()
+        scaleKubernetesClusterCmd.id = cluster_id
+        scaleKubernetesClusterCmd.size = size
+        response = self.apiclient.scaleKubernetesCluster(scaleKubernetesClusterCmd)
+        return response
+
+    def verifyKubernetesClusterState(self, cluster_response, state):
+        """Check if Kubernetes cluster state matches expected state"""
+        self.assertEqual(
+            cluster_response.state,
+            state,
+            "Check KubernetesCluster state {}, expected {}".format(cluster_response.state, state)
+        )
+
+    def verifyKubernetesClusterAffinityGroups(self, cluster, control_aff_grp=None,
+                                               worker_aff_grp=None, etcd_aff_grp=None):
+        """Verify affinity groups are correctly assigned to the cluster"""
+        if control_aff_grp is not None:
+            self.assertEqual(
+                cluster.controlnodeaffinitygroupid,
+                control_aff_grp.id,
+                "Control node affinity group ID mismatch. Expected: {}, Got: {}".format(
+                    control_aff_grp.id, cluster.controlnodeaffinitygroupid)
+            )
+            self.assertEqual(
+                cluster.controlnodeaffinitygroupname,
+                control_aff_grp.name,
+                "Control node affinity group name mismatch. Expected: {}, Got: {}".format(
+                    control_aff_grp.name, cluster.controlnodeaffinitygroupname)
+            )
+        else:
+            self.assertTrue(
+                not hasattr(cluster, 'controlnodeaffinitygroupid') or cluster.controlnodeaffinitygroupid is None,
+                "Control node affinity group should be None"
+            )
+
+        if worker_aff_grp is not None:
+            self.assertEqual(
+                cluster.workernodeaffinitygroupid,
+                worker_aff_grp.id,
+                "Worker node affinity group ID mismatch. Expected: {}, Got: {}".format(
+                    worker_aff_grp.id, cluster.workernodeaffinitygroupid)
+            )
+            self.assertEqual(
+                cluster.workernodeaffinitygroupname,
+                worker_aff_grp.name,
+                "Worker node affinity group name mismatch. Expected: {}, Got: {}".format(
+                    worker_aff_grp.name, cluster.workernodeaffinitygroupname)
+            )
+        else:
+            self.assertTrue(
+                not hasattr(cluster, 'workernodeaffinitygroupid') or cluster.workernodeaffinitygroupid is None,
+                "Worker node affinity group should be None"
+            )
+
+        if etcd_aff_grp is not None:
+            self.assertEqual(
+                cluster.etcdnodeaffinitygroupid,
+                etcd_aff_grp.id,
+                "ETCD node affinity group ID mismatch. Expected: {}, Got: {}".format(
+                    etcd_aff_grp.id, cluster.etcdnodeaffinitygroupid)
+            )
+            self.assertEqual(
+                cluster.etcdnodeaffinitygroupname,
+                etcd_aff_grp.name,
+                "ETCD node affinity group name mismatch. Expected: {}, Got: {}".format(
+                    etcd_aff_grp.name, cluster.etcdnodeaffinitygroupname)
+            )
+        else:
+            self.assertTrue(
+                not hasattr(cluster, 'etcdnodeaffinitygroupid') or cluster.etcdnodeaffinitygroupid is None,
+                "ETCD node affinity group should be None"
+            )
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_01_create_cluster_with_control_node_affinity_group(self):
+        """Test creating a Kubernetes cluster with affinity group for control nodes only
+
+        # Validate the following:
+        # 1. Create an affinity group
+        # 2. Create a Kubernetes cluster with the affinity group for control nodes
+        # 3. Verify cluster is created successfully and affinity group is assigned
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating affinity group for control nodes")
+        control_aff_grp = self.create_aff_grp(aff_grp_name="control-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with control node affinity group")
+        cluster_name = "cks-aff-grp-control-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                control_aff_grp=control_aff_grp
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+            self.verifyKubernetesClusterAffinityGroups(
+                cluster,
+                control_aff_grp=control_aff_grp,
+                worker_aff_grp=None,
+                etcd_aff_grp=None
+            )
+
+            self.debug("Kubernetes cluster with control node affinity group created successfully")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_02_create_cluster_with_worker_node_affinity_group(self):
+        """Test creating a Kubernetes cluster with affinity group for worker nodes only
+
+        # Validate the following:
+        # 1. Create an affinity group
+        # 2. Create a Kubernetes cluster with the affinity group for worker nodes
+        # 3. Verify cluster is created successfully and affinity group is assigned
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating affinity group for worker nodes")
+        worker_aff_grp = self.create_aff_grp(aff_grp_name="worker-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with worker node affinity group")
+        cluster_name = "cks-aff-grp-worker-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                worker_aff_grp=worker_aff_grp
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+            self.verifyKubernetesClusterAffinityGroups(
+                cluster,
+                control_aff_grp=None,
+                worker_aff_grp=worker_aff_grp,
+                etcd_aff_grp=None
+            )
+
+            self.debug("Kubernetes cluster with worker node affinity group created successfully")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_03_create_cluster_with_all_node_type_affinity_groups(self):
+        """Test creating a Kubernetes cluster with different affinity groups for all node types
+
+        # Validate the following:
+        # 1. Create separate affinity groups for control, worker, and etcd nodes
+        # 2. Create a Kubernetes cluster with affinity groups for all node types
+        # 3. Verify cluster is created successfully and all affinity groups are assigned
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating affinity groups for all node types")
+        control_aff_grp = self.create_aff_grp(aff_grp_name="control-aff-grp-" + random_gen())
+        worker_aff_grp = self.create_aff_grp(aff_grp_name="worker-aff-grp-" + random_gen())
+        etcd_aff_grp = self.create_aff_grp(aff_grp_name="etcd-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with affinity groups for all node types")
+        cluster_name = "cks-aff-grp-all-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                etcd_nodes=1,
+                control_aff_grp=control_aff_grp,
+                worker_aff_grp=worker_aff_grp,
+                etcd_aff_grp=etcd_aff_grp
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+            self.verifyKubernetesClusterAffinityGroups(
+                cluster,
+                control_aff_grp=control_aff_grp,
+                worker_aff_grp=worker_aff_grp,
+                etcd_aff_grp=etcd_aff_grp
+            )
+
+            self.debug("Kubernetes cluster with all node type affinity groups created successfully")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_04_create_cluster_with_same_affinity_group_all_node_types(self):
+        """Test creating a Kubernetes cluster with the same affinity group for all node types
+
+        # Validate the following:
+        # 1. Create a single affinity group
+        # 2. Create a Kubernetes cluster using the same affinity group for all node types
+        # 3. Verify cluster is created successfully
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating a single affinity group for all node types")
+        shared_aff_grp = self.create_aff_grp(aff_grp_name="shared-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with same affinity group for all node types")
+        cluster_name = "cks-aff-grp-shared-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                control_aff_grp=shared_aff_grp,
+                worker_aff_grp=shared_aff_grp
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+            self.verifyKubernetesClusterAffinityGroups(
+                cluster,
+                control_aff_grp=shared_aff_grp,
+                worker_aff_grp=shared_aff_grp,
+                etcd_aff_grp=None
+            )
+
+            self.debug("Kubernetes cluster with shared affinity group created successfully")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_05_scale_cluster_respects_worker_affinity_group(self):
+        """Test that scaling a cluster respects the worker node affinity group
+
+        # Validate the following:
+        # 1. Create a cluster with worker node affinity group
+        # 2. Scale up the cluster
+        # 3. Verify affinity group assignments are preserved after scaling
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating affinity group for worker nodes")
+        worker_aff_grp = self.create_aff_grp(aff_grp_name="worker-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with worker node affinity group")
+        cluster_name = "cks-aff-grp-scale-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                worker_aff_grp=worker_aff_grp
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+
+            self.debug("Scaling up Kubernetes cluster from 1 to 2 worker nodes")
+            cluster = self.scaleKubernetesCluster(cluster.id, 2)
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+            self.assertEqual(
+                cluster.size,
+                2,
+                "Cluster size should be 2 after scaling up, got {}".format(cluster.size)
+            )
+
+            # Verify affinity group is still assigned after scaling
+            self.verifyKubernetesClusterAffinityGroups(
+                cluster,
+                control_aff_grp=None,
+                worker_aff_grp=worker_aff_grp,
+                etcd_aff_grp=None
+            )
+
+            self.debug("Kubernetes cluster scaled successfully with affinity group preserved")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_06_stop_start_cluster_preserves_affinity_groups(self):
+        """Test that stopping and starting a cluster preserves affinity groups
+
+        # Validate the following:
+        # 1. Create a cluster with affinity groups
+        # 2. Stop the cluster
+        # 3. Start the cluster
+        # 4. Verify affinity groups are preserved
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating affinity groups")
+        control_aff_grp = self.create_aff_grp(aff_grp_name="control-aff-grp-" + random_gen())
+        worker_aff_grp = self.create_aff_grp(aff_grp_name="worker-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with affinity groups")
+        cluster_name = "cks-aff-grp-lifecycle-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                control_aff_grp=control_aff_grp,
+                worker_aff_grp=worker_aff_grp
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+
+            self.debug("Stopping Kubernetes cluster")
+            self.stopKubernetesCluster(cluster.id)
+
+            cluster = self.listKubernetesCluster(cluster.id)
+            self.assertEqual(
+                cluster.state,
+                'Stopped',
+                "Cluster should be in Stopped state, got {}".format(cluster.state)
+            )
+
+            self.debug("Starting Kubernetes cluster")
+            cluster = self.startKubernetesCluster(cluster.id)
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+
+            # Verify affinity groups are preserved after stop/start
+            self.verifyKubernetesClusterAffinityGroups(
+                cluster,
+                control_aff_grp=control_aff_grp,
+                worker_aff_grp=worker_aff_grp,
+                etcd_aff_grp=None
+            )
+
+            self.debug("Kubernetes cluster stop/start completed with affinity groups preserved")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_07_create_cluster_with_invalid_affinity_group_id(self):
+        """Test creating a cluster with an invalid affinity group ID fails
+
+        # Validate the following:
+        # 1. Attempt to create a cluster with a non-existent affinity group ID
+        # 2. Verify the operation fails with appropriate error
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating Kubernetes cluster with invalid affinity group ID")
+        cluster_name = "cks-aff-grp-invalid-" + random_gen()
+
+        # Create a fake affinity group object with invalid ID
+        class FakeAffinityGroup:
+            def __init__(self):
+                self.id = "invalid-uuid-12345"
+                self.name = "fake-group"
+
+        fake_aff_grp = FakeAffinityGroup()
+
+        with self.assertRaises(Exception) as context:
+            self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                control_aff_grp=fake_aff_grp
+            )
+
+        self.debug("Expected error when creating cluster with invalid affinity group: %s" % context.exception)
+
+        # Clean up any partially created cluster
+        cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+        if cluster is not None:
+            self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_08_cluster_response_includes_affinity_group_details(self):
+        """Test that cluster list response includes affinity group details
+
+        # Validate the following:
+        # 1. Create a cluster with affinity groups
+        # 2. List the cluster
+        # 3. Verify response includes affinity group IDs and names
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating affinity groups")
+        control_aff_grp = self.create_aff_grp(aff_grp_name="control-aff-grp-" + random_gen())
+        worker_aff_grp = self.create_aff_grp(aff_grp_name="worker-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with affinity groups")
+        cluster_name = "cks-aff-grp-response-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1,
+                control_aff_grp=control_aff_grp,
+                worker_aff_grp=worker_aff_grp
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+
+            # List the cluster and verify response
+            listed_cluster = self.listKubernetesCluster(cluster.id)
+
+            self.assertIsNotNone(listed_cluster, "Cluster should be listed")
+
+            # Verify control node affinity group in response
+            self.assertTrue(
+                hasattr(listed_cluster, 'controlnodeaffinitygroupid'),
+                "Response should include controlnodeaffinitygroupid"
+            )
+            self.assertTrue(
+                hasattr(listed_cluster, 'controlnodeaffinitygroupname'),
+                "Response should include controlnodeaffinitygroupname"
+            )
+
+            # Verify worker node affinity group in response
+            self.assertTrue(
+                hasattr(listed_cluster, 'workernodeaffinitygroupid'),
+                "Response should include workernodeaffinitygroupid"
+            )
+            self.assertTrue(
+                hasattr(listed_cluster, 'workernodeaffinitygroupname'),
+                "Response should include workernodeaffinitygroupname"
+            )
+
+            # Verify the values match
+            self.verifyKubernetesClusterAffinityGroups(
+                listed_cluster,
+                control_aff_grp=control_aff_grp,
+                worker_aff_grp=worker_aff_grp,
+                etcd_aff_grp=None
+            )
+
+            self.debug("Cluster response includes correct affinity group details")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_09_create_cluster_without_affinity_groups(self):
+        """Test creating a cluster without any affinity groups
+
+        # Validate the following:
+        # 1. Create a cluster without specifying any affinity groups
+        # 2. Verify cluster is created and affinity group fields are null/empty
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating Kubernetes cluster without affinity groups")
+        cluster_name = "cks-no-aff-grp-" + random_gen()
+        try:
+            cluster = self.createKubernetesCluster(
+                cluster_name,
+                self.kubernetes_version.id,
+                size=1,
+                control_nodes=1
+            )
+
+            self.verifyKubernetesClusterState(cluster, 'Running')
+
+            # Verify no affinity groups are assigned
+            self.verifyKubernetesClusterAffinityGroups(
+                cluster,
+                control_aff_grp=None,
+                worker_aff_grp=None,
+                etcd_aff_grp=None
+            )
+
+            self.debug("Kubernetes cluster created successfully without affinity groups")
+        finally:
+            cluster = self.listKubernetesCluster(cluster_name=cluster_name)
+            if cluster is not None:
+                self.deleteKubernetesClusterAndVerify(cluster.id, False, True)
+
+    @attr(tags=["advanced", "smoke"], required_hardware="true")
+    @skipTestIf("hypervisorNotSupported")
+    def test_10_delete_cluster_with_affinity_groups(self):
+        """Test that deleting a cluster with affinity groups works correctly
+
+        # Validate the following:
+        # 1. Create a cluster with affinity groups
+        # 2. Delete the cluster
+        # 3. Verify cluster is deleted and affinity groups still exist
+        """
+        if self.setup_failed:
+            self.fail("Setup incomplete")
+
+        self.debug("Creating affinity groups")
+        control_aff_grp = self.create_aff_grp(aff_grp_name="control-aff-grp-" + random_gen())
+        worker_aff_grp = self.create_aff_grp(aff_grp_name="worker-aff-grp-" + random_gen())
+
+        self.debug("Creating Kubernetes cluster with affinity groups")
+        cluster_name = "cks-aff-grp-delete-" + random_gen()
+        cluster = self.createKubernetesCluster(
+            cluster_name,
+            self.kubernetes_version.id,
+            size=1,
+            control_nodes=1,
+            control_aff_grp=control_aff_grp,
+            worker_aff_grp=worker_aff_grp
+        )
+
+        self.verifyKubernetesClusterState(cluster, 'Running')
+        cluster_id = cluster.id
+
+        self.debug("Deleting Kubernetes cluster")
+        self.deleteKubernetesClusterAndVerify(cluster_id)
+
+        # Verify cluster is deleted
+        deleted_cluster = self.listKubernetesCluster(cluster_id)
+        self.assertIsNone(deleted_cluster, "Cluster should be deleted")
+
+        # Verify affinity groups still exist
+        control_aff_grp_list = AffinityGroup.list(self.apiclient, id=control_aff_grp.id)
+        self.assertIsNotNone(control_aff_grp_list, "Control affinity group should still exist")
+
+        worker_aff_grp_list = AffinityGroup.list(self.apiclient, id=worker_aff_grp.id)
+        self.assertIsNotNone(worker_aff_grp_list, "Worker affinity group should still exist")
+
+        self.debug("Kubernetes cluster deleted successfully, affinity groups preserved")

ui/public/locales/en.json

@@ -556,6 +556,9 @@
 "label.cks.cluster.size": "Cluster size (Worker nodes)",
 "label.cks.cluster.worker.nodes.offeringid": "Service Offering for Worker Nodes",
 "label.cks.cluster.worker.nodes.templateid": "Template for Worker Nodes",
+"label.cks.cluster.control.nodes.affinitygroupid": "Affinity Groups for Control Nodes",
+"label.cks.cluster.worker.nodes.affinitygroupid": "Affinity Groups for Worker Nodes",
+"label.cks.cluster.etcd.nodes.affinitygroupid": "Affinity Groups for ETCD Nodes",
 "label.cleanup": "Clean up",
 "label.clear": "Clear",
 "label.clear.all": "Clear all",
@@ -2886,6 +2889,9 @@
 "label.edgecluster": "Edge Cluster",
 "label.encryption": "Encryption",
 "label.etcdnodes": "Number of etcd nodes",
+"label.controlaffinitygroupnames": "Control Affinity Groups",
+"label.workeraffinitygroupnames": "Worker Affinity Groups",
+"label.etcdaffinitygroupnames": "ETCD Affinity Groups",
 "label.versioning": "Versioning",
 "label.objectlocking": "Object Lock",
 "label.bucket.policy": "Bucket Policy",

ui/src/config/section/compute.js

@@ -574,7 +574,7 @@ export default {
         const filters = ['cloud.managed', 'external.managed']
         return filters
       },
-      details: ['name', 'description', 'zonename', 'kubernetesversionname', 'autoscalingenabled', 'csienabled', 'minsize', 'maxsize', 'size', 'controlnodes', 'etcdnodes', 'cpunumber', 'memory', 'keypair', 'cniconfigname', 'associatednetworkname', 'account', 'domain', 'zonename', 'clustertype', 'created'],
+      details: ['name', 'description', 'zonename', 'kubernetesversionname', 'autoscalingenabled', 'csienabled', 'minsize', 'maxsize', 'size', 'controlnodes', 'controlaffinitygroupnames', 'etcdnodes', 'etcdaffinitygroupnames', 'workeraffinitygroupnames', 'cpunumber', 'memory', 'keypair', 'cniconfigname', 'associatednetworkname', 'account', 'domain', 'zonename', 'clustertype', 'created'],
       tabs: [
         {
           name: 'k8s',

ui/src/views/compute/CreateKubernetesCluster.vue

@@ -316,7 +316,7 @@
             </a-select-option>
           </a-select>
         </a-form-item>
-        <a-form-item v-if="form.advancedmode && form?.etcdnodes > 0" name="etcdtemplateid" ref="etcdtemplateid">
+        <a-form-item v-if="form.advancedmode && form.etcdnodes && form.etcdnodes > 0" name="etcdtemplateid" ref="etcdtemplateid">
           <template #label>
             <tooltip-label :title="$t('label.cks.cluster.etcd.nodes.templateid')" :tooltip="$t('label.cks.cluster.etcd.nodes.templateid')"/>
           </template>
@@ -335,6 +335,57 @@
             </a-select-option>
           </a-select>
         </a-form-item>
+        <a-form-item v-if="form.advancedmode" name="controlaffinitygroupids" ref="controlaffinitygroupids">
+          <template #label>
+            <tooltip-label :title="$t('label.cks.cluster.control.nodes.affinitygroupid')" :tooltip="$t('label.cks.cluster.control.nodes.affinitygroupid')"/>
+          </template>
+          <a-select
+            v-model:value="controlAffinityGroups"
+            mode="multiple"
+            showSearch
+            optionFilterProp="label"
+            :filterOption="(input, option) => option.label.toLowerCase().indexOf(input.toLowerCase()) >= 0"
+            :loading="affinityGroupLoading"
+            :placeholder="$t('label.cks.cluster.control.nodes.affinitygroupid')">
+            <a-select-option v-for="opt in affinityGroups" :key="opt.id" :label="opt.name">
+              {{ opt.name }}
+            </a-select-option>
+          </a-select>
+        </a-form-item>
+        <a-form-item v-if="form.advancedmode" name="workeraffinitygroupids" ref="workeraffinitygroupids">
+          <template #label>
+            <tooltip-label :title="$t('label.cks.cluster.worker.nodes.affinitygroupid')" :tooltip="$t('label.cks.cluster.worker.nodes.affinitygroupid')"/>
+          </template>
+          <a-select
+            v-model:value="workerAffinityGroups"
+            mode="multiple"
+            showSearch
+            optionFilterProp="label"
+            :filterOption="(input, option) => option.label.toLowerCase().indexOf(input.toLowerCase()) >= 0"
+            :loading="affinityGroupLoading"
+            :placeholder="$t('label.cks.cluster.worker.nodes.affinitygroupid')">
+            <a-select-option v-for="opt in affinityGroups" :key="opt.id" :label="opt.name">
+              {{ opt.name }}
+            </a-select-option>
+          </a-select>
+        </a-form-item>
+        <a-form-item v-if="form.advancedmode && form.etcdnodes && form.etcdnodes > 0" name="etcdaffinitygroupids" ref="etcdaffinitygroupids">
+          <template #label>
+            <tooltip-label :title="$t('label.cks.cluster.etcd.nodes.affinitygroupid')" :tooltip="$t('label.cks.cluster.etcd.nodes.affinitygroupid')"/>
+          </template>
+          <a-select
+            v-model:value="etcdAffinityGroups"
+            mode="multiple"
+            showSearch
+            optionFilterProp="label"
+            :filterOption="(input, option) => option.label.toLowerCase().indexOf(input.toLowerCase()) >= 0"
+            :loading="affinityGroupLoading"
+            :placeholder="$t('label.cks.cluster.etcd.nodes.affinitygroupid')">
+            <a-select-option v-for="opt in affinityGroups" :key="opt.id" :label="opt.name">
+              {{ opt.name }}
+            </a-select-option>
+          </a-select>
+        </a-form-item>
         <a-form-item v-if="form.advancedmode && isASNumberRequired() && !form.networkid" name="asnumber" ref="asnumber">
           <template #label>
               <tooltip-label :title="$t('label.asnumber')" :tooltip="apiParams.asnumber.description"/>
@@ -490,7 +541,12 @@ export default {
       cksNetworkOffering: null,
       asNumbersZone: [],
       asNumberLoading: false,
-      selectedAsNumber: 0
+      selectedAsNumber: 0,
+      affinityGroups: [],
+      affinityGroupLoading: false,
+      controlAffinityGroups: [],
+      workerAffinityGroups: [],
+      etcdAffinityGroups: []
     }
   },
   beforeCreate () {
@@ -560,6 +616,22 @@ export default {
       this.fetchCKSNetworkOfferingName()
       this.fetchCniConfigurations()
     },
+    fetchAffinityGroups () {
+      this.affinityGroups = []
+      const params = {}
+      if (!this.isObjectEmpty(this.selectedZone)) {
+        params.zoneid = this.selectedZone.id
+      }
+      this.affinityGroupLoading = true
+      getAPI('listAffinityGroups', params).then(json => {
+        const groups = json.listaffinitygroupsresponse.affinitygroup
+        if (this.arrayHasItems(groups)) {
+          this.affinityGroups = groups
+        }
+      }).finally(() => {
+        this.affinityGroupLoading = false
+      })
+    },
     isValidValueForKey (obj, key) {
       return key in obj && obj[key] != null
     },
@@ -600,6 +672,7 @@ export default {
       this.fetchNetworkData()
       this.fetchZoneHypervisors()
       this.fetchZoneASNumbers()
+      this.fetchAffinityGroups()
     },
     handleASNumberChange (selectedIndex) {
       this.selectedAsNumber = this.asNumbersZone[selectedIndex].asnumber
@@ -881,6 +954,21 @@ export default {
             advancedTemplates++
           }
         }
+        if (values.advancedmode) {
+          let affinityIndex = 0
+          const addAffinityGroups = (nodeType, groups) => {
+            if (groups && groups.length > 0) {
+              params[`nodeaffinitygroups[${affinityIndex}].node`] = nodeType
+              params[`nodeaffinitygroups[${affinityIndex}].affinitygroup`] = groups.join(',')
+              affinityIndex++
+            }
+          }
+          addAffinityGroups('control', this.controlAffinityGroups)
+          addAffinityGroups('worker', this.workerAffinityGroups)
+          if (values.etcdnodes > 0) {
+            addAffinityGroups('etcd', this.etcdAffinityGroups)
+          }
+        }
         if (this.isValidValueForKey(values, 'noderootdisksize') && values.noderootdisksize > 0) {
           params.noderootdisksize = values.noderootdisksize
         }