Merge pull request #9738 from vespa-engine/hakonhall/use-zoneapi-in-upgraders

Use ZoneApi instead of ZoneId in UpgradePolicy

9 files changed, 81 insertions, 74 deletions

diff --git a/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java b/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java
index d5831efdbaa..1baaac772c8 100644
--- a/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java
+++ b/config-provisioning/src/main/java/com/yahoo/config/provision/zone/UpgradePolicy.java
@@ -13,29 +13,29 @@ import java.util.List;
  */
 public class UpgradePolicy {
 
-    private final List<List<ZoneId>> zones;
+    private final List<List<ZoneApi>> zones;
 
-    private UpgradePolicy(List<List<ZoneId>> zones) {
+    private UpgradePolicy(List<List<ZoneApi>> zones) {
         this.zones = zones;
     }
 
-    public List<List<ZoneId>> asList() {
-        return Collections.unmodifiableList(zones);
+    public List<List<ZoneApi>> asList() {
+        return List.copyOf(zones);
     }
 
-    private UpgradePolicy with(ZoneId... zone) {
-        List<List<ZoneId>> zones = new ArrayList<>(this.zones);
+    private UpgradePolicy with(ZoneApi... zone) {
+        List<List<ZoneApi>> zones = new ArrayList<>(this.zones);
         zones.add(Arrays.asList(zone));
         return new UpgradePolicy(zones);
     }
 
     /** Upgrade given zone as the next step */
-    public UpgradePolicy upgrade(ZoneId zone) {
+    public UpgradePolicy upgrade(ZoneApi zone) {
         return with(zone);
     }
 
     /** Upgrade given zones in parallel as the next step */
-    public UpgradePolicy upgradeInParallel(ZoneId... zone) {
+    public UpgradePolicy upgradeInParallel(ZoneApi... zone) {
         return with(zone);
     }
 
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/InfrastructureUpgrader.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/InfrastructureUpgrader.java
index 8818a441fbd..159eb234aa7 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/InfrastructureUpgrader.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/InfrastructureUpgrader.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.SystemName;
+import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.vespa.hosted.controller.Controller;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.Node;
 import com.yahoo.config.provision.zone.UpgradePolicy;
@@ -42,9 +43,9 @@ public abstract class InfrastructureUpgrader extends Maintainer {
 
     /** Deploy a list of system applications until they converge on the given version */
     private void upgradeAll(Version target, List<SystemApplication> applications) {
-        for (List<ZoneId> zones : upgradePolicy.asList()) {
+        for (List<ZoneApi> zones : upgradePolicy.asList()) {
             boolean converged = true;
-            for (ZoneId zone : zones) {
+            for (ZoneApi zone : zones) {
                 try {
                     converged &= upgradeAll(target, applications, zone);
                 } catch (UnreachableNodeRepositoryException e) {
@@ -62,7 +63,7 @@ public abstract class InfrastructureUpgrader extends Maintainer {
     }
 
     /** Returns whether all applications have converged to the target version in zone */
-    private boolean upgradeAll(Version target, List<SystemApplication> applications, ZoneId zone) {
+    private boolean upgradeAll(Version target, List<SystemApplication> applications, ZoneApi zone) {
         boolean converged = true;
         for (SystemApplication application : applications) {
             if (convergedOn(target, application.dependencies(), zone)) {
@@ -76,28 +77,28 @@ public abstract class InfrastructureUpgrader extends Maintainer {
         return converged;
     }
 
-    private boolean convergedOn(Version target, List<SystemApplication> applications, ZoneId zone) {
+    private boolean convergedOn(Version target, List<SystemApplication> applications, ZoneApi zone) {
         return applications.stream().allMatch(application -> convergedOn(target, application, zone));
     }
 
     /** Upgrade component to target version. Implementation should be idempotent */
-    protected abstract void upgrade(Version target, SystemApplication application, ZoneId zone);
+    protected abstract void upgrade(Version target, SystemApplication application, ZoneApi zone);
 
     /** Returns whether application has converged to target version in zone */
-    protected abstract boolean convergedOn(Version target, SystemApplication application, ZoneId zone);
+    protected abstract boolean convergedOn(Version target, SystemApplication application, ZoneApi zone);
 
     /** Returns the target version for the component upgraded by this, if any */
     protected abstract Optional<Version> targetVersion();
 
     /** Returns whether the upgrader should require given node to upgrade */
-    protected abstract boolean requireUpgradeOf(Node node, SystemApplication application, ZoneId zone);
+    protected abstract boolean requireUpgradeOf(Node node, SystemApplication application, ZoneApi zone);
 
     /** Find the minimum value of a version field in a zone */
-    protected final Optional<Version> minVersion(ZoneId zone, SystemApplication application, Function<Node, Version> versionField) {
+    protected final Optional<Version> minVersion(ZoneApi zone, SystemApplication application, Function<Node, Version> versionField) {
         try {
             return controller().configServer()
                                .nodeRepository()
-                               .list(zone, application.id())
+                               .list(zone.getId(), application.id())
                                .stream()
                                .filter(node -> requireUpgradeOf(node, application, zone))
                                .map(versionField)
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 ed3dd552085..8845f4c652f 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
@@ -4,9 +4,9 @@ package com.yahoo.vespa.hosted.controller.maintenance;
 import com.google.common.collect.ImmutableSet;
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.CloudName;
+import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.vespa.hosted.controller.Controller;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.Node;
-import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.application.SystemApplication;
 import com.yahoo.vespa.hosted.controller.versions.OsVersion;
 
@@ -38,22 +38,22 @@ public class OsUpgrader extends InfrastructureUpgrader {
     }
 
     @Override
-    protected void upgrade(Version target, SystemApplication application, ZoneId zone) {
+    protected void upgrade(Version target, SystemApplication application, ZoneApi zone) {
         if (!application.isEligibleForOsUpgrades() || wantedVersion(zone, application, target).equals(target)) {
             return;
         }
-        log.info(String.format("Upgrading OS of %s to version %s in %s", application.id(), target, zone));
-        controller().configServer().nodeRepository().upgradeOs(zone, application.nodeType(), target);
+        log.info(String.format("Upgrading OS of %s to version %s in %s in cloud %s", application.id(), target, zone.getId(), zone.getCloudName()));
+        controller().configServer().nodeRepository().upgradeOs(zone.getId(), application.nodeType(), target);
     }
 
     @Override
-    protected boolean convergedOn(Version target, SystemApplication application, ZoneId zone) {
+    protected boolean convergedOn(Version target, SystemApplication application, ZoneApi zone) {
         return currentVersion(zone, application, target).equals(target);
     }
 
     @Override
-    protected boolean requireUpgradeOf(Node node, SystemApplication application, ZoneId zone) {
-        return cloud.equals(zone.cloud()) && eligibleForUpgrade(node, application);
+    protected boolean requireUpgradeOf(Node node, SystemApplication application, ZoneApi zone) {
+        return cloud.equals(zone.getCloudName()) && eligibleForUpgrade(node, application);
     }
 
     @Override
@@ -65,11 +65,11 @@ public class OsUpgrader extends InfrastructureUpgrader {
                            .map(OsVersion::version);
     }
 
-    private Version currentVersion(ZoneId zone, SystemApplication application, Version defaultVersion) {
+    private Version currentVersion(ZoneApi zone, SystemApplication application, Version defaultVersion) {
         return minVersion(zone, application, Node::currentOsVersion).orElse(defaultVersion);
     }
 
-    private Version wantedVersion(ZoneId zone, SystemApplication application, Version defaultVersion) {
+    private Version wantedVersion(ZoneApi zone, SystemApplication application, Version defaultVersion) {
         return minVersion(zone, application, Node::wantedOsVersion).orElse(defaultVersion);
     }
 
diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgrader.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgrader.java
index 8ef353248f8..5a40dd591fd 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgrader.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgrader.java
@@ -2,7 +2,7 @@
 package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
-import com.yahoo.config.provision.zone.ZoneId;
+import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.vespa.hosted.controller.Controller;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.Node;
 import com.yahoo.vespa.hosted.controller.application.SystemApplication;
@@ -30,26 +30,26 @@ public class SystemUpgrader extends InfrastructureUpgrader {
     }
 
     @Override
-    protected void upgrade(Version target, SystemApplication application, ZoneId zone) {
+    protected void upgrade(Version target, SystemApplication application, ZoneApi zone) {
         if (minVersion(zone, application, Node::wantedVersion).map(target::isAfter)
                                                               .orElse(true)) {
-            log.info(String.format("Deploying %s version %s in %s", application.id(), target, zone));
-            controller().applications().deploy(application, zone, target);
+            log.info(String.format("Deploying %s version %s in %s", application.id(), target, zone.getId()));
+            controller().applications().deploy(application, zone.toDeprecatedId(), target);
         }
     }
 
     @Override
-    protected boolean convergedOn(Version target, SystemApplication application, ZoneId zone) {
+    protected boolean convergedOn(Version target, SystemApplication application, ZoneApi zone) {
         Optional<Version> minVersion = minVersion(zone, application, Node::currentVersion);
         // Skip application convergence check if there are no nodes belonging to the application in the zone
         if (minVersion.isEmpty()) return true;
 
         return     minVersion.get().equals(target)
-                && application.configConvergedIn(zone, controller(), Optional.of(target));
+                && application.configConvergedIn(zone.toDeprecatedId(), controller(), Optional.of(target));
     }
 
     @Override
-    protected boolean requireUpgradeOf(Node node, SystemApplication application, ZoneId zone) {
+    protected boolean requireUpgradeOf(Node node, SystemApplication application, ZoneApi zone) {
         return eligibleForUpgrade(node);
     }
 
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 f9ce75d2297..f5b9d8263e5 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
@@ -7,6 +7,7 @@ import com.yahoo.config.provision.CloudName;
 import com.yahoo.config.provision.Environment;
 import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.RegionName;
+import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.vespa.hosted.controller.Controller;
 import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.application.SystemApplication;
@@ -70,12 +71,12 @@ public class OsVersionStatus {
             if (!application.isEligibleForOsUpgrades()) {
                 continue; // Avoid querying applications that are not eligible for OS upgrades
             }
-            for (ZoneId zone : zonesToUpgrade(controller)) {
-                controller.configServer().nodeRepository().list(zone, application.id()).stream()
+            for (ZoneApi zone : zonesToUpgrade(controller)) {
+                controller.configServer().nodeRepository().list(zone.getId(), application.id()).stream()
                           .filter(node -> OsUpgrader.eligibleForUpgrade(node, application))
-                          .map(node -> new Node(node.hostname(), node.currentOsVersion(), zone.environment(), zone.region()))
+                          .map(node -> new Node(node.hostname(), node.currentOsVersion(), zone.getEnvironment(), zone.getRegionName()))
                           .forEach(node -> {
-                              var version = new OsVersion(node.version(), zone.cloud());
+                              var version = new OsVersion(node.version(), zone.getCloudName());
                               versions.putIfAbsent(version, new ArrayList<>());
                               versions.get(version).add(node);
                           });
@@ -85,7 +86,7 @@ public class OsVersionStatus {
         return new OsVersionStatus(versions);
     }
 
-    private static List<ZoneId> zonesToUpgrade(Controller controller) {
+    private static List<ZoneApi> zonesToUpgrade(Controller controller) {
         return controller.zoneRegistry().osUpgradePolicies().stream()
                          .flatMap(upgradePolicy -> upgradePolicy.asList().stream())
                          .flatMap(Collection::stream)
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 33bffaae96a..39ff29f4ae0 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
@@ -49,10 +49,10 @@ public class OsUpgraderTest {
     public void upgrade_os() {
         OsUpgrader osUpgrader = osUpgrader(
                 UpgradePolicy.create()
-                             .upgrade(zone1.toDeprecatedId())
-                             .upgradeInParallel(zone2.toDeprecatedId(), zone3.toDeprecatedId())
-                             .upgrade(zone5.toDeprecatedId()) // Belongs to a different cloud and is ignored by this upgrader
-                             .upgrade(zone4.toDeprecatedId()),
+                             .upgrade(zone1)
+                             .upgradeInParallel(zone2, zone3)
+                             .upgrade(zone5) // Belongs to a different cloud and is ignored by this upgrader
+                             .upgrade(zone4),
                 SystemName.cd
         );
 
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
index d8e6f573592..fe7f39fd66d 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/OsVersionStatusUpdaterTest.java
@@ -3,6 +3,7 @@ package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.CloudName;
+import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.vespa.hosted.controller.ControllerTester;
 import com.yahoo.config.provision.zone.UpgradePolicy;
 import com.yahoo.config.provision.zone.ZoneId;
@@ -32,7 +33,7 @@ public class OsVersionStatusUpdaterTest {
                                                                           new JobControl(new MockCuratorDb()));
         // Add all zones to upgrade policy
         UpgradePolicy upgradePolicy = UpgradePolicy.create();
-        for (ZoneId zone : tester.zoneRegistry().zones().controllerUpgraded().ids()) {
+        for (ZoneApi zone : tester.zoneRegistry().zones().controllerUpgraded().zones()) {
             upgradePolicy = upgradePolicy.upgrade(zone);
         }
         tester.zoneRegistry().setOsUpgradePolicy(CloudName.defaultName(), upgradePolicy);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgraderTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgraderTest.java
index c8fd7e90ccd..cb5e7cc90a1 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgraderTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/maintenance/SystemUpgraderTest.java
@@ -3,11 +3,13 @@ package com.yahoo.vespa.hosted.controller.maintenance;
 
 import com.yahoo.component.Version;
 import com.yahoo.config.provision.zone.UpgradePolicy;
+import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.vespa.hosted.controller.api.integration.configserver.Node;
 import com.yahoo.vespa.hosted.controller.application.SystemApplication;
 import com.yahoo.vespa.hosted.controller.deployment.DeploymentTester;
 import com.yahoo.vespa.hosted.controller.integration.NodeRepositoryMock;
+import com.yahoo.vespa.hosted.controller.integration.ZoneApiMock;
 import com.yahoo.vespa.hosted.controller.versions.VespaVersion;
 import org.junit.Before;
 import org.junit.Test;
@@ -25,10 +27,10 @@ import static org.junit.Assert.assertTrue;
  */
 public class SystemUpgraderTest {
 
-    private static final ZoneId zone1 = ZoneId.from("prod", "eu-west-1");
-    private static final ZoneId zone2 = ZoneId.from("prod", "us-west-1");
-    private static final ZoneId zone3 = ZoneId.from("prod", "us-central-1");
-    private static final ZoneId zone4 = ZoneId.from("prod", "us-east-3");
+    private static final ZoneApi zone1 = ZoneApiMock.fromId("prod.eu-west-1");
+    private static final ZoneApi zone2 = ZoneApiMock.fromId("prod.us-west-1");
+    private static final ZoneApi zone3 = ZoneApiMock.fromId("prod.us-central-1");
+    private static final ZoneApi zone4 = ZoneApiMock.fromId("prod.us-east-3");
 
     private DeploymentTester tester;
 
@@ -48,8 +50,8 @@ public class SystemUpgraderTest {
 
         Version version1 = Version.fromString("6.5");
         // Bootstrap a system without host applications
-        tester.configServer().bootstrap(List.of(zone1, zone2, zone3, zone4), SystemApplication.configServer,
-                                        SystemApplication.proxy);
+        tester.configServer().bootstrap(List.of(zone1.toDeprecatedId(), zone2.toDeprecatedId(), zone3.toDeprecatedId(), zone4.toDeprecatedId()),
+                                        SystemApplication.configServer, SystemApplication.proxy);
         // Fail a few nodes. Failed nodes should not affect versions
         failNodeIn(zone1, SystemApplication.configServer);
         failNodeIn(zone3, SystemApplication.proxy);
@@ -142,7 +144,7 @@ public class SystemUpgraderTest {
         SystemUpgrader systemUpgrader = systemUpgrader(UpgradePolicy.create().upgrade(zone1));
 
         // Bootstrap system
-        tester.configServer().bootstrap(List.of(zone1), SystemApplication.configServer,
+        tester.configServer().bootstrap(List.of(zone1.toDeprecatedId()), SystemApplication.configServer,
                                         SystemApplication.proxy);
         Version version1 = Version.fromString("6.5");
         tester.upgradeSystem(version1);
@@ -182,7 +184,7 @@ public class SystemUpgraderTest {
         );
 
         Version version1 = Version.fromString("6.5");
-        tester.configServer().bootstrap(List.of(zone1, zone2, zone3, zone4), SystemApplication.all());
+        tester.configServer().bootstrap(List.of(zone1.toDeprecatedId(), zone2.toDeprecatedId(), zone3.toDeprecatedId(), zone4.toDeprecatedId()), SystemApplication.all());
         tester.upgradeSystem(version1);
         systemUpgrader.maintain();
         assertCurrentVersion(SystemApplication.all(), version1, zone1, zone2, zone3, zone4);
@@ -280,7 +282,7 @@ public class SystemUpgraderTest {
     public void does_not_deploy_proxy_app_in_zones_without_proxy() {
         List<SystemApplication> applications = List.of(
                 SystemApplication.configServerHost, SystemApplication.configServer, SystemApplication.tenantHost);
-        tester.configServer().bootstrap(List.of(zone1), applications);
+        tester.configServer().bootstrap(List.of(zone1.toDeprecatedId()), applications);
         tester.configServer().disallowConvergenceCheck(SystemApplication.proxy.id());
 
         SystemUpgrader systemUpgrader = systemUpgrader(UpgradePolicy.create().upgrade(zone1));
@@ -292,36 +294,38 @@ public class SystemUpgraderTest {
     }
 
     /** Simulate upgrade of nodes allocated to given application. In a real system this is done by the node itself */
-    private void completeUpgrade(SystemApplication application, Version version, ZoneId... zones) {
+    private void completeUpgrade(SystemApplication application, Version version, ZoneApi... zones) {
         assertWantedVersion(application, version, zones);
-        for (ZoneId zone : zones) {
+        for (ZoneApi zone : zones) {
             for (Node node : listNodes(zone, application)) {
-                nodeRepository().putByHostname(zone, new Node(node.hostname(), node.state(), node.type(), node.owner(),
-                                                              node.wantedVersion(), node.wantedVersion()));
+                nodeRepository().putByHostname(
+                        zone.getId(),
+                        new Node(node.hostname(), node.state(), node.type(), node.owner(), node.wantedVersion(), node.wantedVersion()));
             }
 
             assertCurrentVersion(application, version, zone);
         }
     }
 
-    private void convergeServices(SystemApplication application, ZoneId... zones) {
-        for (ZoneId zone : zones) {
-            tester.controllerTester().configServer().convergeServices(application.id(), zone);
+    private void convergeServices(SystemApplication application, ZoneApi... zones) {
+        for (ZoneApi zone : zones) {
+            tester.controllerTester().configServer().convergeServices(application.id(), zone.toDeprecatedId());
         }
     }
 
-    private void completeUpgrade(List<SystemApplication> applications, Version version, ZoneId... zones) {
+    private void completeUpgrade(List<SystemApplication> applications, Version version, ZoneApi... zones) {
         applications.forEach(application -> completeUpgrade(application, version, zones));
     }
 
-    private void failNodeIn(ZoneId zone, SystemApplication application) {
-        List<Node> nodes = nodeRepository().list(zone, application.id());
+    private void failNodeIn(ZoneApi zone, SystemApplication application) {
+        List<Node> nodes = nodeRepository().list(zone.getId(), application.id());
         if (nodes.isEmpty()) {
             throw new IllegalArgumentException("No nodes allocated to " + application.id());
         }
         Node node = nodes.get(0);
-        nodeRepository().putByHostname(zone, new Node(node.hostname(), Node.State.failed, node.type(), node.owner(),
-                                                      node.currentVersion(), node.wantedVersion()));
+        nodeRepository().putByHostname(
+                zone.getId(),
+                new Node(node.hostname(), Node.State.failed, node.type(), node.owner(), node.currentVersion(), node.wantedVersion()));
     }
 
     private void assertSystemVersion(Version version) {
@@ -332,33 +336,33 @@ public class SystemUpgraderTest {
         assertEquals(version, tester.controller().versionStatus().controllerVersion().get().versionNumber());
     }
 
-    private void assertWantedVersion(SystemApplication application, Version version, ZoneId... zones) {
+    private void assertWantedVersion(SystemApplication application, Version version, ZoneApi... zones) {
         assertVersion(application, version, Node::wantedVersion, zones);
     }
 
-    private void assertCurrentVersion(SystemApplication application, Version version, ZoneId... zones) {
+    private void assertCurrentVersion(SystemApplication application, Version version, ZoneApi... zones) {
         assertVersion(application, version, Node::currentVersion, zones);
     }
 
-    private void assertWantedVersion(List<SystemApplication> applications, Version version, ZoneId... zones) {
+    private void assertWantedVersion(List<SystemApplication> applications, Version version, ZoneApi... zones) {
         applications.forEach(application -> assertVersion(application, version, Node::wantedVersion, zones));
     }
 
-    private void assertCurrentVersion(List<SystemApplication> applications, Version version, ZoneId... zones) {
+    private void assertCurrentVersion(List<SystemApplication> applications, Version version, ZoneApi... zones) {
         applications.forEach(application -> assertVersion(application, version, Node::currentVersion, zones));
     }
 
     private void assertVersion(SystemApplication application, Version version, Function<Node, Version> versionField,
-                               ZoneId... zones) {
-        for (ZoneId zone : requireNonEmpty(zones)) {
+                               ZoneApi... zones) {
+        for (ZoneApi zone : requireNonEmpty(zones)) {
             for (Node node : listNodes(zone, application)) {
                 assertEquals(application + " version", version, versionField.apply(node));
             }
         }
     }
 
-    private List<Node> listNodes(ZoneId zone, SystemApplication application) {
-        return nodeRepository().list(zone, application.id()).stream()
+    private List<Node> listNodes(ZoneApi zone, SystemApplication application) {
+        return nodeRepository().list(zone.getId(), application.id()).stream()
                                .filter(SystemUpgrader::eligibleForUpgrade)
                                .collect(Collectors.toList());
     }
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 f7b5a75fe15..b2dfd7b4cb6 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
@@ -53,8 +53,8 @@ public class OsApiTest extends ControllerContainerTest {
         addUserToHostedOperatorRole(operator);
         zoneRegistryMock().setSystemName(SystemName.cd)
                           .setZones(zone1, zone2, zone3)
-                          .setOsUpgradePolicy(cloud1, UpgradePolicy.create().upgrade(zone1.toDeprecatedId()).upgrade(zone2.toDeprecatedId()))
-                          .setOsUpgradePolicy(cloud2, UpgradePolicy.create().upgrade(zone3.toDeprecatedId()));
+                          .setOsUpgradePolicy(cloud1, UpgradePolicy.create().upgrade(zone1).upgrade(zone2))
+                          .setOsUpgradePolicy(cloud2, UpgradePolicy.create().upgrade(zone3));
         osUpgraders = List.of(
                 new OsUpgrader(tester.controller(), Duration.ofDays(1),
                                new JobControl(tester.controller().curator()),