From b53124ea3e77faef1593829648715abfef9a51a1 Mon Sep 17 00:00:00 2001
From: Dejan Pejchev <pejcev.dejan@gmail.com>
Date: Wed, 17 Apr 2024 01:18:38 +0200
Subject: [PATCH] add unit tests for createHeadlessSvcIfNecessary

---
 pkg/controllers/jobset_controller_test.go | 123 ++++++++++++++++++++++
 1 file changed, 123 insertions(+)

diff --git a/pkg/controllers/jobset_controller_test.go b/pkg/controllers/jobset_controller_test.go
index ecd42a2e..6497d958 100644
--- a/pkg/controllers/jobset_controller_test.go
+++ b/pkg/controllers/jobset_controller_test.go
@@ -19,6 +19,13 @@ import (
 	"testing"
 	"time"
 
+	"k8s.io/utils/ptr"
+
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
+	"k8s.io/klog/v2/ktesting"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/stretchr/testify/assert"
@@ -1205,3 +1212,119 @@ func makeJob(args *makeJobArgs) *testutils.JobWrapper {
 		PodAnnotations(annotations)
 	return jobWrapper
 }
+
+func TestCreateHeadlessSvcIfNecessary(t *testing.T) {
+	var (
+		jobSetName = "test-jobset"
+		ns         = "default"
+	)
+
+	tests := []struct {
+		name                           string
+		jobSet                         *jobset.JobSet
+		existingService                bool
+		expectServiceCreate            bool
+		expectServiceName              string
+		expectService                  *corev1.Service
+		expectPublishNotReadyAddresses bool
+	}{
+		{
+			name:            "headless service exists and should not be created",
+			jobSet:          testutils.MakeJobSet(jobSetName, ns).EnableDNSHostnames(true).Obj(),
+			existingService: true,
+		},
+		{
+			name:                           "service does not exist and should be created, subdomain not set",
+			jobSet:                         testutils.MakeJobSet(jobSetName, ns).EnableDNSHostnames(true).Obj(),
+			expectServiceCreate:            true,
+			expectServiceName:              "test-jobset",
+			expectPublishNotReadyAddresses: true,
+		},
+		{
+			name:                           "service does not exist and should be created, subdomain set",
+			jobSet:                         testutils.MakeJobSet(jobSetName, ns).EnableDNSHostnames(true).NetworkSubdomain("test-subdomain").Obj(),
+			expectServiceCreate:            true,
+			expectServiceName:              "test-subdomain",
+			expectPublishNotReadyAddresses: true,
+		},
+		{
+			name:                           "service does exist and should be created, publishNotReadyAddresses is false",
+			jobSet:                         testutils.MakeJobSet(jobSetName, ns).EnableDNSHostnames(true).PublishNotReadyAddresses(false).Obj(),
+			expectServiceCreate:            true,
+			expectServiceName:              "test-jobset",
+			expectPublishNotReadyAddresses: false,
+		},
+		{
+			name:                           "service does not exist and should be created, publishNotReadyAddresses is true",
+			jobSet:                         testutils.MakeJobSet(jobSetName, ns).EnableDNSHostnames(true).PublishNotReadyAddresses(true).Obj(),
+			expectServiceCreate:            true,
+			expectServiceName:              "test-jobset",
+			expectPublishNotReadyAddresses: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			_, ctx := ktesting.NewTestContext(t)
+			scheme := runtime.NewScheme()
+			utilruntime.Must(jobset.AddToScheme(scheme))
+			utilruntime.Must(corev1.AddToScheme(scheme))
+			fakeClientBuilder := fake.NewClientBuilder().WithScheme(scheme)
+			if tc.existingService {
+				svc := &corev1.Service{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "test-jobset",
+						Namespace: ns,
+					},
+				}
+				fakeClientBuilder.WithObjects(svc)
+			}
+			fakeClient := fakeClientBuilder.Build()
+
+			// Create a JobSetReconciler instance with the fake client
+			r := &JobSetReconciler{
+				Client: fakeClient,
+				Scheme: scheme,
+			}
+
+			// Execute the function under test
+			gotErr := r.createHeadlessSvcIfNecessary(ctx, tc.jobSet)
+			if gotErr != nil {
+				t.Errorf("createHeadlessSvcIfNecessary() error = %v", gotErr)
+			}
+			if tc.expectServiceCreate {
+				svc := &corev1.Service{}
+				gotErr = fakeClient.Get(ctx, types.NamespacedName{Name: tc.expectServiceName, Namespace: ns}, svc)
+				if gotErr != nil {
+					t.Errorf("expected service to be created but got an error: %v", gotErr)
+				}
+				if len(svc.OwnerReferences) != 1 {
+					t.Error("expected service to have owner reference set")
+				}
+				expectedOwnerRef := metav1.OwnerReference{
+					APIVersion:         "jobset.x-k8s.io/v1alpha2",
+					Kind:               "JobSet",
+					Name:               "test-jobset",
+					Controller:         ptr.To(true),
+					BlockOwnerDeletion: ptr.To(true),
+				}
+				if diff := cmp.Diff(expectedOwnerRef, svc.OwnerReferences[0]); diff != "" {
+					t.Errorf("unexpected service owner reference value (+got/-want): %s", diff)
+				}
+				if svc.Spec.ClusterIP != corev1.ClusterIPNone {
+					t.Errorf("expected service to have ClusterIP None, got %s", svc.Spec.ClusterIP)
+				}
+				selector, ok := svc.Spec.Selector[jobset.JobSetNameKey]
+				if !ok {
+					t.Errorf("expected service selector to contain %q key", jobset.JobSetNameKey)
+				}
+				if selector != tc.jobSet.Name {
+					t.Errorf("expected service selector to be %q, got %q", tc.jobSet.Name, selector)
+				}
+				if svc.Spec.PublishNotReadyAddresses != tc.expectPublishNotReadyAddresses {
+					t.Errorf("expected PublishNotReadyAddresses to be %t, got %t", tc.expectPublishNotReadyAddresses, svc.Spec.PublishNotReadyAddresses)
+				}
+			}
+		})
+	}
+}