Improve coverage of unit tests (#61)

Author: dgrove-oss

Makefile

@@ -88,7 +88,7 @@ vet: ## Run go vet against code.
 
 .PHONY: test
 test: manifests generate fmt vet envtest ## Run unit tests.
-	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) --bin-dir $(LOCALBIN) -p path)" go test $$(go list ./... | grep -v /e2e) -v -ginkgo.v -coverprofile cover.out
+	KUBEBUILDER_ASSETS="$(shell $(ENVTEST) use $(ENVTEST_K8S_VERSION) --bin-dir $(LOCALBIN) -p path)" go test $$(go list ./api/... ./internal/... ./pkg/...) -v -ginkgo.v -coverprofile cover.out
 
 # Utilize Kind or modify the e2e tests to load the image locally, enabling compatibility with other vendors.
 .PHONY: test-e2e  # Run the e2e tests against a Kind k8s instance that is spun up.

hack/e2e-util.sh

@@ -155,6 +155,9 @@ function configure_cluster {
   echo "Installing cert-manager"
   kubectl apply -f https://github.com/cert-manager/cert-manager/releases/download/$CERTMANAGER_VERSION/cert-manager.yaml
 
+  echo "Installing Kubeflow operator version $KUBEFLOW_VERSION"
+  kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=$KUBEFLOW_VERSION"
+
   # sleep to ensure cert-manager is fully functional
   echo "Waiting for pod in the cert-manager namespace to become ready"
   while [[ $(kubectl get pods -n cert-manager -o 'jsonpath={..status.conditions[?(@.type=="Ready")].status}' | tr ' ' '\n' | sort -u) != "True" ]]
@@ -163,9 +166,6 @@ function configure_cluster {
   done
   echo ""
 
-  echo "Installing Kubeflow operator version $KUBEFLOW_VERSION"
-  kubectl apply -k "github.com/kubeflow/training-operator/manifests/overlays/standalone?ref=$KUBEFLOW_VERSION"
-
   # Sleep until the kubeflow operator is running
   echo "Waiting for pods in the kueueflow namespace to become ready"
   while [[ $(kubectl get pods -n kubeflow -o 'jsonpath={..status.conditions[?(@.type=="Ready")].status}' | tr ' ' '\n' | sort -u) != "True" ]]

internal/controller/appwrapper/appwrapper_controller.go

@@ -160,13 +160,13 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 			Type:    string(workloadv1beta2.QuotaReserved),
 			Status:  metav1.ConditionTrue,
 			Reason:  string(workloadv1beta2.AppWrapperResuming),
-			Message: "AppWrapper was unsuspended by Kueue",
+			Message: "Suspend is false",
 		})
 		meta.SetStatusCondition(&aw.Status.Conditions, metav1.Condition{
 			Type:    string(workloadv1beta2.ResourcesDeployed),
 			Status:  metav1.ConditionTrue,
 			Reason:  string(workloadv1beta2.AppWrapperResuming),
-			Message: "AppWrapper was unsuspended by Kueue",
+			Message: "Suspend is false",
 		})
 		return r.updateStatus(ctx, aw, workloadv1beta2.AppWrapperResuming)
 
@@ -180,7 +180,7 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 				meta.SetStatusCondition(&aw.Status.Conditions, metav1.Condition{
 					Type:    string(workloadv1beta2.PodsReady),
 					Status:  metav1.ConditionFalse,
-					Reason:  "CreatedFailed",
+					Reason:  "CreateFailed",
 					Message: fmt.Sprintf("fatal error creating components: %v", err),
 				})
 				return r.updateStatus(ctx, aw, workloadv1beta2.AppWrapperFailed) // abort on fatal error
@@ -244,14 +244,14 @@ func (r *AppWrapperReconciler) Reconcile(ctx context.Context, req ctrl.Request)
 				Type:    string(workloadv1beta2.ResourcesDeployed),
 				Status:  metav1.ConditionFalse,
 				Reason:  string(workloadv1beta2.AppWrapperSuspended),
-				Message: "AppWrapper was suspended by Kueue",
+				Message: "Suspend is true",
 			})
 		}
 		meta.SetStatusCondition(&aw.Status.Conditions, metav1.Condition{
 			Type:    string(workloadv1beta2.QuotaReserved),
 			Status:  metav1.ConditionFalse,
 			Reason:  string(workloadv1beta2.AppWrapperSuspended),
-			Message: "AppWrapper was suspended by Kueue",
+			Message: "Suspend is true",
 		})
 		return r.updateStatus(ctx, aw, workloadv1beta2.AppWrapperSuspended)
 
@@ -299,6 +299,7 @@ func parseComponent(aw *workloadv1beta2.AppWrapper, raw []byte) (*unstructured.U
 	if namespace == "" {
 		obj.SetNamespace(aw.Namespace)
 	} else if namespace != aw.Namespace {
+		// Should not happen, namespace equality checked by validateAppWrapperInvariants
 		return nil, fmt.Errorf("component namespace \"%s\" is different from appwrapper namespace \"%s\"", namespace, aw.Namespace)
 	}
 	return obj, nil

internal/controller/appwrapper/appwrapper_controller_test.go

@@ -38,17 +38,34 @@ var _ = Describe("AppWrapper Controller", func() {
 	var awReconciler *AppWrapperReconciler
 	var awName types.NamespacedName
 	markerPodSet := podset.PodSetInfo{
-		Labels:      map[string]string{"testkey1": "value1"},
-		Annotations: map[string]string{"test2": "test2"},
+		Labels:       map[string]string{"testkey1": "value1"},
+		Annotations:  map[string]string{"test2": "test2"},
+		NodeSelector: map[string]string{"nodeName": "myNode"},
+		Tolerations:  []v1.Toleration{{Key: "aKey", Operator: "Exists", Effect: "NoSchedule"}},
 	}
 	var kueuePodSets []kueue.PodSet
 
-	advanceToRunning := func() {
+	advanceToResuming := func(components ...workloadv1beta2.AppWrapperComponent) {
+		By("Create an AppWrapper")
+		aw := toAppWrapper(components...)
+		aw.Spec.Suspend = true
+		Expect(k8sClient.Create(ctx, aw)).To(Succeed())
+		awName = types.NamespacedName{
+			Name:      aw.Name,
+			Namespace: aw.Namespace,
+		}
+		awReconciler = &AppWrapperReconciler{
+			Client: k8sClient,
+			Scheme: k8sClient.Scheme(),
+			Config: &config.AppWrapperConfig{ManageJobsWithoutQueueName: true, StandaloneMode: false},
+		}
+		kueuePodSets = (*workload.AppWrapper)(aw).PodSets()
+
 		By("Reconciling: Empty -> Suspended")
 		_, err := awReconciler.Reconcile(ctx, reconcile.Request{NamespacedName: awName})
 		Expect(err).NotTo(HaveOccurred())
 
-		aw := getAppWrapper(awName)
+		aw = getAppWrapper(awName)
 		Expect(aw.Status.Phase).Should(Equal(workloadv1beta2.AppWrapperSuspended))
 		Expect(controllerutil.ContainsFinalizer(aw, AppWrapperFinalizer)).Should(BeTrue())
 
@@ -67,12 +84,14 @@ var _ = Describe("AppWrapper Controller", func() {
 		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.QuotaReserved))).Should(BeTrue())
 		Expect((*workload.AppWrapper)(aw).IsActive()).Should(BeTrue())
 		Expect((*workload.AppWrapper)(aw).IsSuspended()).Should(BeFalse())
+	}
 
+	beginRunning := func() {
 		By("Reconciling: Resuming -> Running")
-		_, err = awReconciler.Reconcile(ctx, reconcile.Request{NamespacedName: awName})
+		_, err := awReconciler.Reconcile(ctx, reconcile.Request{NamespacedName: awName})
 		Expect(err).NotTo(HaveOccurred())
 
-		aw = getAppWrapper(awName)
+		aw := getAppWrapper(awName)
 		Expect(aw.Status.Phase).Should(Equal(workloadv1beta2.AppWrapperRunning))
 		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.ResourcesDeployed))).Should(BeTrue())
 		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.QuotaReserved))).Should(BeTrue())
@@ -83,10 +102,30 @@ var _ = Describe("AppWrapper Controller", func() {
 		Expect(err).NotTo(HaveOccurred())
 		Expect(podStatus.pending).Should(Equal(utils.ExpectedPodCount(aw)))
 
+		By("Simulating first Pod Running")
+		Expect(setPodStatus(aw, v1.PodRunning, 1)).To(Succeed())
+		By("Reconciling: Running -> Running")
+		_, err = awReconciler.Reconcile(ctx, reconcile.Request{NamespacedName: awName})
+		Expect(err).NotTo(HaveOccurred())
+
+		aw = getAppWrapper(awName)
+		Expect(aw.Status.Phase).Should(Equal(workloadv1beta2.AppWrapperRunning))
+		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.ResourcesDeployed))).Should(BeTrue())
+		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.QuotaReserved))).Should(BeTrue())
+		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.PodsReady))).Should(BeFalse())
+		Expect((*workload.AppWrapper)(aw).IsActive()).Should(BeTrue())
+		Expect((*workload.AppWrapper)(aw).IsSuspended()).Should(BeFalse())
+		podStatus, err = awReconciler.workloadStatus(ctx, aw)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(podStatus.pending).Should(Equal(utils.ExpectedPodCount(aw) - 1))
+	}
+
+	fullyRunning := func() {
+		aw := getAppWrapper(awName)
 		By("Simulating all Pods Running")
 		Expect(setPodStatus(aw, v1.PodRunning, utils.ExpectedPodCount(aw))).To(Succeed())
 		By("Reconciling: Running -> Running")
-		_, err = awReconciler.Reconcile(ctx, reconcile.Request{NamespacedName: awName})
+		_, err := awReconciler.Reconcile(ctx, reconcile.Request{NamespacedName: awName})
 		Expect(err).NotTo(HaveOccurred())
 
 		aw = getAppWrapper(awName)
@@ -97,30 +136,13 @@ var _ = Describe("AppWrapper Controller", func() {
 		Expect((*workload.AppWrapper)(aw).IsActive()).Should(BeTrue())
 		Expect((*workload.AppWrapper)(aw).IsSuspended()).Should(BeFalse())
 		Expect((*workload.AppWrapper)(aw).PodsReady()).Should(BeTrue())
-		podStatus, err = awReconciler.workloadStatus(ctx, aw)
+		podStatus, err := awReconciler.workloadStatus(ctx, aw)
 		Expect(err).NotTo(HaveOccurred())
 		Expect(podStatus.running).Should(Equal(utils.ExpectedPodCount(aw)))
 		_, finished := (*workload.AppWrapper)(aw).Finished()
 		Expect(finished).Should(BeFalse())
 	}
 
-	BeforeEach(func() {
-		By("Create an AppWrapper containing two Pods")
-		aw := toAppWrapper(pod(100), pod(100))
-		aw.Spec.Suspend = true
-		Expect(k8sClient.Create(ctx, aw)).To(Succeed())
-		awName = types.NamespacedName{
-			Name:      aw.Name,
-			Namespace: aw.Namespace,
-		}
-		awReconciler = &AppWrapperReconciler{
-			Client: k8sClient,
-			Scheme: k8sClient.Scheme(),
-			Config: &config.AppWrapperConfig{ManageJobsWithoutQueueName: true, StandaloneMode: false},
-		}
-		kueuePodSets = (*workload.AppWrapper)(aw).PodSets()
-	})
-
 	AfterEach(func() {
 		By("Cleanup the AppWrapper and ensure no Pods remain")
 		aw := &workloadv1beta2.AppWrapper{}
@@ -139,7 +161,9 @@ var _ = Describe("AppWrapper Controller", func() {
 	})
 
 	It("Happy Path Lifecycle", func() {
-		advanceToRunning()
+		advanceToResuming(pod(100), pod(100))
+		beginRunning()
+		fullyRunning()
 
 		By("Simulating one Pod Completing")
 		aw := getAppWrapper(awName)
@@ -176,9 +200,11 @@ var _ = Describe("AppWrapper Controller", func() {
 	})
 
 	It("Running Workloads can be Suspended", func() {
-		advanceToRunning()
+		advanceToResuming(pod(100), pod(100))
+		beginRunning()
+		fullyRunning()
 
-		By("Updating aw.Spec by invoking RunWithPodSetsInfo")
+		By("Inoking Suspend and RestorePodSetsInfo")
 		aw := getAppWrapper(awName)
 		(*workload.AppWrapper)(aw).Suspend()
 		Expect((*workload.AppWrapper)(aw).RestorePodSetsInfo(utilslices.Map(kueuePodSets, podset.FromPodSet))).To(BeTrue())
@@ -212,8 +238,10 @@ var _ = Describe("AppWrapper Controller", func() {
 		Expect(podStatus.failed + podStatus.succeeded + podStatus.running + podStatus.pending).Should(Equal(int32(0)))
 	})
 
-	It("A Pod Failure leads to a failed AppWrappers", func() {
-		advanceToRunning()
+	It("A Pod Failure leads to a failed AppWrapper", func() {
+		advanceToResuming(pod(100), pod(100))
+		beginRunning()
+		fullyRunning()
 
 		By("Simulating one Pod Failing")
 		aw := getAppWrapper(awName)
@@ -248,4 +276,23 @@ var _ = Describe("AppWrapper Controller", func() {
 		Expect(finished).Should(BeTrue())
 	})
 
+	It("Failure during resource creation leads to a failed AppWrapper", func() {
+		advanceToResuming(pod(100), malformedPod(100))
+
+		By("Reconciling: Resuming -> Failed")
+		_, err := awReconciler.Reconcile(ctx, reconcile.Request{NamespacedName: awName})
+		Expect(err).NotTo(HaveOccurred())
+
+		aw := getAppWrapper(awName)
+		Expect(aw.Status.Phase).Should(Equal(workloadv1beta2.AppWrapperFailed))
+		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.ResourcesDeployed))).Should(BeTrue())
+		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.QuotaReserved))).Should(BeTrue())
+		Expect(meta.IsStatusConditionTrue(aw.Status.Conditions, string(workloadv1beta2.PodsReady))).Should(BeFalse())
+		Expect((*workload.AppWrapper)(aw).IsActive()).Should(BeTrue())
+		Expect((*workload.AppWrapper)(aw).IsSuspended()).Should(BeFalse())
+		podStatus, err := awReconciler.workloadStatus(ctx, aw)
+		Expect(err).NotTo(HaveOccurred())
+		Expect(podStatus.pending).Should(Equal(int32(1)))
+	})
+
 })

internal/controller/appwrapper/fixtures_test.go

@@ -47,6 +47,7 @@ func randName(baseName string) string {
 
 func toAppWrapper(components ...workloadv1beta2.AppWrapperComponent) *workloadv1beta2.AppWrapper {
 	return &workloadv1beta2.AppWrapper{
+		TypeMeta:   metav1.TypeMeta{APIVersion: workloadv1beta2.GroupVersion.String(), Kind: "AppWrapper"},
 		ObjectMeta: metav1.ObjectMeta{Name: randName("aw"), Namespace: "default"},
 		Spec:       workloadv1beta2.AppWrapperSpec{Components: components},
 	}
@@ -111,26 +112,23 @@ func pod(milliCPU int64) workloadv1beta2.AppWrapperComponent {
 	}
 }
 
-const namespacedPodYAML = `
+const malformedPodYAML = `
 apiVersion: v1
 kind: Pod
 metadata:
   name: %v
-  namespace: %v
 spec:
   restartPolicy: Never
   containers:
   - name: busybox
-    image: quay.io/project-codeflare/busybox:1.36
     command: ["sh", "-c", "sleep 10"]
     resources:
       requests:
         cpu: %v`
 
-func namespacedPod(namespace string, milliCPU int64) workloadv1beta2.AppWrapperComponent {
-	yamlString := fmt.Sprintf(namespacedPodYAML,
+func malformedPod(milliCPU int64) workloadv1beta2.AppWrapperComponent {
+	yamlString := fmt.Sprintf(malformedPodYAML,
 		randName("pod"),
-		namespace,
 		resource.NewMilliQuantity(milliCPU, resource.DecimalSI))
 
 	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
@@ -141,67 +139,3 @@ func namespacedPod(namespace string, milliCPU int64) workloadv1beta2.AppWrapperC
 		Template: runtime.RawExtension{Raw: jsonBytes},
 	}
 }
-
-const serviceYAML = `
-apiVersion: v1
-kind: Service
-metadata:
-  name: %v
-spec:
-  selector:
-    app: test
-  ports:
-  - protocol: TCP
-    port: 80
-    targetPort: 8080`
-
-func service() workloadv1beta2.AppWrapperComponent {
-	yamlString := fmt.Sprintf(serviceYAML, randName("service"))
-	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
-	Expect(err).NotTo(HaveOccurred())
-	return workloadv1beta2.AppWrapperComponent{
-		PodSets:  []workloadv1beta2.AppWrapperPodSet{},
-		Template: runtime.RawExtension{Raw: jsonBytes},
-	}
-}
-
-const deploymentYAML = `
-apiVersion: apps/v1
-kind: Deployment
-metadata:
-  name: %v
-  labels:
-    app: test
-spec:
-  replicas: %v
-  selector:
-    matchLabels:
-      app: test
-  template:
-    metadata:
-      labels:
-        app: test
-    spec:
-      terminationGracePeriodSeconds: 0
-      containers:
-      - name: busybox
-        image: quay.io/project-codeflare/busybox:1.36
-        command: ["sh", "-c", "sleep 10000"]
-        resources:
-          requests:
-            cpu: %v`
-
-func deployment(replicaCount int, milliCPU int64) workloadv1beta2.AppWrapperComponent {
-	yamlString := fmt.Sprintf(deploymentYAML,
-		randName("deployment"),
-		replicaCount,
-		resource.NewMilliQuantity(milliCPU, resource.DecimalSI))
-
-	jsonBytes, err := yaml.YAMLToJSON([]byte(yamlString))
-	Expect(err).NotTo(HaveOccurred())
-	replicas := int32(replicaCount)
-	return workloadv1beta2.AppWrapperComponent{
-		PodSets:  []workloadv1beta2.AppWrapperPodSet{{Replicas: &replicas, Path: "template.spec.template"}},
-		Template: runtime.RawExtension{Raw: jsonBytes},
-	}
-}

internal/controller/workload/workload_controller.go

@@ -78,10 +78,7 @@ func (aw *AppWrapper) PodSets() []kueue.PodSet {
 				continue // Should be unreachable; Template.Raw validated by our AdmissionController
 			}
 			for psIdx, podSet := range component.PodSets {
-				replicas := int32(1)
-				if podSet.Replicas != nil {
-					replicas = *podSet.Replicas
-				}
+				replicas := utils.Replicas(podSet)
 				if template, err := utils.GetPodTemplateSpec(obj, podSet.Path); err == nil {
 					podSets = append(podSets, kueue.PodSet{
 						Name:     fmt.Sprintf("%s-%v-%v", aw.Name, componentIdx, psIdx),