Consolidate HostProvisioner mocks

1 files changed, 173 insertions, 0 deletions

diff --git a/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockHostProvisioner.java b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockHostProvisioner.java
new file mode 100644
index 00000000000..4c9323ca68e
--- /dev/null
+++ b/node-repository/src/main/java/com/yahoo/vespa/hosted/provision/testutils/MockHostProvisioner.java
@@ -0,0 +1,173 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.provision.testutils;
+
+import com.yahoo.component.Version;
+import com.yahoo.config.provision.ApplicationId;
+import com.yahoo.config.provision.Flavor;
+import com.yahoo.config.provision.NodeResources;
+import com.yahoo.config.provision.OutOfCapacityException;
+import com.yahoo.vespa.hosted.provision.Node;
+import com.yahoo.vespa.hosted.provision.node.Address;
+import com.yahoo.vespa.hosted.provision.node.IP;
+import com.yahoo.vespa.hosted.provision.provisioning.FatalProvisioningException;
+import com.yahoo.vespa.hosted.provision.provisioning.HostProvisioner;
+import com.yahoo.vespa.hosted.provision.provisioning.ProvisionedHost;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.EnumSet;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+/**
+ * @author mpolden
+ */
+public class MockHostProvisioner implements HostProvisioner {
+
+    private final List<ProvisionedHost> provisionedHosts = new ArrayList<>();
+    private final List<Flavor> flavors;
+    private final MockNameResolver nameResolver;
+    private final int memoryTaxGb;
+
+    private int deprovisionedHosts = 0;
+    private EnumSet<Behaviour> behaviours = EnumSet.noneOf(Behaviour.class);
+    private Optional<Flavor> hostFlavor = Optional.empty();
+
+    public MockHostProvisioner(List<Flavor> flavors, MockNameResolver nameResolver, int memoryTaxGb) {
+        this.flavors = List.copyOf(flavors);
+        this.nameResolver = nameResolver;
+        this.memoryTaxGb = memoryTaxGb;
+    }
+
+    public MockHostProvisioner(List<Flavor> flavors) {
+        this(flavors, 0);
+    }
+
+    public MockHostProvisioner(List<Flavor> flavors, int memoryTaxGb) {
+        this(flavors, new MockNameResolver().mockAnyLookup(), memoryTaxGb);
+    }
+
+    @Override
+    public List<ProvisionedHost> provisionHosts(List<Integer> provisionIndexes, NodeResources resources,
+                                                ApplicationId applicationId, Version osVersion, HostSharing sharing) {
+        Flavor hostFlavor = this.hostFlavor.orElseGet(() -> flavors.stream().filter(f -> compatible(f, resources))
+                                                                   .findFirst()
+                                                                   .orElseThrow(() -> new OutOfCapacityException("No host flavor matches " + resources)));
+        List<ProvisionedHost> hosts = new ArrayList<>();
+        for (int index : provisionIndexes) {
+            hosts.add(new ProvisionedHost("host" + index,
+                                          "hostname" + index,
+                                          hostFlavor,
+                                          Optional.empty(),
+                                          createAddressesForHost(hostFlavor, index),
+                                          resources,
+                                          osVersion));
+        }
+        provisionedHosts.addAll(hosts);
+        return hosts;
+    }
+
+    @Override
+    public List<Node> provision(Node host, Set<Node> children) throws FatalProvisioningException {
+        if (behaviours.contains(Behaviour.failProvisioning)) throw new FatalProvisioningException("Failed to provision node(s)");
+        if (host.state() != Node.State.provisioned) throw new IllegalStateException("Host to provision must be in " + Node.State.provisioned);
+        List<Node> result = new ArrayList<>();
+        result.add(withIpAssigned(host));
+        for (var child : children) {
+            if (child.state() != Node.State.reserved) throw new IllegalStateException("Child to provisioned must be in " + Node.State.reserved);
+            result.add(withIpAssigned(child));
+        }
+        return result;
+    }
+
+    @Override
+    public void deprovision(Node host) {
+        if (behaviours.contains(Behaviour.failDeprovisioning)) throw new FatalProvisioningException("Failed to deprovision node");
+        provisionedHosts.removeIf(provisionedHost -> provisionedHost.hostHostname().equals(host.hostname()));
+        deprovisionedHosts++;
+    }
+
+    /** Returns the hosts that have been provisioned by this  */
+    public List<ProvisionedHost> provisionedHosts() {
+        return Collections.unmodifiableList(provisionedHosts);
+    }
+
+    /** Returns the number of hosts deprovisioned by this */
+    public int deprovisionedHosts() {
+        return deprovisionedHosts;
+    }
+
+    public MockHostProvisioner with(Behaviour first, Behaviour... rest) {
+        this.behaviours = EnumSet.of(first, rest);
+        return this;
+    }
+
+    public MockHostProvisioner without(Behaviour first, Behaviour... rest) {
+        Set<Behaviour> behaviours = new HashSet<>(this.behaviours);
+        behaviours.removeAll(EnumSet.of(first, rest));
+        this.behaviours = behaviours.isEmpty() ? EnumSet.noneOf(Behaviour.class) : EnumSet.copyOf(behaviours);
+        return this;
+    }
+
+    public MockHostProvisioner overrideHostFlavor(String flavorName) {
+        Flavor flavor = flavors.stream().filter(f -> f.name().equals(flavorName))
+                               .findFirst()
+                               .orElseThrow(() -> new IllegalArgumentException("No such flavor '" + flavorName + "'"));
+        hostFlavor = Optional.of(flavor);
+        return this;
+    }
+
+    public boolean compatible(Flavor flavor, NodeResources resources) {
+        NodeResources resourcesToVerify = resources.withMemoryGb(resources.memoryGb() - memoryTaxGb);
+
+        if (flavor.resources().storageType() == NodeResources.StorageType.remote
+            && flavor.resources().diskGb() >= resources.diskGb())
+            resourcesToVerify = resourcesToVerify.withDiskGb(flavor.resources().diskGb());
+        if (flavor.resources().bandwidthGbps() >= resources.bandwidthGbps())
+            resourcesToVerify = resourcesToVerify.withBandwidthGbps(flavor.resources().bandwidthGbps());
+        return flavor.resources().compatibleWith(resourcesToVerify);
+    }
+
+    private List<Address> createAddressesForHost(Flavor flavor, int hostIndex) {
+        long numAddresses = Math.max(1, Math.round(flavor.resources().bandwidthGbps()));
+        return IntStream.range(0, (int) numAddresses)
+                        .mapToObj(i -> new Address("nodename" + hostIndex + "_" + i))
+                        .collect(Collectors.toList());
+    }
+
+    private Node withIpAssigned(Node node) {
+        if (node.parentHostname().isPresent()) return node;
+        int hostIndex = Integer.parseInt(node.hostname().replaceAll("^[a-z]+|-\\d+$", ""));
+        Set<String> addresses = Set.of("::" + hostIndex + ":0");
+        Set<String> ipAddressPool = new HashSet<>();
+        if (!behaviours.contains(Behaviour.failDnsUpdate)) {
+            nameResolver.addRecord(node.hostname(), addresses.iterator().next());
+            for (int i = 1; i <= 2; i++) {
+                String ip = "::" + hostIndex + ":" + i;
+                ipAddressPool.add(ip);
+                nameResolver.addRecord(node.hostname() + "-" + i, ip);
+            }
+        }
+
+        IP.Pool pool = node.ipConfig().pool().withIpAddresses(ipAddressPool);
+        return node.with(node.ipConfig().withPrimary(addresses).withPool(pool));
+    }
+
+    public enum Behaviour {
+
+        /** Fail all calls to {@link MockHostProvisioner#provision(com.yahoo.vespa.hosted.provision.Node, java.util.Set)} */
+        failProvisioning,
+
+        /** Fail all calls to {@link MockHostProvisioner#deprovision(com.yahoo.vespa.hosted.provision.Node)} */
+        failDeprovisioning,
+
+        /** Fail DNS updates of provisioned hosts */
+        failDnsUpdate,
+
+    }
+
+}