Send OS upgrade budget in applicable clouds

1 files changed, 11 insertions, 1 deletions

diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java
index 902af26e5ba..a453a990855 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneRegistryMock.java
@@ -37,6 +37,7 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
     private final Map<Environment, RegionName> defaultRegionForEnvironment = new HashMap<>();
     private final Map<CloudName, UpgradePolicy> osUpgradePolicies = new HashMap<>();
     private final Map<ZoneApi, List<RoutingMethod>> zoneRoutingMethods = new HashMap<>();
+    private final Map<CloudName, Cloud> clouds = new HashMap<>();
 
     private List<? extends ZoneApi> zones;
     private SystemName system;
@@ -61,6 +62,8 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
                              ZoneApiMock.fromId("prod.us-west-1"),
                              ZoneApiMock.fromId("prod.us-central-1"),
                              ZoneApiMock.fromId("prod.eu-west-1"));
+        var cloud = Cloud.defaultCloud();
+        this.clouds.put(cloud.name(), cloud);
         // All zones use a shared routing method by default
         setRoutingMethod(this.zones, RoutingMethod.shared);
     }
@@ -99,6 +102,13 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
         return this;
     }
 
+    public ZoneRegistryMock addCloud(Cloud... clouds) {
+        for (var cloud : clouds) {
+            this.clouds.put(cloud.name(), cloud);
+        }
+        return this;
+    }
+
     public ZoneRegistryMock exclusiveRoutingIn(ZoneApi... zones) {
         return exclusiveRoutingIn(List.of(zones));
     }
@@ -201,7 +211,7 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
 
     @Override
     public Cloud cloud(CloudName name) {
-        return new Cloud(name, false, true, false, false);
+        return clouds.get(name);
     }
 
     @Override