Always require upgrade budget

12 files changed, 51 insertions, 54 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
index 5ab2c272b00..56752bc8fd2 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/NodeList.java
@@ -64,6 +64,11 @@ public class NodeList extends AbstractFilteringList<Node, NodeList> {
         return matching(node -> ! nodes.contains(node));
     }
 
+    /** Returns the subset of nodes excluding given node */
+    public NodeList except(Node node) {
+        return except(Set.of(node));
+    }
+
     /** Returns the subset of nodes assigned to the given cluster type */
     public NodeList type(ClusterSpec.Type type) {
         return matching(node -> node.allocation().isPresent() && node.allocation().get().membership().cluster().type().equals(type));
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionChange.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionChange.java
index 053bd7c4feb..b9c7adb4914 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionChange.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionChange.java
@@ -40,7 +40,7 @@ public class OsVersionChange {
     }
 
     /** Returns a copy of this with given target added */
-    public OsVersionChange withTarget(Version version, NodeType nodeType, Optional<Duration> upgradeBudget) {
+    public OsVersionChange withTarget(Version version, NodeType nodeType, Duration upgradeBudget) {
         var targets = new HashMap<>(this.targets);
         targets.compute(nodeType, (key, prevTarget) -> {
             Optional<Instant> lastRetiredAt = Optional.ofNullable(prevTarget).flatMap(OsVersionTarget::lastRetiredAt);
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionTarget.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionTarget.java
index 89c49447f17..582cd92d239 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionTarget.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersionTarget.java
@@ -18,14 +18,15 @@ public class OsVersionTarget {
 
     private final NodeType nodeType;
     private final Version version;
-    private final Optional<Duration> upgradeBudget;
+    private final Duration upgradeBudget;
     private final Optional<Instant> lastRetiredAt;
 
-    public OsVersionTarget(NodeType nodeType, Version version, Optional<Duration> upgradeBudget, Optional<Instant> lastRetiredAt) {
+    public OsVersionTarget(NodeType nodeType, Version version, Duration upgradeBudget, Optional<Instant> lastRetiredAt) {
         this.nodeType = Objects.requireNonNull(nodeType);
         this.version = Objects.requireNonNull(version);
-        this.upgradeBudget = requireNotNegative(upgradeBudget);
+        this.upgradeBudget = Objects.requireNonNull(upgradeBudget);
         this.lastRetiredAt = Objects.requireNonNull(lastRetiredAt);
+        if (upgradeBudget.isNegative()) throw new IllegalArgumentException("Upgrade budget cannot be negative");
     }
 
     /** The node type this applies to */
@@ -39,7 +40,7 @@ public class OsVersionTarget {
     }
 
     /** The upgrade budget for this. All nodes targeting this must upgrade within this budget */
-    public Optional<Duration> upgradeBudget() {
+    public Duration upgradeBudget() {
         return upgradeBudget;
     }
 
@@ -64,11 +65,4 @@ public class OsVersionTarget {
         return Objects.hash(nodeType, version, upgradeBudget, lastRetiredAt);
     }
 
-    private static Optional<Duration> requireNotNegative(Optional<Duration> duration) {
-        Objects.requireNonNull(duration);
-        if (duration.isEmpty()) return duration;
-        if (duration.get().isNegative()) throw new IllegalArgumentException("Duration cannot be negative");
-        return duration;
-    }
-
 }
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
index 6862c91b414..2ee7b324582 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/OsVersions.java
@@ -78,10 +78,9 @@ public class OsVersions {
     }
 
     /** Set the target OS version and upgrade budget for nodes of given type */
-    public void setTarget(NodeType nodeType, Version newTarget, Optional<Duration> upgradeBudget, boolean force) {
+    public void setTarget(NodeType nodeType, Version newTarget, Duration upgradeBudget, boolean force) {
         require(nodeType);
         requireNonZero(newTarget);
-        requireUpgradeBudget(upgradeBudget);
         writeChange((change) -> {
             var oldTarget = targetFor(nodeType);
             if (oldTarget.filter(v -> v.equals(newTarget)).isPresent()) {
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java
index f378a4249f4..23e96d65fc1 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/os/RetiringOsUpgrader.java
@@ -15,10 +15,10 @@ import java.util.Optional;
 import java.util.logging.Logger;
 
 /**
- * An upgrader that retires and deprovisions nodes on stale OS versions. Retirement of each node is spread out in time,
- * according to a time budget, to avoid potential service impact of retiring too many nodes close together.
+ * An upgrader that retires and deprovisions hosts on stale OS versions. Retirement of each host is spread out in time,
+ * according to a time budget, to avoid potential service impact of retiring too many hosts close together.
  *
- * Used in clouds where nodes must be re-provisioned to upgrade their OS.
+ * Used in clouds where hosts must be re-provisioned to upgrade their OS.
  *
  * @author mpolden
  */
@@ -40,8 +40,6 @@ public class RetiringOsUpgrader implements OsUpgrader {
 
         Instant now = nodeRepository.clock().instant();
         Duration nodeBudget = target.upgradeBudget()
-                                    .orElseThrow(() -> new IllegalStateException("OS upgrades in this zone requires " +
-                                                                                 "a time budget, but none is set"))
                                     .dividedBy(activeNodes.size());
         Instant retiredAt = target.lastRetiredAt().orElse(Instant.EPOCH);
         if (now.isBefore(retiredAt.plus(nodeBudget))) return; // Budget has not been spent yet
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializer.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializer.java
index c9890194d73..2034e7f25d2 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializer.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializer.java
@@ -40,7 +40,7 @@ public class OsVersionChangeSerializer {
             var targetObject = targetsObject.addObject();
             targetObject.setString(NODE_TYPE_FIELD, NodeSerializer.toString(nodeType));
             targetObject.setString(VERSION_FIELD, target.version().toFullString());
-            target.upgradeBudget().ifPresent(duration -> targetObject.setLong(UPGRADE_BUDGET_FIELD, duration.toMillis()));
+            targetObject.setLong(UPGRADE_BUDGET_FIELD, target.upgradeBudget().toMillis());
             target.lastRetiredAt().ifPresent(instant -> targetObject.setLong(LAST_RETIRED_AT_FIELD, instant.toEpochMilli()));
         });
         try {
@@ -56,7 +56,9 @@ public class OsVersionChangeSerializer {
         inspector.field(TARGETS_FIELD).traverse((ArrayTraverser) (idx, arrayInspector) -> {
             var version = Version.fromString(arrayInspector.field(VERSION_FIELD).asString());
             var nodeType = NodeSerializer.nodeTypeFromString(arrayInspector.field(NODE_TYPE_FIELD).asString());
-            Optional<Duration> budget = optionalLong(arrayInspector.field(UPGRADE_BUDGET_FIELD)).map(Duration::ofMillis);
+            // TODO(mpolden): Require this field after 2021-05-01
+            Duration budget = optionalLong(arrayInspector.field(UPGRADE_BUDGET_FIELD)).map(Duration::ofMillis)
+                                                                                      .orElse(Duration.ZERO);
             Optional<Instant> lastRetiredAt = optionalLong(arrayInspector.field(LAST_RETIRED_AT_FIELD)).map(Instant::ofEpochMilli);
             targets.put(nodeType, new OsVersionTarget(nodeType, version, budget, lastRetiredAt));
         });
diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
index 0875bf3815d..c850962bf53 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiHandler.java
@@ -392,9 +392,10 @@ public class NodesV2ApiHandler extends LoggingRequestHandler {
                                                                    throw new IllegalArgumentException("Invalid duration '" + s + "'", e);
                                                                }
                                                            });
-                nodeRepository.osVersions().setTarget(nodeType, osVersion, upgradeBudget, force);
+                if (upgradeBudget.isEmpty()) throw new IllegalArgumentException("upgradeBudget must be set");
+                nodeRepository.osVersions().setTarget(nodeType, osVersion, upgradeBudget.get(), force);
                 messageParts.add("osVersion to " + osVersion.toFullString());
-                upgradeBudget.ifPresent(d -> messageParts.add("upgradeBudget to " + d));
+                messageParts.add("upgradeBudget to " + upgradeBudget.get());
             }
         }
 
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooterTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooterTest.java
index dbaa5f034f6..1894ef6ed70 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooterTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/NodeRebooterTest.java
@@ -16,7 +16,6 @@ import org.junit.Test;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
-import java.util.Optional;
 import java.util.stream.Collectors;
 
 import static org.junit.Assert.assertEquals;
@@ -123,7 +122,7 @@ public class NodeRebooterTest {
 
     /** Schedule OS upgrade for all host nodes */
     private void scheduleOsUpgrade(NodeRepository nodeRepository) {
-        nodeRepository.osVersions().setTarget(NodeType.host, Version.fromString("7.0"), Optional.empty(), false);
+        nodeRepository.osVersions().setTarget(NodeType.host, Version.fromString("7.0"), Duration.ZERO, false);
     }
 
     /** Simulate completion of an OS upgrade */
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 8b370669956..9c2332662e8 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
@@ -57,8 +57,8 @@ public class OsUpgradeActivatorTest {
 
         // New OS target version is set
         var osVersion0 = Version.fromString("8.0");
-        osVersions.setTarget(NodeType.host, osVersion0, Optional.empty(), false);
-        osVersions.setTarget(NodeType.confighost, osVersion0, Optional.empty(), false);
+        osVersions.setTarget(NodeType.host, osVersion0, Duration.ZERO, false);
+        osVersions.setTarget(NodeType.confighost, osVersion0, Duration.ZERO, false);
 
         // New OS version is activated as there is no ongoing Vespa upgrade
         osUpgradeActivator.maintain();
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 26c6f742052..3b967a9245d 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
@@ -45,7 +45,7 @@ public class OsVersionsTest {
         // Upgrade OS
         assertTrue("No versions set", versions.readChange().targets().isEmpty());
         var version1 = Version.fromString("7.1");
-        versions.setTarget(NodeType.host, version1, Optional.empty(), false);
+        versions.setTarget(NodeType.host, version1, Duration.ZERO, false);
         assertEquals(version1, versions.targetFor(NodeType.host).get());
         assertTrue("Per-node wanted OS version remains unset", hostNodes.get().stream().allMatch(node -> node.status().osVersion().wanted().isEmpty()));
 
@@ -55,15 +55,14 @@ public class OsVersionsTest {
 
         // Upgrade OS again
         var version2 = Version.fromString("7.2");
-        versions.setTarget(NodeType.host, version2, Optional.empty(), false);
+        versions.setTarget(NodeType.host, version2, Duration.ZERO, false);
         assertEquals(version2, versions.targetFor(NodeType.host).get());
 
         // Resume upgrade
         versions.resumeUpgradeOf(NodeType.host, true);
         NodeList allHosts = hostNodes.get();
-        assertTrue("Wanted version is set", allHosts.stream()
-                                                       .filter(node -> !node.equals(hostOnLaterVersion))
-                                                       .allMatch(node -> node.status().osVersion().wanted().isPresent()));
+        assertTrue("Wanted version is set", allHosts.except(hostOnLaterVersion).stream()
+                                                    .allMatch(node -> node.status().osVersion().wanted().isPresent()));
         assertTrue("Wanted version is not set for host on later version",
                    allHosts.first().get().status().osVersion().wanted().isEmpty());
 
@@ -74,12 +73,12 @@ public class OsVersionsTest {
 
         // Downgrading fails
         try {
-            versions.setTarget(NodeType.host, version1, Optional.empty(), false);
+            versions.setTarget(NodeType.host, version1, Duration.ZERO, false);
             fail("Expected exception");
         } catch (IllegalArgumentException ignored) {}
 
         // Forcing downgrade succeeds
-        versions.setTarget(NodeType.host, version1, Optional.empty(), true);
+        versions.setTarget(NodeType.host, version1, Duration.ZERO, true);
         assertEquals(version1, versions.targetFor(NodeType.host).get());
 
         // Target can be removed
@@ -111,7 +110,7 @@ public class OsVersionsTest {
 
         // Set target
         var version1 = Version.fromString("7.1");
-        versions.setTarget(NodeType.host, version1, Optional.empty(), false);
+        versions.setTarget(NodeType.host, version1, Duration.ZERO, false);
         assertEquals(version1, versions.targetFor(NodeType.host).get());
 
         // Activate target
@@ -148,7 +147,7 @@ public class OsVersionsTest {
 
         // Trigger upgrade to next version
         var version2 = Version.fromString("7.2");
-        versions.setTarget(NodeType.host, version2, Optional.empty(), false);
+        versions.setTarget(NodeType.host, version2, Duration.ZERO, false);
         versions.resumeUpgradeOf(NodeType.host, true);
 
         // Wanted version is changed to newest target for all nodes
@@ -174,7 +173,7 @@ public class OsVersionsTest {
         var version1 = Version.fromString("7.1");
         Duration totalBudget = Duration.ofHours(12);
         Duration nodeBudget = totalBudget.dividedBy(hostCount);
-        versions.setTarget(NodeType.host, version1, Optional.of(totalBudget),false);
+        versions.setTarget(NodeType.host, version1, totalBudget,false);
         versions.resumeUpgradeOf(NodeType.host, true);
 
         // One host is deprovisioning
@@ -216,7 +215,7 @@ public class OsVersionsTest {
         // Another upgrade is triggered. Last retirement time is preserved
         clock.advance(Duration.ofDays(1));
         var version2 = Version.fromString("7.2");
-        versions.setTarget(NodeType.host, version2, Optional.of(totalBudget), false);
+        versions.setTarget(NodeType.host, version2, totalBudget, false);
         assertEquals(lastRetiredAt, versions.readChange().targets().get(NodeType.host).lastRetiredAt().get());
     }
 
@@ -231,7 +230,7 @@ public class OsVersionsTest {
 
         // Target is set with zero budget and upgrade started
         var version1 = Version.fromString("7.1");
-        versions.setTarget(NodeType.confighost, version1, Optional.of(Duration.ZERO),false);
+        versions.setTarget(NodeType.confighost, version1, Duration.ZERO,false);
         for (int i = 0; i < hostCount; i++) {
             versions.resumeUpgradeOf(NodeType.confighost, true);
         }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializerTest.java
index 15b9cfec4ca..3b927fae4cd 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/persistence/OsVersionChangeSerializerTest.java
@@ -22,9 +22,9 @@ public class OsVersionChangeSerializerTest {
     @Test
     public void serialization() {
         var change = new OsVersionChange(Map.of(
-                NodeType.host, new OsVersionTarget(NodeType.host, Version.fromString("1.2.3"), Optional.of(Duration.ofHours(1)), Optional.of(Instant.ofEpochMilli(123))),
-                NodeType.proxyhost, new OsVersionTarget(NodeType.proxyhost, Version.fromString("4.5.6"), Optional.empty(), Optional.empty()),
-                NodeType.confighost, new OsVersionTarget(NodeType.confighost, Version.fromString("7.8.9"), Optional.of(Duration.ZERO), Optional.of(Instant.ofEpochMilli(456)))
+                NodeType.host, new OsVersionTarget(NodeType.host, Version.fromString("1.2.3"), Duration.ofHours(1), Optional.of(Instant.ofEpochMilli(123))),
+                NodeType.proxyhost, new OsVersionTarget(NodeType.proxyhost, Version.fromString("4.5.6"), Duration.ofHours(2), Optional.empty()),
+                NodeType.confighost, new OsVersionTarget(NodeType.confighost, Version.fromString("7.8.9"), Duration.ZERO, Optional.of(Instant.ofEpochMilli(456)))
         ));
         var serialized = OsVersionChangeSerializer.fromJson(OsVersionChangeSerializer.toJson(change));
         assertEquals(serialized, change);
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
index e0715702722..a17fcdd7ff1 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/restapi/NodesV2ApiTest.java
@@ -718,13 +718,13 @@ public class NodesV2ApiTest {
 
         // Upgrade OS for confighost and host
         assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/confighost",
-                                   Utf8.toBytes("{\"osVersion\": \"7.5.2\"}"),
+                                   Utf8.toBytes("{\"osVersion\": \"7.5.2\", \"upgradeBudget\": \"PT0S\"}"),
                                    Request.Method.PATCH),
-                       "{\"message\":\"Set osVersion to 7.5.2 for nodes of type confighost\"}");
+                       "{\"message\":\"Set osVersion to 7.5.2, upgradeBudget to PT0S for nodes of type confighost\"}");
         assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/host",
-                                   Utf8.toBytes("{\"osVersion\": \"7.5.2\"}"),
+                                   Utf8.toBytes("{\"osVersion\": \"7.5.2\", \"upgradeBudget\": \"PT0S\"}"),
                                    Request.Method.PATCH),
-                       "{\"message\":\"Set osVersion to 7.5.2 for nodes of type host\"}");
+                       "{\"message\":\"Set osVersion to 7.5.2, upgradeBudget to PT0S for nodes of type host\"}");
 
         // OS versions are set
         assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/"),
@@ -732,29 +732,29 @@ public class NodesV2ApiTest {
 
         // Upgrade OS and Vespa together
         assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/confighost",
-                                   Utf8.toBytes("{\"version\": \"6.124.42\", \"osVersion\": \"7.5.2\"}"),
+                                   Utf8.toBytes("{\"version\": \"6.124.42\", \"osVersion\": \"7.5.2\", \"upgradeBudget\": \"PT0S\"}"),
                                    Request.Method.PATCH),
-                       "{\"message\":\"Set version to 6.124.42, osVersion to 7.5.2 for nodes of type confighost\"}");
+                       "{\"message\":\"Set version to 6.124.42, osVersion to 7.5.2, upgradeBudget to PT0S for nodes of type confighost\"}");
 
         // Attempt to upgrade unsupported node type
         tester.assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/config",
-                                          Utf8.toBytes("{\"osVersion\": \"7.5.2\"}"),
+                                          Utf8.toBytes("{\"osVersion\": \"7.5.2\", \"upgradeBudget\": \"PT0S\"}"),
                                           Request.Method.PATCH),
                               400,
                               "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Node type 'config' does not support OS upgrades\"}");
 
         // Attempt to downgrade OS
         tester.assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/confighost",
-                                          Utf8.toBytes("{\"osVersion\": \"7.4.2\"}"),
+                                          Utf8.toBytes("{\"osVersion\": \"7.4.2\", \"upgradeBudget\": \"PT0S\"}"),
                                           Request.Method.PATCH),
                               400,
                               "{\"error-code\":\"BAD_REQUEST\",\"message\":\"Cannot set target OS version to 7.4.2 without setting 'force', as it's lower than the current version: 7.5.2\"}");
 
         // Downgrading OS with force succeeds
         assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/confighost",
-                                   Utf8.toBytes("{\"osVersion\": \"7.4.2\", \"force\": true}"),
+                                   Utf8.toBytes("{\"osVersion\": \"7.4.2\", \"force\": true, \"upgradeBudget\": \"PT0S\"}"),
                                    Request.Method.PATCH),
-                       "{\"message\":\"Set osVersion to 7.4.2 for nodes of type confighost\"}");
+                       "{\"message\":\"Set osVersion to 7.4.2, upgradeBudget to PT0S for nodes of type confighost\"}");
 
         // Current target is considered bad, remove it
         tester.assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/confighost",
@@ -788,9 +788,9 @@ public class NodesV2ApiTest {
     public void test_os_version() throws Exception {
         // Schedule OS upgrade
         assertResponse(new Request("http://localhost:8080/nodes/v2/upgrade/host",
-                                   Utf8.toBytes("{\"osVersion\": \"7.5.2\"}"),
+                                   Utf8.toBytes("{\"osVersion\": \"7.5.2\", \"upgradeBudget\": \"PT0S\"}"),
                                    Request.Method.PATCH),
-                       "{\"message\":\"Set osVersion to 7.5.2 for nodes of type host\"}");
+                       "{\"message\":\"Set osVersion to 7.5.2, upgradeBudget to PT0S for nodes of type host\"}");
 
         // Activate target
         var nodeRepository = (NodeRepository)tester.container().components().getComponent(MockNodeRepository.class.getName());