Remove Cloud dependency from OsUpgrader

10 files changed, 70 insertions, 78 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java
index 12a7520129a..15ab14e3241 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/Controller.java
@@ -6,7 +6,6 @@ import com.yahoo.component.AbstractComponent;
 import com.yahoo.component.Version;
 import com.yahoo.component.Vtag;
 import com.yahoo.concurrent.maintenance.JobControl;
-import com.yahoo.config.provision.Cloud;
 import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.SystemName;
@@ -205,11 +204,11 @@ public class Controller extends AbstractComponent {
         if (version.isEmpty()) {
             throw new IllegalArgumentException("Invalid version '" + version.toFullString() + "'");
         }
-        Cloud cloud = zoneRegistry.cloud(cloudName);
-        if (cloud == null) {
+        Set<CloudName> clouds = clouds();
+        if (!clouds.contains(cloudName)) {
             throw new IllegalArgumentException("Cloud '" + cloudName + "' does not exist in this system");
         }
-        if (cloud.reprovisionToUpgradeOs() && upgradeBudget.isEmpty()) {
+        if (!zoneRegistry.zones().ofCloud(cloudName).reprovisionToUpgradeOs().ids().isEmpty() && upgradeBudget.isEmpty()) {
             throw new IllegalArgumentException("Cloud '" + cloudName.value() + "' requires a time budget for OS upgrades");
         }
         try (Lock lock = curator.lockOsVersions()) {
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/SystemApplication.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/SystemApplication.java
index 27a97f1ef01..243bca8c027 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/SystemApplication.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/application/SystemApplication.java
@@ -3,7 +3,6 @@ package com.yahoo.vespa.hosted.controller.application;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
-import com.yahoo.config.provision.Cloud;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.Controller;
@@ -65,8 +64,8 @@ public enum SystemApplication {
     }
 
     /** Returns whether this should receive OS upgrades in given cloud */
-    public boolean shouldUpgradeOsIn(Cloud cloud) {
-        if (cloud.reprovisionToUpgradeOs()) {
+    public boolean shouldUpgradeOsIn(ZoneId zone, Controller controller) {
+        if (controller.zoneRegistry().zones().reprovisionToUpgradeOs().ids().contains(zone)) {
             return nodeType == NodeType.host; // TODO(mpolden): Remove once all node types are supported
         }
         return nodeType.isDockerHost();
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java
index 193c3833202..6f3e868dc1a 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/ControllerMaintenance.java
@@ -105,7 +105,7 @@ public class ControllerMaintenance extends AbstractComponent {
                          .map(ZoneApi::getCloudName)
                          .distinct()
                          .sorted()
-                         .map(cloud -> new OsUpgrader(controller, Duration.ofMinutes(1), controller.zoneRegistry().cloud(cloud)))
+                         .map(cloud -> new OsUpgrader(controller, Duration.ofMinutes(1), cloud))
                          .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
     }
 
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
index ecb9a872676..66fb20e8a71 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgrader.java
@@ -2,7 +2,6 @@
 package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
-import com.yahoo.config.provision.Cloud;
 import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.zone.ZoneApi;
@@ -31,10 +30,10 @@ public class OsUpgrader extends InfrastructureUpgrader<OsVersionTarget> {
             Node.State.reserved
     );
 
-    private final Cloud cloud;
+    private final CloudName cloud;
 
-    public OsUpgrader(Controller controller, Duration interval, Cloud cloud) {
-        super(controller, interval, controller.zoneRegistry().osUpgradePolicy(cloud.name()), name(cloud.name()));
+    public OsUpgrader(Controller controller, Duration interval, CloudName cloud) {
+        super(controller, interval, controller.zoneRegistry().osUpgradePolicy(cloud), name(cloud));
         this.cloud = cloud;
     }
 
@@ -58,22 +57,22 @@ public class OsUpgrader extends InfrastructureUpgrader<OsVersionTarget> {
 
     @Override
     protected boolean expectUpgradeOf(Node node, SystemApplication application, ZoneApi zone) {
-        return cloud.name().equals(zone.getCloudName()) && // Cloud is managed by this upgrader
-               application.shouldUpgradeOsIn(cloud) &&     // Application should upgrade in this cloud
-               canUpgrade(node);                           // Node is in an upgradable state
+        return cloud.equals(zone.getCloudName()) &&                  // Cloud is managed by this upgrader
+               application.shouldUpgradeOsIn(zone.getId(), controller()) && // Application should upgrade in this cloud
+               canUpgrade(node);                                            // Node is in an upgradable state
     }
 
     @Override
     protected Optional<OsVersionTarget> targetVersion() {
         // Return target if we have nodes in this cloud on a lower version
-        return controller().osVersionTarget(cloud.name())
-                           .filter(target -> controller().osVersionStatus().nodesIn(cloud.name()).stream()
+        return controller().osVersionTarget(cloud)
+                           .filter(target -> controller().osVersionStatus().nodesIn(cloud).stream()
                                                          .anyMatch(node -> node.currentVersion().isBefore(target.osVersion().version())));
     }
 
     @Override
     protected boolean changeTargetTo(OsVersionTarget target, SystemApplication application, ZoneApi zone) {
-        if (!application.shouldUpgradeOsIn(cloud)) return false;
+        if (!application.shouldUpgradeOsIn(zone.getId(), controller())) return false;
         return controller().serviceRegistry().configServer().nodeRepository()
                            .targetVersionsOf(zone.getId())
                            .osVersion(application.nodeType())
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java
index ad8a6611e48..c505dbfe1c6 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/versions/OsVersionStatus.java
@@ -63,8 +63,7 @@ public class OsVersionStatus {
 
         for (var application : SystemApplication.all()) {
             for (var zone : zonesToUpgrade(controller)) {
-                var cloud = controller.zoneRegistry().cloud(zone.getCloudName());
-                if (!application.shouldUpgradeOsIn(cloud)) continue;
+                if (!application.shouldUpgradeOsIn(zone.getId(), controller)) continue;
                 var targetOsVersion = controller.serviceRegistry().configServer().nodeRepository()
                                                 .targetVersionsOf(zone.getId())
                                                 .osVersion(application.nodeType())
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneFilterMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneFilterMock.java
index 23f59683f4b..841624cb011 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneFilterMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/ZoneFilterMock.java
@@ -26,21 +26,23 @@ public class ZoneFilterMock implements ZoneList {
 
     private final List<ZoneApi> zones;
     private final Map<ZoneApi, List<RoutingMethod>> zoneRoutingMethods;
+    private final Set<ZoneApi> reprovisionToUpgradeOs;
     private final boolean negate;
 
-    private ZoneFilterMock(List<ZoneApi> zones, Map<ZoneApi, List<RoutingMethod>> zoneRoutingMethods, boolean negate) {
+    private ZoneFilterMock(List<ZoneApi> zones, Map<ZoneApi, List<RoutingMethod>> zoneRoutingMethods, Set<ZoneApi> reprovisionToUpgradeOs, boolean negate) {
         this.zones = zones;
         this.zoneRoutingMethods = zoneRoutingMethods;
+        this.reprovisionToUpgradeOs = reprovisionToUpgradeOs;
         this.negate = negate;
     }
 
-    public static ZoneFilter from(Collection<? extends ZoneApi> zones, Map<ZoneApi, List<RoutingMethod>> routingMethods) {
-        return new ZoneFilterMock(List.copyOf(zones), Map.copyOf(routingMethods), false);
+    public static ZoneFilter from(Collection<? extends ZoneApi> zones, Map<ZoneApi, List<RoutingMethod>> routingMethods, Set<ZoneApi> reprovisionToUpgradeOs) {
+        return new ZoneFilterMock(List.copyOf(zones), Map.copyOf(routingMethods), reprovisionToUpgradeOs, false);
     }
 
     @Override
     public ZoneList not() {
-        return new ZoneFilterMock(zones, zoneRoutingMethods, ! negate);
+        return new ZoneFilterMock(zones, zoneRoutingMethods, reprovisionToUpgradeOs, ! negate);
     }
 
     @Override
@@ -69,6 +71,11 @@ public class ZoneFilterMock implements ZoneList {
     }
 
     @Override
+    public ZoneList reprovisionToUpgradeOs() {
+        return filter(reprovisionToUpgradeOs::contains);
+    }
+
+    @Override
     public ZoneList in(Environment... environments) {
         return filter(zone -> Set.of(environments).contains(zone.getEnvironment()));
     }
@@ -100,7 +107,7 @@ public class ZoneFilterMock implements ZoneList {
                                 condition.negate().test(zone) :
                                 condition.test(zone))
                         .collect(Collectors.toList()),
-                zoneRoutingMethods, false);
+                zoneRoutingMethods, reprovisionToUpgradeOs, false);
     }
 
 }
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 a453a990855..8fcd9527517 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
@@ -4,7 +4,6 @@ package com.yahoo.vespa.hosted.controller.integration;
 import com.yahoo.component.AbstractComponent;
 import com.yahoo.config.provision.ApplicationId;
 import com.yahoo.config.provision.AthenzDomain;
-import com.yahoo.config.provision.Cloud;
 import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.NodeType;
@@ -24,9 +23,11 @@ import com.yahoo.vespa.hosted.controller.api.integration.zone.ZoneRegistry;
 import java.net.URI;
 import java.time.Duration;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 
 /**
  * @author mpolden
@@ -37,7 +38,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 final Set<ZoneApi> reprovisionToUpgradeOs = new HashSet<>();
 
     private List<? extends ZoneApi> zones;
     private SystemName system;
@@ -62,8 +63,6 @@ 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);
     }
@@ -102,13 +101,6 @@ 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));
     }
@@ -134,6 +126,15 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
         return this;
     }
 
+    public ZoneRegistryMock reprovisionToUpgradeOsIn(ZoneApi... zones) {
+        return reprovisionToUpgradeOsIn(List.of(zones));
+    }
+
+    public ZoneRegistryMock reprovisionToUpgradeOsIn(List<ZoneApi> zones) {
+        this.reprovisionToUpgradeOs.addAll(zones);
+        return this;
+    }
+
     @Override
     public SystemName system() {
         return system;
@@ -141,7 +142,7 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
 
     @Override
     public ZoneFilter zones() {
-        return ZoneFilterMock.from(zones, zoneRoutingMethods);
+        return ZoneFilterMock.from(zones, zoneRoutingMethods, reprovisionToUpgradeOs);
     }
 
     @Override
@@ -210,11 +211,6 @@ public class ZoneRegistryMock extends AbstractComponent implements ZoneRegistry
     }
 
     @Override
-    public Cloud cloud(CloudName name) {
-        return clouds.get(name);
-    }
-
-    @Override
     public boolean hasZone(ZoneId zoneId) {
         return zones.stream().anyMatch(zone -> zone.getId().equals(zoneId));
     }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java
index 6eeabc98932..ab274dc5ef8 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/MetricsReporterTest.java
@@ -3,7 +3,6 @@ package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.ApplicationId;
-import com.yahoo.config.provision.Cloud;
 import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.HostName;
@@ -352,7 +351,7 @@ public class MetricsReporterTest {
         var zone = ZoneId.from("prod.eu-west-1");
         var cloud = CloudName.defaultName();
         tester.zoneRegistry().setOsUpgradePolicy(cloud, UpgradePolicy.create().upgrade(ZoneApiMock.from(zone)));
-        var osUpgrader = new OsUpgrader(tester.controller(), Duration.ofDays(1), Cloud.defaultCloud());
+        var osUpgrader = new OsUpgrader(tester.controller(), Duration.ofDays(1), CloudName.defaultName());
         var statusUpdater = new OsVersionStatusUpdater(tester.controller(), Duration.ofDays(1)
         );
         tester.configServer().bootstrap(List.of(zone), SystemApplication.configServerHost, SystemApplication.tenantHost);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
index 0bc29157c06..cb906d61a3b 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsUpgraderTest.java
@@ -2,7 +2,6 @@
 package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
-import com.yahoo.config.provision.Cloud;
 import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.zone.UpgradePolicy;
@@ -37,8 +36,8 @@ public class OsUpgraderTest {
 
     @Test
     public void upgrade_os() {
-        Cloud cloud1 = Cloud.builder().name(CloudName.from("c1")).requireAccessControl(true).build();
-        Cloud cloud2 = Cloud.builder().name(CloudName.from("c2")).requireAccessControl(true).build();
+        CloudName cloud1 = CloudName.from("c1");
+        CloudName cloud2 = CloudName.from("c2");
         ZoneApi zone1 = zone("prod.eu-west-1", cloud1);
         ZoneApi zone2 = zone("prod.us-west-1", cloud1);
         ZoneApi zone3 = zone("prod.us-central-1", cloud1);
@@ -49,7 +48,7 @@ public class OsUpgraderTest {
                                                    .upgradeInParallel(zone2, zone3)
                                                    .upgrade(zone5) // Belongs to a different cloud and is ignored by this upgrader
                                                    .upgrade(zone4);
-        OsUpgrader osUpgrader = osUpgrader(upgradePolicy, SystemName.cd, cloud1, cloud2);
+        OsUpgrader osUpgrader = osUpgrader(upgradePolicy, SystemName.cd, cloud1, false);
 
         // Bootstrap system
         tester.configServer().bootstrap(List.of(zone1.getId(), zone2.getId(), zone3.getId(), zone4.getId(), zone5.getId()),
@@ -65,8 +64,8 @@ public class OsUpgraderTest {
 
         // New OS version released
         Version version1 = Version.fromString("7.1");
-        tester.controller().upgradeOsIn(cloud1.name(), Version.fromString("7.0"), Optional.empty(), false);
-        tester.controller().upgradeOsIn(cloud1.name(), version1, Optional.empty(), false);
+        tester.controller().upgradeOsIn(cloud1, Version.fromString("7.0"), Optional.empty(), false);
+        tester.controller().upgradeOsIn(cloud1, version1, Optional.empty(), false);
         assertEquals(1, tester.controller().osVersionTargets().size()); // Only allows one version per cloud
         statusUpdater.maintain();
 
@@ -104,13 +103,13 @@ public class OsUpgraderTest {
         osUpgrader.maintain();
         assertWanted(version1, SystemApplication.tenantHost, zone1.getId(), zone2.getId(), zone3.getId(), zone4.getId());
         statusUpdater.maintain();
-        assertTrue("All nodes on target version", tester.controller().osVersionStatus().nodesIn(cloud1.name()).stream()
+        assertTrue("All nodes on target version", tester.controller().osVersionStatus().nodesIn(cloud1).stream()
                                                         .allMatch(node -> node.currentVersion().equals(version1)));
     }
 
     @Test
     public void upgrade_os_with_budget() {
-        Cloud cloud = Cloud.builder().name(CloudName.from("cloud")).reprovisionToUpgradeOs(true).requireAccessControl(true).build();
+        CloudName cloud = CloudName.from("cloud");
         ZoneApi zone1 = zone("dev.us-east-1", cloud);
         ZoneApi zone2 = zone("prod.us-west-1", cloud);
         ZoneApi zone3 = zone("prod.us-central-1", cloud);
@@ -119,7 +118,7 @@ public class OsUpgraderTest {
                                                    .upgrade(zone1)
                                                    .upgradeInParallel(zone2, zone3)
                                                    .upgrade(zone4);
-        OsUpgrader osUpgrader = osUpgrader(upgradePolicy, SystemName.cd, cloud);
+        OsUpgrader osUpgrader = osUpgrader(upgradePolicy, SystemName.cd, cloud, true);
 
         // Bootstrap system
         tester.configServer().bootstrap(List.of(zone1.getId(), zone2.getId(), zone3.getId(), zone4.getId()),
@@ -130,13 +129,13 @@ public class OsUpgraderTest {
         // Upgrade without budget fails
         Version version = Version.fromString("7.1");
         try {
-            tester.controller().upgradeOsIn(cloud.name(), version, Optional.empty(), false);
+            tester.controller().upgradeOsIn(cloud, version, Optional.empty(), false);
             fail("Expected exception");
         } catch (IllegalArgumentException ignored) {}
 
         // Upgrade with budget
-        tester.controller().upgradeOsIn(cloud.name(), version, Optional.of(Duration.ofHours(12)), false);
-        assertEquals(Duration.ofHours(12), tester.controller().osVersionTarget(cloud.name()).get().upgradeBudget().get());
+        tester.controller().upgradeOsIn(cloud, version, Optional.of(Duration.ofHours(12)), false);
+        assertEquals(Duration.ofHours(12), tester.controller().osVersionTarget(cloud).get().upgradeBudget().get());
         statusUpdater.maintain();
         osUpgrader.maintain();
 
@@ -161,7 +160,7 @@ public class OsUpgraderTest {
 
         // All host applications upgraded
         statusUpdater.maintain();
-        assertTrue("All nodes on target version", tester.controller().osVersionStatus().nodesIn(cloud.name()).stream()
+        assertTrue("All nodes on target version", tester.controller().osVersionStatus().nodesIn(cloud).stream()
                                                         .allMatch(node -> node.currentVersion().equals(version)));
     }
 
@@ -230,19 +229,20 @@ public class OsUpgraderTest {
         return tester.configServer().nodeRepository();
     }
 
-    private OsUpgrader osUpgrader(UpgradePolicy upgradePolicy, SystemName system, Cloud managedCloud, Cloud... otherClouds) {
+    private OsUpgrader osUpgrader(UpgradePolicy upgradePolicy, SystemName system, CloudName cloud, boolean reprovisionToUpgradeOs) {
         var zones = upgradePolicy.asList().stream().flatMap(Collection::stream).collect(Collectors.toList());
         tester.zoneRegistry()
               .setZones(zones)
-              .addCloud(managedCloud)
-              .addCloud(otherClouds)
               .setSystemName(system)
-              .setOsUpgradePolicy(managedCloud.name(), upgradePolicy);
-        return new OsUpgrader(tester.controller(), Duration.ofDays(1), managedCloud);
+              .setOsUpgradePolicy(cloud, upgradePolicy);
+        if (reprovisionToUpgradeOs) {
+            tester.zoneRegistry().reprovisionToUpgradeOsIn(zones);
+        }
+        return new OsUpgrader(tester.controller(), Duration.ofDays(1), cloud);
     }
 
-    private static ZoneApi zone(String id, Cloud cloud) {
-        return ZoneApiMock.newBuilder().withId(id).withCloud(cloud.name().value()).build();
+    private static ZoneApi zone(String id, CloudName cloud) {
+        return ZoneApiMock.newBuilder().withId(id).with(cloud).build();
     }
 
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
index 8131a850d3b..857530df1fe 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/OsApiTest.java
@@ -2,7 +2,6 @@
 package com.yahoo.vespa.hosted.controller.restapi.os;
 
 import com.yahoo.application.container.handler.Request;
-import com.yahoo.config.provision.Cloud;
 import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.zone.UpgradePolicy;
@@ -37,16 +36,11 @@ public class OsApiTest extends ControllerContainerTest {
 
     private static final String responses = "src/test/java/com/yahoo/vespa/hosted/controller/restapi/os/responses/";
     private static final AthenzIdentity operator = AthenzUser.fromUserId("operatorUser");
-    private static final Cloud cloud1 = Cloud.builder().name(CloudName.from("cloud1")).build();
-    private static final Cloud cloud2 = Cloud.builder().name(CloudName.from("cloud2"))
-                                             .dynamicProvisioning(true)
-                                             .allowHostSharing(false)
-                                             .reprovisionToUpgradeOs(true)
-                                             .requireAccessControl(true)
-                                             .build();
-    private static final ZoneApi zone1 = ZoneApiMock.newBuilder().withId("prod.us-east-3").with(cloud1.name()).build();
-    private static final ZoneApi zone2 = ZoneApiMock.newBuilder().withId("prod.us-west-1").with(cloud1.name()).build();
-    private static final ZoneApi zone3 = ZoneApiMock.newBuilder().withId("prod.eu-west-1").with(cloud2.name()).build();
+    private static final CloudName cloud1 = CloudName.from("cloud1");
+    private static final CloudName cloud2 = CloudName.from("cloud2");
+    private static final ZoneApi zone1 = ZoneApiMock.newBuilder().withId("prod.us-east-3").with(cloud1).build();
+    private static final ZoneApi zone2 = ZoneApiMock.newBuilder().withId("prod.us-west-1").with(cloud1).build();
+    private static final ZoneApi zone3 = ZoneApiMock.newBuilder().withId("prod.eu-west-1").with(cloud2).build();
 
     private ContainerTester tester;
     private List<OsUpgrader> osUpgraders;
@@ -57,9 +51,9 @@ public class OsApiTest extends ControllerContainerTest {
         addUserToHostedOperatorRole(operator);
         zoneRegistryMock().setSystemName(SystemName.cd)
                           .setZones(zone1, zone2, zone3)
-                          .addCloud(cloud1, cloud2)
-                          .setOsUpgradePolicy(cloud1.name(), UpgradePolicy.create().upgrade(zone1).upgrade(zone2))
-                          .setOsUpgradePolicy(cloud2.name(), UpgradePolicy.create().upgrade(zone3));
+                          .reprovisionToUpgradeOsIn(zone3)
+                          .setOsUpgradePolicy(cloud1, UpgradePolicy.create().upgrade(zone1).upgrade(zone2))
+                          .setOsUpgradePolicy(cloud2, UpgradePolicy.create().upgrade(zone3));
         osUpgraders = List.of(
                 new OsUpgrader(tester.controller(), Duration.ofDays(1),
                                cloud1),