Simplify testing with custom zones and system

5 files changed, 34 insertions, 27 deletions

diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTest.java
index 422c856ca01..87888baa947 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTest.java
@@ -599,7 +599,7 @@ public class ControllerTest {
         var context = tester.newDeploymentContext();
         ZoneId zone = ZoneId.from("dev", "us-east-1");
         tester.controllerTester().zoneRegistry()
-                .setRoutingMethod(ZoneApiMock.from(zone), RoutingMethod.shared, RoutingMethod.sharedLayer4);
+              .setRoutingMethod(ZoneApiMock.from(zone), RoutingMethod.shared, RoutingMethod.sharedLayer4);
 
         // Deploy
         context.runJob(zone, applicationPackage);
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java
index 3e25a09b7d7..6c7abcc530e 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/ControllerTester.java
@@ -8,6 +8,7 @@ import com.yahoo.config.provision.HostName;
 import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.TenantName;
+import com.yahoo.config.provision.zone.RoutingMethod;
 import com.yahoo.config.provision.zone.ZoneApi;
 import com.yahoo.config.provision.zone.ZoneId;
 import com.yahoo.test.ManualClock;
@@ -38,6 +39,7 @@ import com.yahoo.vespa.hosted.controller.integration.ConfigServerMock;
 import com.yahoo.vespa.hosted.controller.integration.MetricsMock;
 import com.yahoo.vespa.hosted.controller.integration.SecretStoreMock;
 import com.yahoo.vespa.hosted.controller.integration.ServiceRegistryMock;
+import com.yahoo.vespa.hosted.controller.integration.ZoneApiMock;
 import com.yahoo.vespa.hosted.controller.integration.ZoneRegistryMock;
 import com.yahoo.vespa.hosted.controller.persistence.CuratorDb;
 import com.yahoo.vespa.hosted.controller.persistence.MockCuratorDb;
@@ -68,6 +70,7 @@ import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
 import java.util.logging.Handler;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
@@ -197,6 +200,21 @@ public final class ControllerTester {
         return new Version(current.getMajor(), nextMinorVersion.getAndIncrement(), current.getMicro());
     }
 
+    /** Set the zones and system for this and bootstrap infrastructure nodes */
+    public ControllerTester setZones(List<ZoneId> zones, SystemName system) {
+        zoneRegistry().setZones(zones.stream().map(ZoneApiMock::from).collect(Collectors.toList()))
+                      .setSystemName(system);
+        configServer().bootstrap(zones, SystemApplication.notController());
+        return this;
+    }
+
+    /** Set the routing method for given zones */
+    public ControllerTester setRoutingMethod(List<ZoneId> zones, RoutingMethod routingMethod) {
+        zoneRegistry().setRoutingMethod(zones.stream().map(ZoneApiMock::from).collect(Collectors.toList()),
+                                        routingMethod);
+        return this;
+    }
+
     /** Create a new controller instance. Useful to verify that controller state is rebuilt from persistence */
     public final void createNewController() {
         if (inContainer)
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/ApplicationPackageBuilder.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/ApplicationPackageBuilder.java
index 9eaa15cdbe3..fc7a99eb2f0 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/ApplicationPackageBuilder.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/ApplicationPackageBuilder.java
@@ -186,27 +186,28 @@ public class ApplicationPackageBuilder {
         return this;
     }
 
+    /** Add a trusted certificate to security/clients.pem */
     public ApplicationPackageBuilder trust(X509Certificate certificate) {
         this.trustedCertificates.add(certificate);
         return this;
     }
 
+    /** Add a default trusted certificate to security/clients.pem */
     public ApplicationPackageBuilder trustDefaultCertificate() {
         try {
             var generator = KeyPairGenerator.getInstance("RSA");
-            var builder = X509CertificateBuilder.fromKeypair(
+            var certificate = X509CertificateBuilder.fromKeypair(
                     generator.generateKeyPair(),
                     new X500Principal("CN=name"),
                     Instant.now(),
                     Instant.now().plusMillis(300_000),
                     SignatureAlgorithm.SHA256_WITH_RSA,
                     X509CertificateBuilder.generateRandomSerialNumber()
-            );
-            this.trustedCertificates.add(builder.build());
+            ).build();
+            return trust(certificate);
         } catch (NoSuchAlgorithmException e) {
             throw new RuntimeException(e);
         }
-        return this;
     }
 
     private byte[] deploymentSpec() {
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentTriggerTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentTriggerTest.java
index 81c9f51278e..c8b4eaa5236 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentTriggerTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/DeploymentTriggerTest.java
@@ -5,15 +5,11 @@ import com.yahoo.component.Version;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.zone.RoutingMethod;
 import com.yahoo.config.provision.zone.ZoneId;
-import com.yahoo.vespa.hosted.controller.ControllerTester;
 import com.yahoo.vespa.hosted.controller.api.integration.deployment.ApplicationVersion;
 import com.yahoo.vespa.hosted.controller.api.integration.deployment.JobType;
 import com.yahoo.vespa.hosted.controller.api.integration.deployment.RunId;
 import com.yahoo.vespa.hosted.controller.application.ApplicationPackage;
 import com.yahoo.vespa.hosted.controller.application.Change;
-import com.yahoo.vespa.hosted.controller.application.SystemApplication;
-import com.yahoo.vespa.hosted.controller.integration.ServiceRegistryMock;
-import com.yahoo.vespa.hosted.controller.integration.ZoneApiMock;
 import com.yahoo.vespa.hosted.controller.versions.VespaVersion;
 import org.junit.Test;
 
@@ -1061,17 +1057,13 @@ public class DeploymentTriggerTest {
         ApplicationPackage cdPackage = new ApplicationPackageBuilder().region("cd-us-central-1")
                                                                       .region("cd-aws-us-east-1a")
                                                                       .build();
-        ServiceRegistryMock services = new ServiceRegistryMock();
-        var zones = List.of(ZoneApiMock.fromId("test.cd-us-central-1"),
-                            ZoneApiMock.fromId("staging.cd-us-central-1"),
-                            ZoneApiMock.fromId("prod.cd-us-central-1"),
-                            ZoneApiMock.fromId("prod.cd-aws-us-east-1a"));
-        services.zoneRegistry()
-                .setSystemName(SystemName.cd)
-                .setZones(zones)
-                .setRoutingMethod(zones, RoutingMethod.shared);
-        tester = new DeploymentTester(new ControllerTester(services));
-        tester.configServer().bootstrap(services.zoneRegistry().zones().all().ids(), SystemApplication.values());
+        var zones = List.of(ZoneId.from("test.cd-us-central-1"),
+                            ZoneId.from("staging.cd-us-central-1"),
+                            ZoneId.from("prod.cd-us-central-1"),
+                            ZoneId.from("prod.cd-aws-us-east-1a"));
+        tester.controllerTester()
+              .setZones(zones, SystemName.cd)
+              .setRoutingMethod(zones, RoutingMethod.shared);
         tester.controllerTester().upgradeSystem(Version.fromString("6.1"));
         tester.controllerTester().computeVersionStatus();
         var app = tester.newDeploymentContext();
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java
index bfb749eb681..bf27fa005bc 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/routing/RoutingPoliciesTest.java
@@ -576,14 +576,10 @@ public class RoutingPoliciesTest {
         var prodZone = ZoneId.from(Environment.prod, sharedRegion);
         var stagingZone = ZoneId.from(Environment.staging, sharedRegion);
         var testZone = ZoneId.from(Environment.test, sharedRegion);
-        var zones = List.of(ZoneApiMock.from(prodZone),
-                            ZoneApiMock.from(stagingZone),
-                            ZoneApiMock.from(testZone));
-        tester.controllerTester().zoneRegistry()
-              .setZones(zones)
+        var zones = List.of(prodZone, stagingZone, testZone);
+        tester.controllerTester()
+              .setZones(zones, SystemName.Public)
               .setRoutingMethod(zones, RoutingMethod.exclusive);
-        tester.controllerTester().configServer().bootstrap(List.of(prodZone, stagingZone, testZone),
-                                                           SystemApplication.notController());
 
         var context = tester.tester.newDeploymentContext();
         var endpointId = EndpointId.of("r0");