refactored code and added tests

Author: seaunderwater

controllers/controlplane/controlplane.go

@@ -116,9 +116,7 @@ func (r *CritControlPlaneReconciler) reconcileControlPlane(
 		// created control plane will be spread out over failure domains
 		fds := controlPlane.NFailureDomainsWithFewestMachines(ccp.Spec.Replicas)
 		for i := 0; i < ccp.Spec.Replicas; i++ {
-			// this allows cycling through fds when replicas > fds
-			fd := fds[i%len(fds)]
-			if err := r.cloneConfigsAndGenerateMachine(ctx, cluster, ccp, bootstrapSpec, &fd); err != nil {
+			if err := r.cloneConfigsAndGenerateMachine(ctx, cluster, ccp, bootstrapSpec, fds[i]); err != nil {
 				logger.Error(err, "failed to create initial control plane Machine")
 				r.recorder.Eventf(ccp, corev1.EventTypeWarning, "FailedInitialization", "Failed to create initial control plane Machine for cluster %s/%s control plane: %v", cluster.Namespace, cluster.Name, err)
 				return ctrl.Result{}, err

internal/control_plane.go

@@ -164,11 +164,8 @@ func (c *ControlPlane) FailureDomainWithFewestMachines() *string {
 
 // NFailureDomainWithFewestMachines returns N failure domains with the fewest number of machines.
 // Used to distribute control plane replicas over several failure domains
-func (c *ControlPlane) NFailureDomainsWithFewestMachines(replicas int) []string {
-	if len(c.Cluster.Status.FailureDomains.FilterControlPlane()) == 0 {
-		return nil
-	}
-	return NPickFewest(c.FailureDomains().FilterControlPlane(), c.Machines, replicas)
+func (c *ControlPlane) NFailureDomainsWithFewestMachines(replicas int) []*string {
+	return PickNFewest(c.FailureDomains().FilterControlPlane(), c.Machines, replicas)
 }
 
 // GenerateCritConfig generates a new crit config for creating new control plane nodes.

internal/failure_domain.go

@@ -87,19 +87,17 @@ func PickFewest(failureDomains clusterv1.FailureDomains, machines FilterableMach
 	return pointer.StringPtr(aggregations[0].id)
 }
 
-// NPickFewest returns N failure domains with the fewest number of machines.
-func NPickFewest(failureDomains clusterv1.FailureDomains, machines FilterableMachineCollection, replicas int) []string {
+// PickNFewest returns N failure domains with the fewest number of machines.
+func PickNFewest(failureDomains clusterv1.FailureDomains, machines FilterableMachineCollection, n int) []*string {
+	fds := make([]*string, n)
 	aggregations := pick(failureDomains, machines)
 	if len(aggregations) == 0 {
-		return nil
+		return fds
 	}
 	sort.Sort(aggregations)
-	fds := make([]string, 0)
-	for i := 0; i < len(aggregations); i++ {
-		fds = append(fds, aggregations[i].id)
-	}
-	if len(aggregations) <= replicas {
-		return fds[:len(aggregations)]
+	for i := 0; i < n; i++ {
+		// allows cycling through fds when n > len(aggregations)
+		fds[i] = pointer.StringPtr(aggregations[i%len(aggregations)].id)
 	}
 	return fds
 }

internal/failure_domain_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package internal
 
 import (
+	"github.com/onsi/gomega/types"
 	"testing"
 
 	. "github.com/onsi/gomega"
@@ -97,6 +98,111 @@ func TestNewFailureDomainPicker(t *testing.T) {
 	}
 }
 
+func TestNewFailureDomainPickNFewest(t *testing.T) {
+	a := pointer.StringPtr("us-west-1a")
+	b := pointer.StringPtr("us-west-1b")
+	c := pointer.StringPtr("us-west-1c")
+	fds := clusterv1.FailureDomains{
+		*a: clusterv1.FailureDomainSpec{},
+		*b: clusterv1.FailureDomainSpec{},
+		*c: clusterv1.FailureDomainSpec{},
+	}
+	machinea := &clusterv1.Machine{Spec: clusterv1.MachineSpec{FailureDomain: a}}
+	machineb := &clusterv1.Machine{Spec: clusterv1.MachineSpec{FailureDomain: b}}
+	machinec := &clusterv1.Machine{Spec: clusterv1.MachineSpec{FailureDomain: c}}
+	machinenil := &clusterv1.Machine{Spec: clusterv1.MachineSpec{FailureDomain: nil}}
+
+	testcases := []struct {
+		name     string
+		fds      clusterv1.FailureDomains
+		machines FilterableMachineCollection
+		n        int
+		expected types.GomegaMatcher
+	}{
+		{
+			name:     "simple",
+			n:        1,
+			expected: Equal([]*string{nil}),
+		},
+		{
+			name:     "simple 3",
+			n:        3,
+			expected: Equal([]*string{nil, nil, nil}),
+		},
+		{
+			name: "machines and no failure domain",
+			machines: NewFilterableMachineCollection(
+				machinea.DeepCopy(),
+				machineb.DeepCopy(),
+				machinec.DeepCopy()),
+			n:        3,
+			expected: Equal([]*string{nil, nil, nil}),
+		},
+		{
+			name: "failure domains and no machines",
+			fds:  fds,
+			n:    2,
+			expected: Or(
+				ConsistOf([]*string{a, b}),
+				ConsistOf([]*string{b, c}),
+				ConsistOf([]*string{c, a}),
+			),
+		},
+		{
+			name: "machines in all but one failure domain",
+			fds:  fds,
+			machines: NewFilterableMachineCollection(
+				machinea.DeepCopy(),
+				machineb.DeepCopy()),
+			n: 2,
+			expected: Or(
+				ConsistOf([]*string{c, b}),
+				ConsistOf([]*string{c, a}),
+			),
+		},
+		{
+			name:     "n greater than number of failure domains",
+			fds:      fds,
+			n:        6,
+			expected: ConsistOf([]*string{a, b, c, a, b, c}),
+		},
+		{
+			name: "no failure domain specified on machine",
+			fds: clusterv1.FailureDomains{
+				*a: clusterv1.FailureDomainSpec{},
+			},
+			machines: NewFilterableMachineCollection(machinenil.DeepCopy()),
+			n:        3,
+			expected: Equal([]*string{a, a, a}),
+		},
+		{
+			name: "mismatched failure domain on machine",
+			fds: clusterv1.FailureDomains{
+				*a: clusterv1.FailureDomainSpec{},
+			},
+			machines: NewFilterableMachineCollection(
+				machineb.DeepCopy(),
+				machinec.DeepCopy(),
+			),
+			n:        3,
+			expected: Equal([]*string{a, a, a}),
+		},
+	}
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			g := NewWithT(t)
+
+			fds := PickNFewest(tc.fds, tc.machines, tc.n)
+			if tc.expected == nil {
+				g.Expect(fds).To(BeNil())
+			} else {
+				g.Expect(fds).To(tc.expected)
+				g.Expect(len(fds)).To(Equal(tc.n))
+			}
+		})
+	}
+}
+
 func TestNewFailureDomainPickMost(t *testing.T) {
 	a := pointer.StringPtr("us-west-1a")
 	b := pointer.StringPtr("us-west-1b")