Add tests

Author: LimKianAn

pkg/lbmanager/lbmanager.go

@@ -39,23 +39,12 @@ func (m *LBManager) CreateSvcLBIfNone(ctx context.Context, in *api.Postgres) err
 			return fmt.Errorf("failed to fetch Service of type LoadBalancer: %w", err)
 		}
 
-		existingLBIP, nextFreePort, err := m.nextFreeSocket(ctx)
+		nextFreePort, err := m.nextFreePort(ctx)
 		if err != nil {
-			return fmt.Errorf("failed to get a free port for creating Service of type LoadBalancer: %w", err)
-		}
-		var lbIPToUse string
-		if m.LBIP != "" {
-			// a specific IP was configured in the config, so use that one
-			lbIPToUse = m.LBIP
-		} else if existingLBIP != "" {
-			// no ip was configured, but one is already in use, so use the existing one
-			lbIPToUse = existingLBIP
-		} else {
-			// nothing was configured, nothing exists yet, so use an empty address so a new loadbalancer will be created and assigned
-			lbIPToUse = ""
+			return fmt.Errorf("failed to get the next free port: %w", err)
 		}
 
-		if err := m.Create(ctx, in.ToSvcLB(lbIPToUse, nextFreePort)); err != nil {
+		if err := m.Create(ctx, in.ToSvcLB(m.LBIP, nextFreePort)); err != nil {
 			return fmt.Errorf("failed to create Service of type LoadBalancer: %w", err)
 		}
 		return nil
@@ -75,35 +64,27 @@ func (m *LBManager) DeleteSvcLB(ctx context.Context, in *api.Postgres) error {
 }
 
 // nextFreeSocket finds any existing LoadBalancerIP and the next free port out of the configure port range.
-func (m *LBManager) nextFreeSocket(ctx context.Context) (string, int32, error) {
+func (m *LBManager) nextFreePort(ctx context.Context) (int32, error) {
 	// TODO prevent concurrency issues when calculating port / ip.
 
-	anyExistingLBIP := ""
-
 	// Fetch all services managed by this postgreslet
 	lbs := &corev1.ServiceList{}
 	if err := m.List(ctx, lbs, client.MatchingLabels(api.SvcLoadBalancerLabel)); err != nil {
-		return anyExistingLBIP, 0, fmt.Errorf("failed to fetch the list of services of type LoadBalancer: %w", err)
+		return 0, fmt.Errorf("failed to fetch the list of services of type LoadBalancer: %w", err)
 	}
 
 	// If there are none, this will be the first (managed) service we create, so start with PortRangeStart and return
 	if len(lbs.Items) == 0 {
-		return anyExistingLBIP, m.PortRangeStart, nil
+		return m.PortRangeStart, nil
 	}
 
 	// If there are already any managed services, store all the used ports in a slice.
-	// Also store the LoadBalancerIP.
-	portsInUse := make([]int32, len(lbs.Items))
+	portsInUse := []int32{}
 	for i := range lbs.Items {
 		svc := lbs.Items[i]
 		if len(svc.Spec.Ports) > 0 {
 			portsInUse = append(portsInUse, svc.Spec.Ports[0].Port)
 		}
-		if svc.Spec.LoadBalancerIP != "" {
-			// Technically, we only store the IP of the last Service in this list.
-			// As there should only be one IP per postgreslet and one postgreslet per cluster, this is good enough.
-			anyExistingLBIP = svc.Spec.LoadBalancerIP
-		}
 	}
 
 	// Now try all ports in the configured port range to find a free one.
@@ -115,11 +96,11 @@ func (m *LBManager) nextFreeSocket(ctx context.Context) (string, int32, error) {
 			continue
 		}
 		// The postgreslet hasn't assigned this port yet, so use it.
-		return anyExistingLBIP, port, nil
+		return port, nil
 	}
 
 	// If we made it this far, no free port could be found.
-	return anyExistingLBIP, 0, errors.New("no free port in the configured port range found")
+	return 0, errors.New("no free port in the configured port range found")
 }
 
 func containsElem(s []int32, v int32) bool {

pkg/lbmanager/lbmanager_test.go

@@ -0,0 +1,110 @@
+package lbmanager
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	api "github.com/fi-ts/postgreslet/api/v1"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+)
+
+func TestLBManager_nextFreePort(t *testing.T) {
+	portRangeStart := int32(0)
+	portRangeSize := int32(5)
+
+	tests := []struct {
+		name    string
+		lbMgr   *LBManager
+		want    int32
+		wantErr bool
+	}{
+		{
+			name: "no svc in the cluster",
+			lbMgr: &LBManager{
+				Client:         fake.NewClientBuilder().WithScheme(scheme()).WithLists(svcListWithPorts()).Build(),
+				LBIP:           "0.0.0.0",
+				PortRangeStart: portRangeStart,
+				PortRangeSize:  portRangeSize,
+			},
+			want:    0,
+			wantErr: false,
+		},
+		{
+			name: "one svc already in the cluster",
+			lbMgr: &LBManager{
+				Client:         fake.NewClientBuilder().WithScheme(scheme()).WithLists(svcListWithPorts(0)).Build(),
+				LBIP:           "0.0.0.0",
+				PortRangeStart: portRangeStart,
+				PortRangeSize:  portRangeSize,
+			},
+			want:    1,
+			wantErr: false,
+		},
+		{
+			name: "last free port left",
+			lbMgr: &LBManager{
+				Client:         fake.NewClientBuilder().WithScheme(scheme()).WithLists(svcListWithPorts(0, 1, 2, 3)).Build(),
+				LBIP:           "0.0.0.0",
+				PortRangeStart: portRangeStart,
+				PortRangeSize:  portRangeSize,
+			},
+			want:    4,
+			wantErr: false,
+		},
+		{
+			name: "no free port",
+			lbMgr: &LBManager{
+				Client:         fake.NewClientBuilder().WithScheme(scheme()).WithLists(svcListWithPorts(0, 1, 2, 3, 4)).Build(),
+				LBIP:           "0.0.0.0",
+				PortRangeStart: portRangeStart,
+				PortRangeSize:  portRangeSize,
+			},
+			want:    0,
+			wantErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := tt.lbMgr.nextFreePort(context.Background())
+			if (err != nil) != tt.wantErr {
+				t.Errorf("LBManager.nextFreePort() error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("LBManager.nextFreePort() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func scheme() *runtime.Scheme {
+	scheme := runtime.NewScheme()
+	_ = corev1.AddToScheme(scheme)
+
+	return scheme
+}
+
+// svcListWithPorts generates a `ServiceList` containing `Service`s with ports respectively
+func svcListWithPorts(ports ...int32) *corev1.ServiceList {
+	svcList := &corev1.ServiceList{}
+	for _, port := range ports {
+		svcList.Items = append(svcList.Items, *svcWithPort(port))
+	}
+	return svcList
+}
+
+func svcWithPort(port int32) *corev1.Service {
+	svc := corev1.Service{}
+	svc.Name = fmt.Sprintf("svc-with-port-%d", port)
+	svc.Labels = api.SvcLoadBalancerLabel
+	svc.Spec.Ports = []corev1.ServicePort{
+		{
+			Port: port,
+		},
+	}
+	return &svc
+}