Use NodeRepoMock in test, no functional changes

5 files changed, 55 insertions, 46 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdater.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdater.java
index c24b2261f42..726e23a3fe4 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdater.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdater.java
@@ -102,6 +102,8 @@ public class NodeAdminStateUpdater {
      * with respect to: freeze, Orchestrator, and services running.
      */
     public void converge(State wantedState) {
+        NodeSpec node = nodeRepository.getNode(hostHostname);
+        boolean hostIsActiveInNR = node.state() == NodeState.active;
         if (wantedState == RESUMED) {
             adjustNodeAgentsToRunFromNodeRepository();
         } else if (currentState == TRANSITIONING && nodeAdmin.subsystemFreezeDuration().compareTo(FREEZE_CONVERGENCE_TIMEOUT) > 0) {
@@ -110,8 +112,7 @@ public class NodeAdminStateUpdater {
             adjustNodeAgentsToRunFromNodeRepository();
             nodeAdmin.setFrozen(false);
 
-            NodeState currentNodeState = nodeRepository.getNode(hostHostname).state();
-            if (currentNodeState == NodeState.active) orchestrator.resume(hostHostname);
+            if (hostIsActiveInNR) orchestrator.resume(hostHostname);
 
             throw new ConvergenceException("Timed out trying to freeze all nodes: will force an unfrozen tick");
         }
@@ -120,11 +121,9 @@ public class NodeAdminStateUpdater {
         currentState = TRANSITIONING;
 
         boolean wantFrozen = wantedState != RESUMED;
-        if (!nodeAdmin.setFrozen(wantFrozen)) {
+        if (!nodeAdmin.setFrozen(wantFrozen))
             throw new ConvergenceException("NodeAdmin is not yet " + (wantFrozen ? "frozen" : "unfrozen"));
-        }
 
-        boolean hostIsActiveInNR = nodeRepository.getNode(hostHostname).state() == NodeState.active;
         switch (wantedState) {
             case RESUMED:
                 if (hostIsActiveInNR) orchestrator.resume(hostHostname);
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/ContainerTester.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/ContainerTester.java
index af30d3cbe56..1f3ab416db8 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/ContainerTester.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/ContainerTester.java
@@ -78,7 +78,7 @@ public class ContainerTester implements AutoCloseable {
         for (int i = 1; i < 4; i++) ipAddresses.addAddress("host" + i + ".test.yahoo.com", "f000::" + i);
 
         NodeSpec hostSpec = NodeSpec.Builder.testSpec(HOST_HOSTNAME.value()).type(NodeType.host).build();
-        nodeRepository.updateNodeRepositoryNode(hostSpec);
+        nodeRepository.updateNodeSpec(hostSpec);
 
         Clock clock = Clock.systemUTC();
         Metrics metrics = new Metrics();
@@ -122,7 +122,7 @@ public class ContainerTester implements AutoCloseable {
                                                    ", but that image does not exist in the container engine");
             }
         }
-        nodeRepository.updateNodeRepositoryNode(new NodeSpec.Builder(nodeSpec)
+        nodeRepository.updateNodeSpec(new NodeSpec.Builder(nodeSpec)
                 .parentHostname(HOST_HOSTNAME.value())
                 .build());
     }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/NodeRepoMock.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/NodeRepoMock.java
index 5722de4cf90..0c986929de1 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/NodeRepoMock.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/integration/NodeRepoMock.java
@@ -3,15 +3,17 @@ package com.yahoo.vespa.hosted.node.admin.integration;
 
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.Acl;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.AddNode;
+import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NoSuchNodeException;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeAttributes;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeRepository;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeState;
 
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 /**
@@ -20,55 +22,60 @@ import java.util.stream.Collectors;
  * @author dybis
  */
 public class NodeRepoMock implements NodeRepository {
-    private static final Object monitor = new Object();
 
-    private final Map<String, NodeSpec> nodeRepositoryNodesByHostname = new HashMap<>();
+    private final Map<String, NodeSpec> nodeSpecByHostname = new ConcurrentHashMap<>();
+    private volatile Map<String, Acl> aclByHostname = Map.of();
 
     @Override
     public void addNodes(List<AddNode> nodes) { }
 
     @Override
     public List<NodeSpec> getNodes(String baseHostName) {
-        synchronized (monitor) {
-            return nodeRepositoryNodesByHostname.values().stream()
-                    .filter(node -> baseHostName.equals(node.parentHostname().orElse(null)))
-                    .collect(Collectors.toList());
-        }
+        return nodeSpecByHostname.values().stream()
+                .filter(node -> baseHostName.equals(node.parentHostname().orElse(null)))
+                .collect(Collectors.toList());
     }
 
     @Override
     public Optional<NodeSpec> getOptionalNode(String hostName) {
-        synchronized (monitor) {
-            return Optional.ofNullable(nodeRepositoryNodesByHostname.get(hostName));
-        }
+        return Optional.ofNullable(nodeSpecByHostname.get(hostName));
     }
 
     @Override
     public Map<String, Acl> getAcls(String hostname) {
-        return Map.of();
+        return aclByHostname;
     }
 
     @Override
     public void updateNodeAttributes(String hostName, NodeAttributes nodeAttributes) {
-        synchronized (monitor) {
-            updateNodeRepositoryNode(new NodeSpec.Builder(getNode(hostName))
-                    .updateFromNodeAttributes(nodeAttributes)
-                    .build());
-        }
+        updateNodeSpec(new NodeSpec.Builder(getNode(hostName))
+                .updateFromNodeAttributes(nodeAttributes)
+                .build());
     }
 
     @Override
     public void setNodeState(String hostName, NodeState nodeState) {
-        synchronized (monitor) {
-            updateNodeRepositoryNode(new NodeSpec.Builder(getNode(hostName))
-                    .state(nodeState)
-                    .build());
-        }
+        updateNodeSpec(new NodeSpec.Builder(getNode(hostName))
+                .state(nodeState)
+                .build());
     }
 
-    public void updateNodeRepositoryNode(NodeSpec nodeSpec) {
-        synchronized (monitor) {
-            nodeRepositoryNodesByHostname.put(nodeSpec.hostname(), nodeSpec);
-        }
+    public void updateNodeSpec(NodeSpec nodeSpec) {
+        nodeSpecByHostname.put(nodeSpec.hostname(), nodeSpec);
+    }
+
+    public void updateNodeSpec(String hostname, Function<NodeSpec.Builder, NodeSpec.Builder> mapper) {
+        nodeSpecByHostname.compute(hostname, (__, nodeSpec) -> {
+            if (nodeSpec == null) throw new NoSuchNodeException(hostname);
+            return mapper.apply(new NodeSpec.Builder(nodeSpec)).build();
+        });
+    }
+
+    public void resetNodeSpecs() {
+        nodeSpecByHostname.clear();
+    }
+
+    public void setAcl(Map<String, Acl> aclByHostname) {
+        this.aclByHostname = Map.copyOf(aclByHostname);
     }
 }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java
index 53e283abb42..19d7e294367 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/maintenance/servicedump/VespaServiceDumperImplTest.java
@@ -65,7 +65,7 @@ class VespaServiceDumperImplTest {
                 .report(ServiceDumpReport.REPORT_ID, request.toJsonNode())
                 .archiveUri(URI.create("s3://uri-1/tenant1/"))
                 .build();
-        nodeRepository.updateNodeRepositoryNode(initialSpec);
+        nodeRepository.updateNodeSpec(initialSpec);
 
         // Create dumper and invoke tested method
         VespaServiceDumper reporter = new VespaServiceDumperImpl(operations, syncClient, nodeRepository, clock);
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java
index 8ee3a95744b..9ab340a2421 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/nodeadmin/NodeAdminStateUpdaterTest.java
@@ -2,13 +2,15 @@
 package com.yahoo.vespa.hosted.node.admin.nodeadmin;
 
 import com.yahoo.config.provision.HostName;
+import com.yahoo.config.provision.NodeType;
 import com.yahoo.test.ManualClock;
 import com.yahoo.vespa.flags.InMemoryFlagSource;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.Acl;
-import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeRepository;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeSpec;
 import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.NodeState;
+import com.yahoo.vespa.hosted.node.admin.configserver.noderepository.OrchestratorStatus;
 import com.yahoo.vespa.hosted.node.admin.configserver.orchestrator.Orchestrator;
+import com.yahoo.vespa.hosted.node.admin.integration.NodeRepoMock;
 import com.yahoo.vespa.hosted.node.admin.nodeagent.NodeAgentContextFactory;
 import org.junit.Test;
 
@@ -16,7 +18,6 @@ import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.Map;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -45,7 +46,7 @@ import static org.mockito.Mockito.when;
  */
 public class NodeAdminStateUpdaterTest {
     private final NodeAgentContextFactory nodeAgentContextFactory = mock(NodeAgentContextFactory.class);
-    private final NodeRepository nodeRepository = mock(NodeRepository.class);
+    private final NodeRepoMock nodeRepository = spy(new NodeRepoMock());
     private final Orchestrator orchestrator = mock(Orchestrator.class);
     private final NodeAdmin nodeAdmin = mock(NodeAdmin.class);
     private final HostName hostHostname = HostName.from("basehost1.test.yahoo.com");
@@ -238,20 +239,22 @@ public class NodeAdminStateUpdaterTest {
     }
 
     private void mockNodeRepo(NodeState hostState, int numberOfNodes) {
-        List<NodeSpec> containersToRun = IntStream.range(1, numberOfNodes + 1)
-                .mapToObj(i -> NodeSpec.Builder.testSpec("host" + i + ".yahoo.com").build())
-                .collect(Collectors.toList());
+        nodeRepository.resetNodeSpecs();
+
+        IntStream.rangeClosed(1, numberOfNodes)
+                .mapToObj(i -> NodeSpec.Builder.testSpec("host" + i + ".yahoo.com").parentHostname(hostHostname.value()).build())
+                .forEach(nodeRepository::updateNodeSpec);
 
-        when(nodeRepository.getNodes(eq(hostHostname.value()))).thenReturn(containersToRun);
-        when(nodeRepository.getNode(eq(hostHostname.value()))).thenReturn(
-                NodeSpec.Builder.testSpec(hostHostname.value(), hostState).build());
+        nodeRepository.updateNodeSpec(NodeSpec.Builder.testSpec(hostHostname.value(), hostState).type(NodeType.host).build());
     }
 
     private void mockAcl(Acl acl, int... nodeIds) {
-        Map<String, Acl> aclByHostname = Arrays.stream(nodeIds)
+        nodeRepository.setAcl(Arrays.stream(nodeIds)
                 .mapToObj(i -> "host" + i + ".yahoo.com")
-                .collect(Collectors.toMap(Function.identity(), h -> acl));
+                .collect(Collectors.toMap(Function.identity(), h -> acl)));
+    }
 
-        when(nodeRepository.getAcls(eq(hostHostname.value()))).thenReturn(aclByHostname);
+    private void setHostOrchestratorStatus(HostName hostname, OrchestratorStatus orchestratorStatus) {
+        nodeRepository.updateNodeSpec(hostname.value(), node -> node.orchestratorStatus(orchestratorStatus));
     }
 }