Add unit test to DockerHostCapacity

2 files changed, 170 insertions, 1 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerHostCapacity.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerHostCapacity.java
index 980948ae48e..9ab5d54bb54 100644
--- a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerHostCapacity.java
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/provisioning/DockerHostCapacity.java
@@ -38,6 +38,7 @@ public class DockerHostCapacity {
         if (comp == 0) {
             comp = freeCapacityOf(hostB, false).compare(freeCapacityOf(hostA, false));
             if (comp == 0) {
+                // If resources are equal - we want to assign to the one with the most IPaddresses free
                 comp = freeIPs(hostB) - freeIPs(hostA);
             }
         }
@@ -98,9 +99,12 @@ public class DockerHostCapacity {
      * Calculate the remaining capacity for the dockerHost.
      */
     private ResourceCapacity freeCapacityOf(Node dockerHost, boolean includeHeadroom) {
+        // Only hosts have free capacity
+        if (!dockerHost.type().equals(NodeType.host)) return new ResourceCapacity();
+
         ResourceCapacity hostCapacity = new ResourceCapacity(dockerHost);
         for (Node container : allNodes.childNodes(dockerHost).asList()) {
-            // Until we have migrated we might have docker containers unallocated
+            // Until we have migrated we might have docker containers unallocated - TODO check off if headroom tenant is safe
             if (includeHeadroom || !(container.allocation().isPresent() && container.allocation().get().owner().tenant().value().equals("headroom"))) {
                 hostCapacity.subtract(container);
             }
diff --git a/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DockerHostCapacityTest.java b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DockerHostCapacityTest.java
new file mode 100644
index 00000000000..2b843929b6e
--- /dev/null
+++ b/node-repository/src/test/java/com/yahoo/vespa/hosted/provision/provisioning/DockerHostCapacityTest.java
@@ -0,0 +1,165 @@
+package com.yahoo.vespa.hosted.provision.provisioning;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.ClusterMembership;
+import com.yahoo.config.provision.Flavor;
+import com.yahoo.config.provision.NodeFlavors;
+import com.yahoo.config.provision.NodeType;
+import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.node.Allocation;
+import com.yahoo.vespa.hosted.provision.node.Generation;
+import com.yahoo.vespa.hosted.provision.testutils.FlavorConfigBuilder;
+import org.junit.Before;
+import org.junit.Test;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+
+import static junit.framework.TestCase.assertTrue;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+
+/**
+ * @author smorgrav
+ */
+public class DockerHostCapacityTest {
+
+    private static final String HEADROOM_TENANT = "-__!@#$$%THISisHEADroom";
+
+    private DockerHostCapacity capacity;
+    private List<Node> nodes;
+    private Node host1, host2, host3;
+    Node nodeA, nodeB, nodeC, nodeD, nodeE;
+    Flavor flavorDocker, flavorDocker2;
+
+    @Before
+    public void setup() {
+        // Create flavors
+        NodeFlavors nodeFlavors = FlavorConfigBuilder.createDummies("host", "docker", "docker2");
+        flavorDocker = nodeFlavors.getFlavorOrThrow("docker");
+        flavorDocker2 = nodeFlavors.getFlavorOrThrow("docker2");
+
+        // Create three docker hosts
+        host1 = Node.create("host1", Collections.singleton("::1"), generateIPs(2, 4), "host1", Optional.empty(), nodeFlavors.getFlavorOrThrow("host"), NodeType.host);
+        host2 = Node.create("host2", Collections.singleton("::11"), generateIPs(12, 3), "host2", Optional.empty(), nodeFlavors.getFlavorOrThrow("host"), NodeType.host);
+        host3 = Node.create("host3", Collections.singleton("::21"), generateIPs(22, 1), "host3", Optional.empty(), nodeFlavors.getFlavorOrThrow("host"), NodeType.host);
+
+        // Add two containers to host1
+        nodeA = Node.create("nodeA", Collections.singleton("::2"), Collections.emptySet(), "nodeA", Optional.of("host1"), flavorDocker, NodeType.tenant);
+        nodeB = Node.create("nodeB", Collections.singleton("::3"), Collections.emptySet(), "nodeB", Optional.of("host1"), flavorDocker, NodeType.tenant);
+
+        // Add two containers to host 2 (same as host 1)
+        nodeC = Node.create("nodeC", Collections.singleton("::12"), Collections.emptySet(), "nodeC", Optional.of("host2"), flavorDocker, NodeType.tenant);
+        nodeD = Node.create("nodeD", Collections.singleton("::13"), Collections.emptySet(), "nodeD", Optional.of("host2"), flavorDocker, NodeType.tenant);
+
+        // Add a larger container to host3
+        nodeE = Node.create("nodeE", Collections.singleton("::22"), Collections.emptySet(), "nodeE", Optional.of("host3"), flavorDocker2, NodeType.tenant);
+
+        // init docker host capacity
+        nodes = new ArrayList<>();
+        Collections.addAll(nodes, host1, host2, host3, nodeA, nodeB, nodeC, nodeD, nodeE);
+        capacity = new DockerHostCapacity(nodes);
+    }
+
+    @Test
+    public void compare_used_to_sort_in_decending_order() {
+        assertEquals(host1, nodes.get(0)); //Make sure it is unsorted here
+        Collections.sort(nodes, capacity::compare);
+        assertEquals(host3, nodes.get(0));
+        assertEquals(host1, nodes.get(1));
+        assertEquals(host2, nodes.get(2));
+
+        // Add a headroom node to host1 and the host2 should be prioritized
+        Allocation allocation = new Allocation(app(HEADROOM_TENANT), ClusterMembership.from("container/id1/3", new Version()), Generation.inital(), false);
+        Node nodeF = Node.create("nodeF", Collections.singleton("::6"), Collections.emptySet(), "nodeF", Optional.of("host1"), flavorDocker, NodeType.tenant);
+        nodeF.with(allocation);
+        nodes.add(nodeF);
+        capacity = new DockerHostCapacity(nodes);
+        Collections.sort(nodes, capacity::compare);
+        assertEquals(host3, nodes.get(0));
+        assertEquals(host2, nodes.get(1));
+        assertEquals(host1, nodes.get(2));
+
+        // Remove a node from host1 and it should be prioritized again
+        nodes.remove(nodeA);
+        capacity = new DockerHostCapacity(nodes);
+        Collections.sort(nodes, capacity::compare);
+        assertEquals(host3, nodes.get(0));
+        assertEquals(host1, nodes.get(1));
+        assertEquals(host2, nodes.get(2));
+    }
+
+    @Test
+    public void hasCapacity() {
+        assertTrue(capacity.hasCapacity(host1, flavorDocker));
+        assertTrue(capacity.hasCapacity(host1, flavorDocker2));
+        assertTrue(capacity.hasCapacity(host2, flavorDocker));
+        assertTrue(capacity.hasCapacity(host2, flavorDocker2));
+        assertFalse(capacity.hasCapacity(host3, flavorDocker));  // No ip available
+        assertFalse(capacity.hasCapacity(host3, flavorDocker2)); // No ip available
+
+        // Add a new node to host1 to deplete the memory resource
+        Node nodeF = Node.create("nodeF", Collections.singleton("::6"), Collections.emptySet(),
+                "nodeF", Optional.of("host1"), flavorDocker, NodeType.tenant);
+        nodes.add(nodeF);
+        capacity = new DockerHostCapacity(nodes);
+        assertFalse(capacity.hasCapacity(host1, flavorDocker));
+        assertFalse(capacity.hasCapacity(host1, flavorDocker2));
+    }
+
+    @Test
+    public void freeIPs() {
+        assertEquals(2, capacity.freeIPs(host1));
+        assertEquals(1, capacity.freeIPs(host2));
+        assertEquals(0, capacity.freeIPs(host3));
+    }
+
+    @Test
+    public void getCapacityTotal() {
+        ResourceCapacity total = capacity.getCapacityTotal();
+        assertEquals(21.0, total.getCpu(), 0.1);
+        assertEquals(30.0, total.getMemory(), 0.1);
+        assertEquals(36.0, total.getDisk(), 0.1);
+    }
+
+    @Test
+    public void getFreeCapacityTotal() {
+        ResourceCapacity totalFree = capacity.getFreeCapacityTotal();
+        assertEquals(15.0, totalFree.getCpu(), 0.1);
+        assertEquals(14.0, totalFree.getMemory(), 0.1);
+        assertEquals(24.0, totalFree.getDisk(), 0.1);
+    }
+
+    @Test
+    public void freeCapacityInFlavorEquivalence() {
+        assertEquals(2, capacity.freeCapacityInFlavorEquivalence(flavorDocker));
+        assertEquals(2, capacity.freeCapacityInFlavorEquivalence(flavorDocker2));
+    }
+
+    @Test
+    public void getNofHostsAvailableFor() {
+        assertEquals(2, capacity.getNofHostsAvailableFor(flavorDocker));
+        assertEquals(2, capacity.getNofHostsAvailableFor(flavorDocker2));
+    }
+
+    private Set<String> generateIPs(int start, int count) {
+        // Allow 4 containers
+        Set<String> additionalIps = new HashSet<>();
+        for (int i = start; i < (start + count); i++) {
+            additionalIps.add("::" + i);
+        }
+        return additionalIps;
+    }
+
+    private ApplicationId app(String tenant) {
+        return new ApplicationId.Builder()
+                .tenant(tenant)
+                .applicationName("test")
+                .instanceName("default").build();
+    }
+}