1 files changed, 111 insertions, 5 deletions
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/NodeStateChangeCheckerTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/NodeStateChangeCheckerTest.java
index 45ca07c88e4..08265248959 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/NodeStateChangeCheckerTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/NodeStateChangeCheckerTest.java
@@ -114,7 +114,7 @@ public class NodeStateChangeCheckerTest {
     }
 
     @Test
-    void testCanUpgradeForce() {
+    void testCanUpgradeWithForce() {
         var nodeStateChangeChecker = createChangeChecker(createCluster(1));
         NodeState newState = new NodeState(STORAGE, INITIALIZING);
         Result result = nodeStateChangeChecker.evaluateTransition(
@@ -152,7 +152,7 @@ public class NodeStateChangeCheckerTest {
     void testSafeMaintenanceDisallowedWhenOtherStorageNodeInFlatClusterIsSuspended() {
         // Nodes 0-3, storage node 0 being in maintenance with "Orchestrator" description.
         ContentCluster cluster = createCluster(4);
-        cluster.clusterInfo().getStorageNodeInfo(0).setWantedState(new NodeState(STORAGE, MAINTENANCE).setDescription("Orchestrator"));
+        setStorageNodeWantedStateToMaintenance(cluster, 0);
         var nodeStateChangeChecker = createChangeChecker(cluster);
         ClusterState clusterStateWith0InMaintenance = clusterState(String.format(
                 "version:%d distributor:4 storage:4 .0.s:m",
@@ -168,6 +168,101 @@ public class NodeStateChangeCheckerTest {
     }
 
     @Test
+    void testMaintenanceAllowedFor2Of4Groups() {
+        // 4 groups with 1 node in each group
+        Collection<ConfiguredNode> nodes = createNodes(4);
+        StorDistributionConfig config = createDistributionConfig(4, 4);
+
+        int maxNumberOfGroupsAllowedToBeDown = 2;
+        var cluster = new ContentCluster("Clustername", nodes, new Distribution(config), maxNumberOfGroupsAllowedToBeDown);
+        setAllNodesUp(cluster, HostInfo.createHostInfo(createDistributorHostInfo(4, 5, 6)));
+        var nodeStateChangeChecker = createChangeChecker(cluster);
+
+        // All nodes up, set a storage node in group 0 to maintenance
+        {
+            int nodeIndex = 0;
+            checkSettingToMaintenanceIsAllowed(nodeIndex, nodeStateChangeChecker, defaultAllUpClusterState());
+            setStorageNodeWantedStateToMaintenance(cluster, nodeIndex);
+        }
+
+        // Node in group 0 in maintenance, set storage node in group 1 to maintenance
+        {
+            ClusterState clusterState = clusterState(String.format("version:%d distributor:4 .0.s:d storage:4 .0.s:m", currentClusterStateVersion));
+            int nodeIndex = 1;
+            checkSettingToMaintenanceIsAllowed(nodeIndex, nodeStateChangeChecker, clusterState);
+            setStorageNodeWantedStateToMaintenance(cluster, nodeIndex);
+        }
+
+        // Nodes in group 0 and 1 in maintenance, try to set storage node in group 2 to maintenance, should fail
+        {
+            ClusterState clusterState = clusterState(String.format("version:%d distributor:4 storage:4 .0.s:m .1.s:m", currentClusterStateVersion));
+            int nodeIndex = 2;
+            Node node = new Node(STORAGE, nodeIndex);
+            Result result = nodeStateChangeChecker.evaluateTransition(node, clusterState, SAFE, UP_NODE_STATE, MAINTENANCE_NODE_STATE);
+            assertFalse(result.settingWantedStateIsAllowed(), result.toString());
+            assertFalse(result.wantedStateAlreadySet());
+            assertEquals("At most nodes in 2 groups can have wanted state", result.getReason());
+        }
+
+    }
+
+    @Test
+    void testMaintenanceAllowedFor2Of4Groups8Nodes() {
+        // 4 groups with 2 node in each group
+        Collection<ConfiguredNode> nodes = createNodes(8);
+        StorDistributionConfig config = createDistributionConfig(8, 4);
+
+        int maxNumberOfGroupsAllowedToBeDown = 2;
+        var cluster = new ContentCluster("Clustername", nodes, new Distribution(config), maxNumberOfGroupsAllowedToBeDown);
+        setAllNodesUp(cluster, HostInfo.createHostInfo(createDistributorHostInfo(4, 5, 6)));
+        var nodeStateChangeChecker = createChangeChecker(cluster);
+
+        // All nodes up, set a storage node in group 0 to maintenance
+        {
+            ClusterState clusterState = defaultAllUpClusterState(8);
+            int nodeIndex = 0;
+            checkSettingToMaintenanceIsAllowed(nodeIndex, nodeStateChangeChecker, clusterState);
+            setStorageNodeWantedStateToMaintenance(cluster, nodeIndex);
+        }
+
+        // 1 Node in group 0 in maintenance, try to set node 1 in group 0 to maintenance
+        {
+            ClusterState clusterState = clusterState(String.format("version:%d distributor:8 .0.s:d storage:8 .0.s:m", currentClusterStateVersion));
+            int nodeIndex = 1;
+            checkSettingToMaintenanceIsAllowed(nodeIndex, nodeStateChangeChecker, clusterState);
+            setStorageNodeWantedStateToMaintenance(cluster, nodeIndex);
+        }
+
+        // Nodes in group 0 in maintenance, try to set storage node 2 in group 1 to maintenance
+        {
+            ClusterState clusterState = clusterState(String.format("version:%d distributor:8 storage:8 .0.s:m .1.s:m", currentClusterStateVersion));
+            int nodeIndex = 2;
+            checkSettingToMaintenanceIsAllowed(nodeIndex, nodeStateChangeChecker, clusterState);
+            setStorageNodeWantedStateToMaintenance(cluster, nodeIndex);
+        }
+
+        // 2 nodes in group 0 and 1 in group 1 in maintenance, try to set storage node 4 in group 2 to maintenance, should fail (different group)
+        {
+            ClusterState clusterState = clusterState(String.format("version:%d distributor:8 storage:8 .0.s:m .1.s:m .2.s:m", currentClusterStateVersion));
+            int nodeIndex = 4;
+            Node node = new Node(STORAGE, nodeIndex);
+            Result result = nodeStateChangeChecker.evaluateTransition(node, clusterState, SAFE, UP_NODE_STATE, MAINTENANCE_NODE_STATE);
+            assertFalse(result.settingWantedStateIsAllowed(), result.toString());
+            assertFalse(result.wantedStateAlreadySet());
+            assertEquals("At most nodes in 2 groups can have wanted state", result.getReason());
+        }
+
+        // 2 nodes in group 0 and 1 in group 1 in maintenance, try to set storage node 3 in group 1 to maintenance
+        {
+            ClusterState clusterState = clusterState(String.format("version:%d distributor:8 storage:8 .0.s:m .1.s:m .2.s:m", currentClusterStateVersion));
+            int nodeIndex = 3;
+            checkSettingToMaintenanceIsAllowed(nodeIndex, nodeStateChangeChecker, clusterState);
+            setStorageNodeWantedStateToMaintenance(cluster, nodeIndex);
+        }
+
+    }
+
+    @Test
     void testSafeMaintenanceDisallowedWhenOtherDistributorInFlatClusterIsSuspended() {
         // Nodes 0-3, distributor 0 being down with "Orchestrator" description.
         ContentCluster cluster = createCluster(4);
@@ -439,10 +534,9 @@ public class NodeStateChangeCheckerTest {
         NodeStateChangeChecker nodeStateChangeChecker = createChangeChecker(cluster);
 
         for (int x = 0; x < cluster.clusterInfo().getConfiguredNodes().size(); x++) {
-            State state = UP;
-            cluster.clusterInfo().getDistributorNodeInfo(x).setReportedState(new NodeState(DISTRIBUTOR, state), 0);
+            cluster.clusterInfo().getDistributorNodeInfo(x).setReportedState(new NodeState(DISTRIBUTOR, UP), 0);
             cluster.clusterInfo().getDistributorNodeInfo(x).setHostInfo(HostInfo.createHostInfo(createDistributorHostInfo(4, 5, 6)));
-            cluster.clusterInfo().getStorageNodeInfo(x).setReportedState(new NodeState(STORAGE, state), 0);
+            cluster.clusterInfo().getStorageNodeInfo(x).setReportedState(new NodeState(STORAGE, UP), 0);
         }
 
         return nodeStateChangeChecker.evaluateTransition(
@@ -763,6 +857,18 @@ public class NodeStateChangeCheckerTest {
         return configBuilder.build();
     }
 
+    private void checkSettingToMaintenanceIsAllowed(int nodeIndex, NodeStateChangeChecker nodeStateChangeChecker, ClusterState clusterState) {
+        Node node = new Node(STORAGE, nodeIndex);
+        Result result = nodeStateChangeChecker.evaluateTransition(node, clusterState, SAFE, UP_NODE_STATE, MAINTENANCE_NODE_STATE);
+        assertTrue(result.settingWantedStateIsAllowed());
+        assertFalse(result.wantedStateAlreadySet());
+        assertEquals("Preconditions fulfilled and new state different", result.getReason());
+    }
+
+    private void setStorageNodeWantedStateToMaintenance(ContentCluster cluster, int nodeIndex) {
+        cluster.clusterInfo().getStorageNodeInfo(nodeIndex).setWantedState(MAINTENANCE_NODE_STATE.setDescription("Orchestrator"));
+    }
+
     private void setStorageNodeWantedState(ContentCluster cluster, int nodeIndex, State state, String description) {
         NodeState nodeState = new NodeState(STORAGE, state);
         cluster.clusterInfo().getStorageNodeInfo(nodeIndex).setWantedState(nodeState.setDescription(description));