Retire 1 node per clusterid

2 files changed, 6 insertions, 4 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirer.java
index 0372a961f88..7290bc2b9e6 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirer.java
@@ -117,11 +117,11 @@ public class NodeRetirer extends Maintainer {
         // Get all the nodes that we could retire along with their deployments
         Map<Deployment, Set<Node>> nodesToRetireByDeployment = new HashMap<>();
         for (ApplicationId applicationId : activeApplications) {
-            Map<ClusterSpec, Set<Node>> nodesByCluster = getNodesBelongingToApplication(allNodes, applicationId).stream()
+            Map<ClusterSpec.Id, Set<Node>> nodesByCluster = getNodesBelongingToApplication(allNodes, applicationId).stream()
                     .collect(Collectors.groupingBy(
-                            node -> node.allocation().get().membership().cluster(),
+                            node -> node.allocation().get().membership().cluster().id(),
                             Collectors.toSet()));
-            Map<ClusterSpec, Set<Node>> retireableNodesByCluster = nodesByCluster.entrySet().stream()
+            Map<ClusterSpec.Id, Set<Node>> retireableNodesByCluster = nodesByCluster.entrySet().stream()
                     .collect(Collectors.toMap(
                             Map.Entry::getKey,
                             entry -> filterRetireableNodes(entry.getValue())));
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirerTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirerTester.java
index 6949f6da301..81535fa8d1b 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirerTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRetirerTester.java
@@ -101,7 +101,9 @@ public class NodeRetirerTester {
             Flavor flavor = flavors.get(flavorIds[i]);
             ClusterSpec cluster = ClusterSpec.request(ClusterSpec.Type.container, ClusterSpec.Id.from("cluster-" + i), Version.fromString("6.99"));
             Capacity capacity = Capacity.fromNodeCount(numNodes[i], flavor.name());
-            clusterContexts.add(new MockDeployer.ClusterContext(applicationId, cluster, capacity, 1));
+            // If the number of node the app wants is divisible by 2, make it into 2 groups, otherwise as 1
+            int numGroups = numNodes[i] % 2 == 0 ? 2 : 1;
+            clusterContexts.add(new MockDeployer.ClusterContext(applicationId, cluster, capacity, numGroups));
         }
 
         apps.put(applicationId, new MockDeployer.ApplicationContext(applicationId, clusterContexts));