Allocate nodes of type controller to controller app

2 files changed, 42 insertions, 17 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisioner.java
index b6955195dcf..f61d4158253 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisioner.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisioner.java
@@ -12,6 +12,7 @@ import com.yahoo.vespa.hosted.provision.Node;
 import com.yahoo.vespa.hosted.provision.NodeRepository;
 import com.yahoo.vespa.service.monitor.application.ConfigServerApplication;
 import com.yahoo.vespa.service.monitor.application.ConfigServerHostApplication;
+import com.yahoo.vespa.service.monitor.application.ControllerApplication;
 import com.yahoo.vespa.service.monitor.application.HostedVespaApplication;
 import com.yahoo.vespa.service.monitor.application.ProxyHostApplication;
 
@@ -35,7 +36,8 @@ public class InfrastructureProvisioner extends Maintainer {
     private static final List<HostedVespaApplication> HOSTED_VESPA_APPLICATIONS = Arrays.asList(
             ConfigServerApplication.CONFIG_SERVER_APPLICATION,
             ConfigServerHostApplication.CONFIG_SERVER_HOST_APPLICATION,
-            ProxyHostApplication.PROXY_HOST_APPLICATION);
+            ProxyHostApplication.PROXY_HOST_APPLICATION,
+            ControllerApplication.CONTROLLER_APPLICATION);
 
     private final Provisioner provisioner;
     private final InfrastructureVersions infrastructureVersions;
@@ -101,4 +103,5 @@ public class InfrastructureProvisioner extends Maintainer {
         }
         return targetVersion;
     }
+
 }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisionerTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisionerTest.java
index f0c4ad2ef2d..642e6adfc75 100644
--- a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisionerTest.java
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/maintenance/InfrastructureProvisionerTest.java
@@ -2,8 +2,8 @@
 package com.yahoo.vespa.hosted.provision.maintenance;
 
 import com.yahoo.component.Version;
-import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.ClusterMembership;
+import com.yahoo.config.provision.ClusterSpec;
 import com.yahoo.config.provision.NodeType;
 import com.yahoo.config.provision.Provisioner;
 import com.yahoo.vespa.hosted.provision.Node;
@@ -13,11 +13,17 @@ import com.yahoo.vespa.hosted.provision.node.Agent;
 import com.yahoo.vespa.hosted.provision.node.Allocation;
 import com.yahoo.vespa.hosted.provision.node.Generation;
 import com.yahoo.vespa.service.monitor.application.ConfigServerApplication;
+import com.yahoo.vespa.service.monitor.application.ControllerApplication;
+import com.yahoo.vespa.service.monitor.application.HostedVespaApplication;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
 
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.Optional;
 
-import org.junit.Test;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
@@ -26,47 +32,63 @@ import static org.mockito.Mockito.when;
 /**
  * @author freva
  */
+@RunWith(Parameterized.class)
 public class InfrastructureProvisionerTest {
 
-    private final NodeRepositoryTester tester = new NodeRepositoryTester();
+    @Parameters(name = "application={0}")
+    public static Iterable<Object[]> parameters() {
+        return Arrays.asList(
+                new HostedVespaApplication[]{ConfigServerApplication.CONFIG_SERVER_APPLICATION},
+                new HostedVespaApplication[]{ControllerApplication.CONTROLLER_APPLICATION}
+        );
+    }
 
+    private final NodeRepositoryTester tester = new NodeRepositoryTester();
     private final Provisioner provisioner = mock(Provisioner.class);
     private final NodeRepository nodeRepository = tester.nodeRepository();
     private final InfrastructureVersions infrastructureVersions = mock(InfrastructureVersions.class);
     private final InfrastructureProvisioner infrastructureProvisioner = new InfrastructureProvisioner(
             provisioner, nodeRepository, infrastructureVersions, Duration.ofDays(99), new JobControl(nodeRepository.database()));
 
+    private final HostedVespaApplication application;
+    private final NodeType nodeType;
+
+    public InfrastructureProvisionerTest(HostedVespaApplication application) {
+        this.application = application;
+        this.nodeType = application.getCapacity().type();
+    }
+
     @Test
     public void returns_version_if_usable_nodes_on_old_version() {
         Version target = Version.fromString("6.123.456");
         Version oldVersion = Version.fromString("6.122.333");
-        when(infrastructureVersions.getTargetVersionFor(eq(NodeType.config))).thenReturn(Optional.of(target));
+        when(infrastructureVersions.getTargetVersionFor(eq(nodeType))).thenReturn(Optional.of(target));
 
         addNode(1, Node.State.failed, Optional.of(oldVersion));
         addNode(2, Node.State.dirty, Optional.empty());
         addNode(3, Node.State.active, Optional.of(oldVersion));
 
-        assertEquals(Optional.of(target), infrastructureProvisioner.getTargetVersion(NodeType.config));
+        assertEquals(Optional.of(target), infrastructureProvisioner.getTargetVersion(nodeType));
     }
 
     @Test
     public void returns_version_if_has_usable_nodes_without_version() {
         Version target = Version.fromString("6.123.456");
         Version oldVersion = Version.fromString("6.122.333");
-        when(infrastructureVersions.getTargetVersionFor(eq(NodeType.config))).thenReturn(Optional.of(target));
+        when(infrastructureVersions.getTargetVersionFor(eq(nodeType))).thenReturn(Optional.of(target));
 
         addNode(1, Node.State.failed, Optional.of(oldVersion));
         addNode(2, Node.State.ready, Optional.empty());
         addNode(3, Node.State.active, Optional.of(target));
 
-        assertEquals(Optional.of(target), infrastructureProvisioner.getTargetVersion(NodeType.config));
+        assertEquals(Optional.of(target), infrastructureProvisioner.getTargetVersion(nodeType));
     }
 
     @Test
     public void returns_empty_if_usable_nodes_on_target_version() {
         Version target = Version.fromString("6.123.456");
         Version oldVersion = Version.fromString("6.122.333");
-        when(infrastructureVersions.getTargetVersionFor(eq(NodeType.config))).thenReturn(Optional.of(target));
+        when(infrastructureVersions.getTargetVersionFor(eq(nodeType))).thenReturn(Optional.of(target));
 
         addNode(1, Node.State.failed, Optional.of(oldVersion));
         addNode(2, Node.State.parked, Optional.of(target));
@@ -74,32 +96,31 @@ public class InfrastructureProvisionerTest {
         addNode(4, Node.State.inactive, Optional.of(target));
         addNode(5, Node.State.dirty, Optional.empty());
 
-        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(NodeType.config));
+        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(nodeType));
     }
 
     @Test
     public void returns_empty_if_no_usable_nodes() {
-        when(infrastructureVersions.getTargetVersionFor(eq(NodeType.config))).thenReturn(Optional.of(Version.fromString("6.123.456")));
+        when(infrastructureVersions.getTargetVersionFor(eq(nodeType))).thenReturn(Optional.of(Version.fromString("6.123.456")));
 
         // No nodes in node repo
-        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(NodeType.config));
+        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(nodeType));
 
         // Add nodes in non-provisionable states
         addNode(1, Node.State.dirty, Optional.empty());
         addNode(2, Node.State.failed, Optional.empty());
-        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(NodeType.config));
+        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(nodeType));
     }
 
     @Test
     public void returns_empty_if_target_version_not_set() {
-        when(infrastructureVersions.getTargetVersionFor(eq(NodeType.config))).thenReturn(Optional.empty());
-        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(NodeType.config));
+        when(infrastructureVersions.getTargetVersionFor(eq(nodeType))).thenReturn(Optional.empty());
+        assertEquals(Optional.empty(), infrastructureProvisioner.getTargetVersion(nodeType));
     }
 
     private Node addNode(int id, Node.State state, Optional<Version> wantedVespaVersion) {
-        Node node = tester.addNode("id-" + id, "node-" + id, "default", NodeType.config);
+        Node node = tester.addNode("id-" + id, "node-" + id, "default", nodeType);
         Optional<Node> nodeWithAllocation = wantedVespaVersion.map(version -> {
-            ConfigServerApplication application = ConfigServerApplication.CONFIG_SERVER_APPLICATION;
             ClusterSpec clusterSpec = ClusterSpec.from(application.getClusterType(), application.getClusterId(), ClusterSpec.Group.from(0), version, false);
             ClusterMembership membership = ClusterMembership.from(clusterSpec, 1);
             Allocation allocation = new Allocation(application.getApplicationId(), membership, new Generation(0, 0), false);
@@ -107,4 +128,5 @@ public class InfrastructureProvisionerTest {
         });
         return nodeRepository.database().writeTo(state, nodeWithAllocation.orElse(node), Agent.system, Optional.empty());
     }
+
 }