Make RoutingGeneratorMock modifiable

3 files changed, 40 insertions, 20 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 6b5e5b0199c..85c7ea0af1e 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
@@ -12,11 +12,9 @@ import com.yahoo.config.provision.RegionName;
 import com.yahoo.config.provision.SystemName;
 import com.yahoo.config.provision.TenantName;
 import com.yahoo.vespa.athenz.api.NToken;
-import com.yahoo.vespa.config.SlimeUtils;
 import com.yahoo.vespa.hosted.controller.api.application.v4.model.DeployOptions;
 import com.yahoo.vespa.hosted.controller.api.application.v4.model.EndpointStatus;
 import com.yahoo.vespa.hosted.controller.api.identifiers.DeploymentId;
-import com.yahoo.vespa.hosted.controller.api.integration.deployment.JobType;
 import com.yahoo.vespa.hosted.controller.api.integration.dns.Record;
 import com.yahoo.vespa.hosted.controller.api.integration.dns.RecordName;
 import com.yahoo.vespa.hosted.controller.api.integration.zone.ZoneId;
@@ -34,13 +32,10 @@ import com.yahoo.vespa.hosted.controller.rotation.RotationLock;
 import org.junit.Test;
 
 import java.io.IOException;
-import java.nio.file.Files;
-import java.nio.file.Paths;
 import java.time.Duration;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.Supplier;
 
 import static com.yahoo.config.provision.SystemName.main;
 import static com.yahoo.vespa.hosted.controller.api.integration.deployment.JobType.component;
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 324ae4440cd..106cd280098 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
@@ -23,6 +23,7 @@ import com.yahoo.vespa.hosted.controller.api.integration.entity.MemoryEntityServ
 import com.yahoo.vespa.hosted.controller.api.integration.github.GitHubMock;
 import com.yahoo.vespa.hosted.controller.api.integration.organization.MockOrganization;
 import com.yahoo.vespa.hosted.controller.api.integration.routing.MemoryGlobalRoutingService;
+import com.yahoo.vespa.hosted.controller.api.integration.routing.RoutingGenerator;
 import com.yahoo.vespa.hosted.controller.api.integration.stubs.MockBuildService;
 import com.yahoo.vespa.hosted.controller.api.integration.stubs.MockLogStore;
 import com.yahoo.vespa.hosted.controller.api.integration.zone.ZoneId;
@@ -69,6 +70,7 @@ public final class ControllerTester {
     private final EntityService entityService;
     private final MockBuildService buildService;
     private final MetricsServiceMock metricsService;
+    private final RoutingGeneratorMock routingGenerator;
 
     private Controller controller;
 
@@ -77,7 +79,7 @@ public final class ControllerTester {
         this(new AthenzDbMock(), clock, new ConfigServerMock(new ZoneRegistryMock()),
              new ZoneRegistryMock(), new GitHubMock(), curatorDb, rotationsConfig,
              new MemoryNameService(), new ArtifactRepositoryMock(), new MemoryEntityService(), new MockBuildService(),
-             metricsService);
+             metricsService, new RoutingGeneratorMock());
     }
 
     public ControllerTester(ManualClock clock) {
@@ -101,7 +103,7 @@ public final class ControllerTester {
                              GitHubMock gitHub, CuratorDb curator, RotationsConfig rotationsConfig,
                              MemoryNameService nameService, ArtifactRepositoryMock artifactRepository,
                              EntityService entityService, MockBuildService buildService,
-                             MetricsServiceMock metricsService) {
+                             MetricsServiceMock metricsService, RoutingGeneratorMock routingGenerator) {
         this.athenzDb = athenzDb;
         this.clock = clock;
         this.configServer = configServer;
@@ -114,9 +116,10 @@ public final class ControllerTester {
         this.entityService = entityService;
         this.buildService = buildService;
         this.metricsService = metricsService;
+        this.routingGenerator = routingGenerator;
         this.controller = createController(curator, rotationsConfig, configServer, clock, gitHub, zoneRegistry,
                                            athenzDb, nameService, artifactRepository, entityService, buildService,
-                                           metricsService);
+                                           metricsService, routingGenerator);
 
         // Make root logger use time from manual clock
         Logger.getLogger("").getHandlers()[0].setFilter(
@@ -150,10 +153,13 @@ public final class ControllerTester {
 
     public MetricsServiceMock metricsService() { return metricsService; }
 
+    public RoutingGeneratorMock routingGenerator() { return routingGenerator; }
+
     /** Create a new controller instance. Useful to verify that controller state is rebuilt from persistence */
     public final void createNewController() {
         controller = createController(curator, rotationsConfig, configServer, clock, gitHub, zoneRegistry, athenzDb,
-                                      nameService, artifactRepository, entityService, buildService, metricsService);
+                                      nameService, artifactRepository, entityService, buildService, metricsService,
+                                      routingGenerator);
     }
 
     /** Creates the given tenant and application and deploys it */
@@ -267,7 +273,8 @@ public final class ControllerTester {
                                                GitHubMock gitHub, ZoneRegistryMock zoneRegistryMock,
                                                AthenzDbMock athensDb, MemoryNameService nameService,
                                                ArtifactRepository artifactRepository, EntityService entityService,
-                                               BuildService buildService, MetricsServiceMock metricsService) {
+                                               BuildService buildService, MetricsServiceMock metricsService,
+                                               RoutingGenerator routingGenerator) {
         Controller controller = new Controller(curator,
                                                rotationsConfig,
                                                gitHub,
@@ -278,7 +285,7 @@ public final class ControllerTester {
                                                configServer,
                                                metricsService,
                                                nameService,
-                                               new RoutingGeneratorMock(),
+                                               routingGenerator,
                                                new ChefMock(),
                                                clock,
                                                new AthenzClientFactoryMock(athensDb),
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/RoutingGeneratorMock.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/RoutingGeneratorMock.java
index bfc50a5a93b..46e41eea6e1 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/RoutingGeneratorMock.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/integration/RoutingGeneratorMock.java
@@ -1,28 +1,46 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.controller.integration;
 
+import com.yahoo.component.AbstractComponent;
 import com.yahoo.vespa.hosted.controller.api.integration.routing.RoutingEndpoint;
 import com.yahoo.vespa.hosted.controller.api.integration.routing.RoutingGenerator;
 import com.yahoo.vespa.hosted.controller.api.identifiers.DeploymentId;
 
-import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
+ * Returns a default set of endpoints on every query if it has no mappings, or those added by the user, otherwise.
+ *
  * @author bratseth
+ * @author jonmv
  */
 public class RoutingGeneratorMock implements RoutingGenerator {
 
+    private final Map<DeploymentId, List<RoutingEndpoint>> allEndpoints = new ConcurrentHashMap<>();
+    private static final List<RoutingEndpoint> defaultEndpoints =
+            Arrays.asList(new RoutingEndpoint("http://old-endpoint.vespa.yahooapis.com:4080", false),
+                          new RoutingEndpoint("http://qrs-endpoint.vespa.yahooapis.com:4080", "host1", false),
+                          new RoutingEndpoint("http://feeding-endpoint.vespa.yahooapis.com:4080", "host2", false),
+                          new RoutingEndpoint("http://global-endpoint.vespa.yahooapis.com:4080", "host1", true),
+                          new RoutingEndpoint("http://alias-endpoint.vespa.yahooapis.com:4080", "host1", true));
+
     @Override
     public List<RoutingEndpoint> endpoints(DeploymentId deployment) {
-        List<RoutingEndpoint> endpoints = new ArrayList<>();
-        // TODO: TLS: Update to HTTPS when ready.
-        endpoints.add(new RoutingEndpoint("http://old-endpoint.vespa.yahooapis.com:4080", false));
-        endpoints.add(new RoutingEndpoint("http://qrs-endpoint.vespa.yahooapis.com:4080", "host1", false));
-        endpoints.add(new RoutingEndpoint("http://feeding-endpoint.vespa.yahooapis.com:4080", "host2", false));
-        endpoints.add(new RoutingEndpoint("http://global-endpoint.vespa.yahooapis.com:4080", "host1", true));
-        endpoints.add(new RoutingEndpoint("http://alias-endpoint.vespa.yahooapis.com:4080", "host1", true));
-        return endpoints;
+        return allEndpoints.isEmpty()
+                ? defaultEndpoints
+                : allEndpoints.getOrDefault(deployment, Collections.emptyList());
+    }
+
+    public void putEndpoints(DeploymentId deployment, List<RoutingEndpoint> endpoints) {
+        allEndpoints.put(deployment, endpoints);
+    }
+
+    public void removeEndpoints(DeploymentId deployment) {
+        allEndpoints.remove(deployment);
     }
 
 }