1 files changed, 8 insertions, 3 deletions
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/NodeRepositoryMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/NodeRepositoryMock.java
index afb56f10c38..4079591730d 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/NodeRepositoryMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/NodeRepositoryMock.java
@@ -59,9 +59,14 @@ public class NodeRepositoryMock implements NodeRepository {
 
     /** Add or update given nodes in zone */
     public void putNodes(ZoneId zone, List<Node> nodes) {
-        nodeRepository.putIfAbsent(zone, new HashMap<>());
-        nodeRepository.get(zone).putAll(nodes.stream().collect(Collectors.toMap(Node::hostname,
-                                                                                Function.identity())));
+        Map<HostName, Node> zoneNodes = nodeRepository.computeIfAbsent(zone, (k) -> new HashMap<>());
+        for (var node : nodes) {
+            zoneNodes.put(node.hostname(), node);
+        }
+    }
+
+    public void putNode(ZoneId zone, Node node) {
+        nodeRepository.computeIfAbsent(zone, (k) -> new HashMap<>()).put(node.hostname(), node);
     }
 
     public void putApplication(ZoneId zone, Application application) {