feat(ws): set status.pauseTime on workspace (#26)

* feat(ws): populate the status.pauseTime on the workspace Signed-off-by: Adem Baccara <[email protected]> * test(ws): add unit test Signed-off-by: Adem Baccara <[email protected]> * mathew commit 1 Signed-off-by: Mathew Wicks <[email protected]> --------- Signed-off-by: Adem Baccara <[email protected]> Signed-off-by: Mathew Wicks <[email protected]> Co-authored-by: Mathew Wicks <[email protected]>
kubeflow · Aug 6, 2024 · 38b8955 · 38b8955
1 parent e4e0aff
commit 38b8955
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 6 deletions.
diff --git a/workspaces/controller/internal/controller/workspace_controller.go b/workspaces/controller/internal/controller/workspace_controller.go
@@ -371,10 +371,6 @@ func (r *WorkspaceReconciler) Reconcile(ctx context.Context, req ctrl.Request) (
 		}
 	}
 
-	//
-	// TODO: figure out how to set `status.pauseTime`, it will probably have to be done in a webhook
-	//
-
 	// update Workspace status
 	workspaceStatus, err := r.generateWorkspaceStatus(ctx, log, workspace, pod, statefulSet)
 	if err != nil {
@@ -927,6 +923,18 @@ func generateService(workspace *kubefloworgv1beta1.Workspace, imageConfigSpec ku
 func (r *WorkspaceReconciler) generateWorkspaceStatus(ctx context.Context, log logr.Logger, workspace *kubefloworgv1beta1.Workspace, pod *corev1.Pod, statefulSet *appsv1.StatefulSet) (kubefloworgv1beta1.WorkspaceStatus, error) {
 	status := workspace.Status
 
+	// if workspace is paused, update the `status.pauseTime`
+	// NOTE: when the workspace is not paused, the pauseTime should be 0
+	if *workspace.Spec.Paused {
+		if status.PauseTime == 0 {
+			status.PauseTime = metav1.Now().Unix()
+		}
+	} else {
+		if status.PauseTime != 0 {
+			status.PauseTime = 0
+		}
+	}
+
 	// cases where the Pod does not exist
 	if pod == nil {
 		// STATUS: Paused

diff --git a/workspaces/controller/internal/controller/workspace_controller_test.go b/workspaces/controller/internal/controller/workspace_controller_test.go
@@ -20,6 +20,8 @@ import (
 	"fmt"
 	"time"
 
+	"k8s.io/utils/ptr"
+
 	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
 
@@ -130,6 +132,38 @@ var _ = Describe("Workspace Controller", func() {
 			workspace := NewExampleWorkspace1(workspaceName, namespaceName, workspaceKindName)
 			Expect(k8sClient.Create(ctx, workspace)).To(Succeed())
 
+			By("pausing the Workspace")
+			patch := client.MergeFrom(workspace.DeepCopy())
+			newWorkspace := workspace.DeepCopy()
+			newWorkspace.Spec.Paused = ptr.To(true)
+			Expect(k8sClient.Patch(ctx, newWorkspace, patch)).To(Succeed())
+
+			By("setting the Workspace `status.pauseTime` to the current time")
+			tolerance := int64(5)
+			currentTime := time.Now().Unix()
+			Eventually(func() (int64, error) {
+				err := k8sClient.Get(ctx, types.NamespacedName{Name: workspaceName, Namespace: namespaceName}, workspace)
+				if err != nil {
+					return 0, err
+				}
+				return workspace.Status.PauseTime, nil
+			}, timeout, interval).Should(BeNumerically("~", currentTime, tolerance))
+
+			By("un-pausing the Workspace")
+			patch = client.MergeFrom(workspace.DeepCopy())
+			newWorkspace = workspace.DeepCopy()
+			newWorkspace.Spec.Paused = ptr.To(false)
+			Expect(k8sClient.Patch(ctx, newWorkspace, patch)).To(Succeed())
+
+			By("setting the Workspace `status.pauseTime` to 0")
+			Eventually(func() (int64, error) {
+				err := k8sClient.Get(ctx, types.NamespacedName{Name: workspaceName, Namespace: namespaceName}, workspace)
+				if err != nil {
+					return 0, err
+				}
+				return workspace.Status.PauseTime, nil
+			}, timeout, interval).Should(BeZero())
+
 			By("creating a StatefulSet")
 			statefulSetList := &appsv1.StatefulSetList{}
 			Eventually(func() ([]appsv1.StatefulSet, error) {
@@ -138,7 +172,7 @@ var _ = Describe("Workspace Controller", func() {
 					return nil, err
 				}
 				return statefulSetList.Items, nil
-			}).Should(HaveLen(1))
+			}, timeout, interval).Should(HaveLen(1))
 
 			// TODO: use this to get the StatefulSet
 			//statefulSet := statefulSetList.Items[0]
@@ -151,7 +185,7 @@ var _ = Describe("Workspace Controller", func() {
 					return nil, err
 				}
 				return serviceList.Items, nil
-			}).Should(HaveLen(1))
+			}, timeout, interval).Should(HaveLen(1))
 
 			// TODO: use this to get the Service
 			//service := serviceList.Items[0]