Require distribution to be non-null and fix tests

3 files changed, 5 insertions, 7 deletions

diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java
index 2535589395d..2e8e2707166 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/ContentCluster.java
@@ -25,12 +25,11 @@ public class ContentCluster {
     private final String clusterName;
     private final ClusterInfo clusterInfo = new ClusterInfo();
     private final Map<Node, Long> nodeStartTimestamps = new TreeMap<>();
+    private final int maxNumberOfGroupsAllowedToBeDown;
 
     private int slobrokGenerationCount = 0;
     private Distribution distribution;
 
-    private final int maxNumberOfGroupsAllowedToBeDown;
-
     public ContentCluster(String clusterName, Collection<ConfiguredNode> configuredNodes, Distribution distribution) {
         this(clusterName, configuredNodes, distribution, -1);
     }
@@ -45,7 +44,7 @@ public class ContentCluster {
                    int maxNumberOfGroupsAllowedToBeDown) {
         if (configuredNodes == null) throw new IllegalArgumentException("Nodes must be set");
         this.clusterName = clusterName;
-        this.distribution = distribution;
+        this.distribution = Objects.requireNonNull(distribution, "distribution must be non-null");
         setNodes(configuredNodes, new NodeListener() {});
         this.maxNumberOfGroupsAllowedToBeDown = maxNumberOfGroupsAllowedToBeDown;
     }
diff --git a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/GroupAvailabilityCalculator.java b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/GroupAvailabilityCalculator.java
index a697d8f9868..35e3ae8d063 100644
--- a/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/GroupAvailabilityCalculator.java
+++ b/clustercontroller-core/src/main/java/com/yahoo/vespa/clustercontroller/core/GroupAvailabilityCalculator.java
@@ -12,6 +12,7 @@ import com.yahoo.vdslib.state.NodeType;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -27,7 +28,7 @@ class GroupAvailabilityCalculator {
                                         double minNodeRatioPerGroup,
                                         int safeMaintenanceGroupThreshold,
                                         List<Integer> nodesSafelySetToMaintenance) {
-        this.distribution = distribution;
+        this.distribution = Objects.requireNonNull(distribution, "distribution must be non-null");
         this.minNodeRatioPerGroup = minNodeRatioPerGroup;
         this.safeMaintenanceGroupThreshold = safeMaintenanceGroupThreshold;
         this.nodesSafelySetToMaintenance = nodesSafelySetToMaintenance;
@@ -181,9 +182,6 @@ class GroupAvailabilityCalculator {
     }
 
     public Result calculate(ClusterState state) {
-        if (distribution == null) { // FIXME: for tests that don't set distribution properly!
-            return new Result();
-        }
         if (isFlatCluster(distribution.getRootGroup())) {
             // Implicit group takedown only applies to hierarchic cluster setups.
             return new Result();
diff --git a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java
index 238dfd42da5..238a64459ca 100644
--- a/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java
+++ b/clustercontroller-core/src/test/java/com/yahoo/vespa/clustercontroller/core/FleetControllerTest.java
@@ -84,6 +84,7 @@ public abstract class FleetControllerTest implements Waiter {
     protected static FleetControllerOptions.Builder defaultOptions(String clusterName, Collection<ConfiguredNode> nodes) {
         var builder = new FleetControllerOptions.Builder(clusterName, nodes);
         builder.enableTwoPhaseClusterStateActivation(true); // Enable by default, tests can explicitly disable.
+        builder.setStorageDistribution(DistributionBuilder.forFlatCluster(builder.nodes().size()));
         return builder;
     }