Unify node patch methods in tests

8 files changed, 34 insertions, 40 deletions

diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/autoscale/AutoscalingTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/autoscale/AutoscalingTester.java
index 19911076e69..cc755c01405 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/autoscale/AutoscalingTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/autoscale/AutoscalingTester.java
@@ -93,7 +93,7 @@ class AutoscalingTester {
 
     public void makeReady(String hostname) {
         Node node = nodeRepository().getNode(hostname).get();
-        nodeRepository().write(node.with(new IP.Config(Set.of("::" + 0 + ":0"), Set.of())), nodeRepository().lock(node));
+        provisioningTester.patchNode(node, (n) -> n.with(new IP.Config(Set.of("::" + 0 + ":0"), Set.of())));
         Node host = nodeRepository().getNode(node.parentHostname().get()).get();
         host = host.with(new IP.Config(Set.of("::" + 0 + ":0"), Set.of("::" + 0 + ":2")));
         if (host.state() == Node.State.provisioned)
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InactiveAndFailedExpirerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InactiveAndFailedExpirerTest.java
index 3d17cbf0217..056fe041377 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InactiveAndFailedExpirerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InactiveAndFailedExpirerTest.java
@@ -133,7 +133,7 @@ public class InactiveAndFailedExpirerTest {
         // Flag one node for retirement and redeploy
         {
             Node toRetire = tester.getNodes(applicationId, Node.State.active).asList().get(0);
-            tester.patchNode(toRetire.withWantToRetire(true, Agent.operator, tester.clock().instant()));
+            tester.patchNode(toRetire, (node) -> node.withWantToRetire(true, Agent.operator, tester.clock().instant()));
             List<HostSpec> hostSpecs = tester.prepare(applicationId, cluster, Capacity.from(new ClusterResources(2, 1, nodeResources)));
             tester.activate(applicationId, new HashSet<>(hostSpecs));
         }
@@ -203,9 +203,7 @@ public class InactiveAndFailedExpirerTest {
         assertEquals(2, inactiveNodes.size());
 
         // Nodes marked for deprovisioning are moved to parked
-        tester.nodeRepository().write(inactiveNodes.stream()
-                                                   .map(node -> node.withWantToRetire(true, true, Agent.system, tester.clock().instant()))
-                                                   .collect(Collectors.toList()), () -> {});
+        tester.patchNodes(inactiveNodes, (node) -> node.withWantToRetire(true, true, Agent.system, tester.clock().instant()));
         tester.advanceTime(Duration.ofMinutes(11));
         new InactiveExpirer(tester.nodeRepository(), tester.clock(), Duration.ofMinutes(10), new TestMetric()).run();
         assertEquals(2, tester.nodeRepository().getNodes(Node.State.parked).size());
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/OsUpgradeActivatorTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/OsUpgradeActivatorTest.java
index 218812f9a3d..f795dbaaa1c 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/OsUpgradeActivatorTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/OsUpgradeActivatorTest.java
@@ -97,13 +97,7 @@ public class OsUpgradeActivatorTest {
     }
 
     private void completeUpgradeOf(List<Node> nodes) {
-        for (var node : nodes) {
-            try (var lock = tester.nodeRepository().lock(node)) {
-                node =  tester.nodeRepository().getNode(node.hostname()).get();
-                node = node.with(node.status().withVespaVersion(node.allocation().get().membership().cluster().vespaVersion()));
-                tester.nodeRepository().write(node, lock);
-            }
-        }
+        tester.patchNodes(nodes, (node) -> node.with(node.status().withVespaVersion(node.allocation().get().membership().cluster().vespaVersion())));
     }
 
     private Stream<Node> streamUpdatedNodes(List<Node> nodes) {
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
index 6a41e766ace..715ecdb5949 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/os/OsVersionsTest.java
@@ -21,7 +21,6 @@ import java.util.List;
 import java.util.Optional;
 import java.util.function.Function;
 import java.util.function.Supplier;
-import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 
 import static org.junit.Assert.assertEquals;
@@ -273,20 +272,11 @@ public class OsVersionsTest {
     }
 
     private void setWantedVersion(List<Node> nodes, Version wantedVersion) {
-        writeNode(nodes, node -> node.with(node.status().withOsVersion(node.status().osVersion().withWanted(Optional.of(wantedVersion)))));
+        tester.patchNodes(nodes, node -> node.with(node.status().withOsVersion(node.status().osVersion().withWanted(Optional.of(wantedVersion)))));
     }
 
     private void setCurrentVersion(List<Node> nodes, Version currentVersion) {
-        writeNode(nodes, node -> node.with(node.status().withOsVersion(node.status().osVersion().withCurrent(Optional.of(currentVersion)))));
-    }
-
-    private void writeNode(List<Node> nodes, UnaryOperator<Node> updateFunc) {
-        for (var node : nodes) {
-            try (var lock = tester.nodeRepository().lock(node)) {
-                node = tester.nodeRepository().getNode(node.hostname()).get();
-                tester.nodeRepository().write(updateFunc.apply(node), lock);
-            }
-        }
+        tester.patchNodes(nodes, node -> node.with(node.status().withOsVersion(node.status().osVersion().withCurrent(Optional.of(currentVersion)))));
     }
 
     private void completeUpgradeOf(List<Node> nodes) {
@@ -294,7 +284,8 @@ public class OsVersionsTest {
     }
 
     private void completeUpgradeOf(List<Node> nodes, NodeType nodeType) {
-        writeNode(nodes, (node) -> {
+        // Complete upgrade by deprovisioning stale hosts and provisioning new ones
+        tester.patchNodes(nodes, (node) -> {
             Optional<Version> wantedOsVersion = node.status().osVersion().wanted();
             if (node.status().wantToDeprovision()) {
                 // Complete upgrade by deprovisioning stale hosts and provisioning new ones
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerProvisionTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerProvisionTest.java
index 411283abf33..57c7c46c2d9 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerProvisionTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DynamicDockerProvisionTest.java
@@ -128,8 +128,7 @@ public class DynamicDockerProvisionTest {
         NodeResources resources = new NodeResources(10, 10, 10, 10);
         ApplicationId app = tester.makeApplicationId();
 
-        Function<Node, Node> retireNode = node ->
-                tester.nodeRepository().write(node.withWantToRetire(true, Agent.system, Instant.now()), () -> {});
+        Function<Node, Node> retireNode = node -> tester.patchNode(node, (n) -> n.withWantToRetire(true, Agent.system, Instant.now()));
         Function<Integer, Node> getNodeInGroup = group -> tester.nodeRepository().getNodes(app).stream()
                 .filter(node -> node.allocation().get().membership().cluster().group().get().index() == group)
                 .findAny().orElseThrow();
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java
index d53a0732c89..6cf5a2c8342 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTest.java
@@ -673,7 +673,7 @@ public class ProvisioningTest {
         ApplicationId application = tester.makeApplicationId();
         // Flag all nodes for retirement
         List<Node> readyNodes = tester.makeReadyNodes(5, defaultResources);
-        readyNodes.forEach(node -> tester.patchNode(node.withWantToRetire(true, Agent.system, tester.clock().instant())));
+        tester.patchNodes(readyNodes, (node) -> node.withWantToRetire(true, Agent.system, tester.clock().instant()));
 
         try {
             prepare(application, 2, 0, 2, 0, defaultResources, tester);
@@ -701,7 +701,7 @@ public class ProvisioningTest {
         assertEquals(0, NodeList.copyOf(tester.nodeRepository().getNodes(application, Node.State.active)).retired().size());
 
         // Mark the nodes as want to retire
-        tester.nodeRepository().getNodes(application, Node.State.active).forEach(node -> tester.patchNode(node.withWantToRetire(true, Agent.system, tester.clock().instant())));
+        tester.nodeRepository().getNodes(application, Node.State.active).forEach(node -> tester.patchNode(node, (n) -> n.withWantToRetire(true, Agent.system, tester.clock().instant())));
         // redeploy without allow failing
         tester.activate(application, tester.prepare(application, cluster, capacityFORCED));
 
@@ -786,7 +786,7 @@ public class ProvisioningTest {
         // Retire some nodes and redeploy
         {
             List<Node> nodesToRetire = tester.getNodes(application, Node.State.active).asList().subList(0, 2);
-            nodesToRetire.forEach(node -> tester.patchNode(node.withWantToRetire(true, Agent.system, tester.clock().instant())));
+            tester.patchNodes(nodesToRetire, (node) -> node.withWantToRetire(true, Agent.system, tester.clock().instant()));
 
             SystemState state = prepare(application, 2, 0, 2, 0, defaultResources, tester);
             tester.activate(application, state.allHosts);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
index 0a3c85d3702..4c8d5caad43 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/ProvisioningTester.java
@@ -53,6 +53,7 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
 import java.util.function.Function;
+import java.util.function.UnaryOperator;
 import java.util.logging.Level;
 import java.util.stream.Collectors;
 
@@ -139,7 +140,22 @@ public class ProvisioningTester {
     public CapacityPolicies capacityPolicies() { return capacityPolicies; }
     public NodeList getNodes(ApplicationId id, Node.State ... inState) { return NodeList.copyOf(nodeRepository.getNodes(id, inState)); }
 
-    public void patchNode(Node node) { nodeRepository.write(node, () -> {}); }
+    public Node patchNode(Node node, UnaryOperator<Node> patcher) {
+        return patchNodes(List.of(node), patcher).get(0);
+    }
+
+    public List<Node> patchNodes(List<Node> nodes, UnaryOperator<Node> patcher) {
+        List<Node> updated = new ArrayList<>();
+        for (var node : nodes) {
+            try (var lock = nodeRepository.lock(node)) {
+                node = nodeRepository.getNode(node.hostname()).get();
+                node = patcher.apply(node);
+                nodeRepository.write(node, lock);
+                updated.add(node);
+            }
+        }
+        return updated;
+    }
 
     public List<HostSpec> prepare(ApplicationId application, ClusterSpec cluster, int nodeCount, int groups, NodeResources resources) {
         return prepare(application, cluster, nodeCount, groups, false, resources);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/VirtualNodeProvisioningTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/VirtualNodeProvisioningTest.java
index 51261c29a71..0ef7071b095 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/VirtualNodeProvisioningTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/VirtualNodeProvisioningTest.java
@@ -129,9 +129,7 @@ public class VirtualNodeProvisioningTest {
         assertDistinctParentHosts(nodes, ClusterSpec.Type.container, containerNodeCount);
         assertDistinctParentHosts(nodes, ClusterSpec.Type.content, contentNodeCount);
 
-        for (Node n : nodes) {
-          tester.patchNode(n.withParentHostname("clashing"));
-        }
+        tester.patchNodes(nodes, (n) -> n.withParentHostname("clashing"));
         containerHosts = prepare(containerClusterSpec, containerNodeCount, groups);
         contentHosts = prepare(contentClusterSpec, contentNodeCount, groups);
         activate(containerHosts, contentHosts);
@@ -160,9 +158,7 @@ public class VirtualNodeProvisioningTest {
         assertDistinctParentHosts(nodes, ClusterSpec.Type.container, containerNodeCount);
         assertDistinctParentHosts(nodes, ClusterSpec.Type.content, contentNodeCount);
 
-        for (Node n : nodes) {
-          tester.patchNode(n.withParentHostname("clashing"));
-        }
+        tester.patchNodes(nodes, (n) -> n.withParentHostname("clashing"));
         OutOfCapacityException expected = null;
         try {
           containerHosts = prepare(containerClusterSpec, containerNodeCount, groups);
@@ -216,9 +212,9 @@ public class VirtualNodeProvisioningTest {
         assertEquals(3, nodes.size());
 
         // Set indistinct parents
-        tester.patchNode(nodes.get(0).withParentHostname("parentHost1"));
-        tester.patchNode(nodes.get(1).withParentHostname("parentHost1"));
-        tester.patchNode(nodes.get(2).withParentHostname("parentHost2"));
+        tester.patchNode(nodes.get(0), (n) -> n.withParentHostname("parentHost1"));
+        tester.patchNode(nodes.get(1), (n) -> n.withParentHostname("parentHost1"));
+        tester.patchNode(nodes.get(2), (n) -> n.withParentHostname("parentHost2"));
         nodes = getNodes(applicationId);
         assertEquals(3, nodes.stream().filter(n -> n.parentHostname().isPresent()).count());