1 files changed, 32 insertions, 9 deletions

diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequestTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequestTest.java
index 9239b8774b0..7161fb1be79 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequestTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/restapiv2/requests/SetNodeStateRequestTest.java
@@ -23,20 +23,23 @@ import java.util.Map;
 import java.util.Optional;
 
 import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class SetNodeStateRequestTest {
-    public static final String REASON = "operator";
-    ContentCluster cluster = mock(ContentCluster.class);
-    SetUnitStateRequest.Condition condition = SetUnitStateRequest.Condition.SAFE;
-    Map<String, UnitState> newStates = new HashMap<>();
-    UnitState unitState = mock(UnitState.class);
+    private static final String REASON = "operator";
+    private ContentCluster cluster = mock(ContentCluster.class);
+    private SetUnitStateRequest.Condition condition = SetUnitStateRequest.Condition.SAFE;
+    private Map<String, UnitState> newStates = new HashMap<>();
+    private UnitState unitState = mock(UnitState.class);
     private final int NODE_INDEX = 2;
-    Node storageNode = new Node(NodeType.STORAGE, NODE_INDEX);
-    NodeStateOrHostInfoChangeHandler stateListener = mock(NodeStateOrHostInfoChangeHandler.class);
-    ClusterState currentClusterState = mock(ClusterState.class);
+    private Node storageNode = new Node(NodeType.STORAGE, NODE_INDEX);
+    private NodeStateOrHostInfoChangeHandler stateListener = mock(NodeStateOrHostInfoChangeHandler.class);
+    private ClusterState currentClusterState = mock(ClusterState.class);
+    private boolean probe = false;
 
     @Before
     public void setUp() {
@@ -53,6 +56,16 @@ public class SetNodeStateRequestTest {
     }
 
     @Test
+    public void testProbingDoesntChangeState() throws StateRestApiException {
+        probe = true;
+        testSetStateRequest(
+                "maintenance",
+                State.UP, State.UP,
+                NodeStateChangeChecker.Result.allowSettingOfWantedState(),
+                Optional.empty(), Optional.empty());
+    }
+
+    @Test
     public void testUpToDown() throws StateRestApiException {
         testSetStateRequest(
                 "down",
@@ -124,6 +137,9 @@ public class SetNodeStateRequestTest {
         when(cluster.getNodeInfo(distributorNode)).thenReturn(distributorNodeInfo);
 
         NodeState distributorNodeState = new NodeState(distributorNode.getType(), distributorWantedState);
+        if (distributorWantedState != State.UP) {
+            distributorNodeState.setDescription(REASON);
+        }
         when(distributorNodeInfo.getUserWantedState()).thenReturn(distributorNodeState);
 
         setWantedState();
@@ -133,6 +149,9 @@ public class SetNodeStateRequestTest {
                     new NodeState(NodeType.STORAGE, expectedNewStorageWantedState.get());
             verify(storageNodeInfo).setWantedState(expectedNewStorageNodeState);
             verify(stateListener).handleNewWantedNodeState(storageNodeInfo, expectedNewStorageNodeState);
+        } else {
+            verify(storageNodeInfo, times(0)).setWantedState(any());
+            verify(stateListener, times(0)).handleNewWantedNodeState(eq(storageNodeInfo), any());
         }
 
         if (expectedNewDistributorWantedState.isPresent()) {
@@ -140,6 +159,9 @@ public class SetNodeStateRequestTest {
                     new NodeState(NodeType.DISTRIBUTOR, expectedNewDistributorWantedState.get());
             verify(distributorNodeInfo).setWantedState(expectedNewDistributorNodeState);
             verify(stateListener).handleNewWantedNodeState(distributorNodeInfo, expectedNewDistributorNodeState);
+        } else {
+            verify(distributorNodeInfo, times(0)).setWantedState(any());
+            verify(stateListener, times(0)).handleNewWantedNodeState(eq(distributorNodeInfo), any());
         }
     }
 
@@ -150,6 +172,7 @@ public class SetNodeStateRequestTest {
                 newStates,
                 storageNode,
                 stateListener,
-                currentClusterState);
+                currentClusterState,
+                probe);
     }
 }
\ No newline at end of file