HDDS-12455. Merge the SimpleMockNodeManager with MockNodeManager

hadoop-hdds/server-scm/src/test/java/org/apache/hadoop/hdds/scm/container/MockNodeManager.java

@@ -106,6 +106,7 @@ public class MockNodeManager implements NodeManager {
   private final Map<DatanodeDetails, SCMNodeStat> nodeMetricMap;
   private final SCMNodeStat aggregateStat;
   private final Map<DatanodeID, List<SCMCommand<?>>> commandMap;
+  private final ConcurrentMap<DatanodeID, DatanodeInfo> nodeMap;
   private Node2PipelineMap node2PipelineMap;
   private final NodeStateMap node2ContainerMap;
   private NetworkTopology clusterMap;
@@ -120,6 +121,7 @@ public class MockNodeManager implements NodeManager {
     this.nodeMetricMap = new HashMap<>();
     this.node2PipelineMap = new Node2PipelineMap();
     this.node2ContainerMap = new NodeStateMap();
+    this.nodeMap = new ConcurrentHashMap<>();
     this.dnsToUuidMap = new ConcurrentHashMap<>();
     this.aggregateStat = new SCMNodeStat();
     this.clusterMap = new NetworkTopologyImpl(new OzoneConfiguration());
@@ -207,6 +209,15 @@ public MockNodeManager(
         NUM_PIPELINE_PER_METADATA_DISK;
   }
 
+  private DatanodeInfo getRequiredInfo(DatanodeDetails dd)
+      throws NodeNotFoundException {
+    DatanodeInfo info = nodeMap.get(dd.getID());
+    if (info == null) {
+      throw new NodeNotFoundException(dd.getID());
+    }
+    return info;
+  }
+
   /**
    * Invoked from ctor to create some node Metrics.
    *
@@ -417,6 +428,11 @@ public SCMNodeMetric getNodeStat(DatanodeDetails datanodeDetails) {
   @Override
   public NodeStatus getNodeStatus(DatanodeDetails dd)
       throws NodeNotFoundException {
+    DatanodeInfo info = nodeMap.get(dd.getID());
+    if (info != null) {
+      return info.getNodeStatus();
+    }
+
     if (healthyNodes.contains(dd)) {
       return NodeStatus.inServiceHealthy();
     } else if (staleNodes.contains(dd)) {
@@ -436,6 +452,7 @@ public NodeStatus getNodeStatus(DatanodeDetails dd)
   @Override
   public void setNodeOperationalState(DatanodeDetails datanodeDetails,
       HddsProtos.NodeOperationalState newState) throws NodeNotFoundException {
+    setNodeOperationalState(datanodeDetails, newState, 0);
   }
 
   /**
@@ -447,6 +464,19 @@ public void setNodeOperationalState(DatanodeDetails datanodeDetails,
   public void setNodeOperationalState(DatanodeDetails datanodeDetails,
       HddsProtos.NodeOperationalState newState, long opStateExpiryEpocSec)
       throws NodeNotFoundException {
+    DatanodeInfo info = getRequiredInfo(datanodeDetails);
+    NodeStatus old = info.getNodeStatus();
+
+    NodeStatus updated = NodeStatus.valueOf(
+        newState,
+        old.getHealth(),
+        opStateExpiryEpocSec);
+
+    datanodeDetails.setPersistedOpState(updated.getOperationalState());
+    datanodeDetails.setPersistedOpStateExpiryEpochSec(
+        updated.getOpStateExpiryEpochSeconds());
+
+    info.setNodeStatus(updated);
   }
 
   /**
@@ -652,6 +682,12 @@ public void setNodeState(DatanodeDetails dn, HddsProtos.NodeState state) {
     } else {
       deadNodes.add(dn);
     }
+    DatanodeInfo info = nodeMap.get(dn.getID());
+    if (info != null) {
+      NodeStatus old = info.getNodeStatus();
+      info.setNodeStatus(NodeStatus.valueOf(
+          old.getOperationalState(), state, old.getOpStateExpiryEpochSeconds()));
+    }
   }
 
   /**
@@ -695,7 +731,12 @@ public RegisteredCommand register(DatanodeDetails datanodeDetails,
                                     NodeReportProto nodeReport,
                                     PipelineReportsProto pipelineReportsProto,
                                     LayoutVersionProto layoutInfo) {
+    NodeStatus s = NodeStatus.inServiceHealthy();
+    datanodeDetails.setPersistedOpState(s.getOperationalState());
+    datanodeDetails.setPersistedOpStateExpiryEpochSec(
+        s.getOpStateExpiryEpochSeconds());
     final DatanodeInfo info = new DatanodeInfo(datanodeDetails, NodeStatus.inServiceHealthy(), layoutInfo);
+    nodeMap.put(datanodeDetails.getID(), info);
     try {
       node2ContainerMap.addNode(info);
       addEntryTodnsToUuidMap(datanodeDetails.getIpAddress(),
@@ -741,9 +782,9 @@ public List<SCMCommand<?>> processHeartbeat(DatanodeDetails datanodeDetails,
   }
 
   @Override
-  public Boolean isNodeRegistered(
-      DatanodeDetails datanodeDetails) {
-    return healthyNodes.contains(datanodeDetails);
+  public Boolean isNodeRegistered(DatanodeDetails datanodeDetails) {
+    return nodeMap.containsKey(datanodeDetails.getID()) ||
+        healthyNodes.contains(datanodeDetails);
   }
 
   @Override

hadoop-hdds/server-scm/src/test/java/org/apache/hadoop/hdds/scm/container/TestMockNodeManagerOperationalState.java

@@ -0,0 +1,58 @@
+/*
+ * 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 org.apache.hadoop.hdds.scm.container;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.apache.hadoop.hdds.protocol.DatanodeDetails;
+import org.apache.hadoop.hdds.protocol.MockDatanodeDetails;
+import org.apache.hadoop.hdds.protocol.proto.HddsProtos;
+import org.apache.hadoop.hdds.scm.node.NodeStatus;
+import org.apache.hadoop.hdds.scm.node.states.NodeNotFoundException;
+import org.junit.jupiter.api.Test;
+
+class TestMockNodeManagerOperationalState {
+
+  @Test
+  void testSetNodeOperationalStateUpdatesNodeStatus() throws Exception {
+    MockNodeManager nm = new MockNodeManager(false, 0);
+    DatanodeDetails dn = MockDatanodeDetails.randomDatanodeDetails();
+
+    // register via NodeManager default overload (what ctors do)
+    nm.register(dn, null, null);
+
+    nm.setNodeOperationalState(dn,
+        HddsProtos.NodeOperationalState.DECOMMISSIONING, 123L);
+
+    NodeStatus s = nm.getNodeStatus(dn);
+    assertEquals(HddsProtos.NodeOperationalState.DECOMMISSIONING,
+        s.getOperationalState());
+    assertEquals(123L, s.getOpStateExpiryEpochSeconds());
+  }
+
+  @Test
+  void testSetNodeOperationalStateForUnknownNodeThrows() {
+    MockNodeManager nm = new MockNodeManager(false, 0);
+    DatanodeDetails dn = MockDatanodeDetails.randomDatanodeDetails();
+
+    assertThrows(NodeNotFoundException.class, () ->
+        nm.setNodeOperationalState(dn,
+            HddsProtos.NodeOperationalState.DECOMMISSIONING, 123L));
+  }
+}