fix: add custom template (#168)

Add custom pod template to e2e tests
Azure · Nov 15, 2023 · b97a288 · b97a288
1 parent fa60f47
commit b97a288
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 0 deletions.
diff --git a/pkg/controllers/workspace_controller.go b/pkg/controllers/workspace_controller.go
@@ -96,6 +96,7 @@ func (c *WorkspaceReconciler) addOrUpdateWorkspace(ctx context.Context, wObj *ka
 	if wObj.Inference.Preset != nil && strings.Contains(string(wObj.Inference.Preset.Name), "llama") {
 		useHeadlessService = true
 	}
+
 	if err := c.ensureService(ctx, wObj, useHeadlessService); err != nil {
 		if updateErr := c.updateStatusConditionIfNotMatch(ctx, wObj, kaitov1alpha1.WorkspaceConditionTypeReady, metav1.ConditionFalse,
 			"workspaceFailed", err.Error()); updateErr != nil {

diff --git a/pkg/resources/manifests.go b/pkg/resources/manifests.go
@@ -270,6 +270,7 @@ func GenerateDeploymentManifestWithPodTemplate(ctx context.Context, workspaceObj
 	}
 
 	templateCopy := workspaceObj.Inference.Template.DeepCopy()
+
 	if templateCopy.ObjectMeta.Labels == nil {
 		templateCopy.ObjectMeta.Labels = make(map[string]string)
 	}

diff --git a/test/e2e/preset_test.go b/test/e2e/preset_test.go
@@ -67,6 +67,20 @@ func createLlama13BWorkspaceWithPresetPrivateMode(registry, registrySecret, imag
 	return workspaceObj
 }
 
+func createCustomWorkspaceWithPresetCustomMode(imageName string, numOfNode int) *kaitov1alpha1.Workspace {
+	workspaceObj := &kaitov1alpha1.Workspace{}
+	By("Creating a workspace CR with custom workspace mode", func() {
+		uniqueID := fmt.Sprint("preset-", rand.Intn(1000))
+		workspaceObj = utils.GenerateWorkspaceManifest(uniqueID, namespaceName, "",
+			numOfNode, "Standard_D4s_v3", &metav1.LabelSelector{
+				MatchLabels: map[string]string{"kaito-workspace": "private-preset-e2e-test"},
+			}, nil, "", utils.InferenceModeCustomTemplate, nil, utils.GeneratePodTemplate(uniqueID, namespaceName, imageName, nil))
+
+		createAndValidateWorkspace(workspaceObj)
+	})
+	return workspaceObj
+}
+
 func createAndValidateWorkspace(workspaceObj *kaitov1alpha1.Workspace) {
 	By("Creating workspace", func() {
 		Eventually(func() error {
@@ -363,4 +377,22 @@ var _ = Describe("Workspace Preset", func() {
 		validateWorkspaceReadiness(workspaceObj)
 	})
 
+	It("should create a custom template workspace successfully", func() {
+		numOfNode := 1
+		imageName := "nginx:latest"
+		workspaceObj := createCustomWorkspaceWithPresetCustomMode(imageName, numOfNode)
+
+		defer cleanupResources(workspaceObj)
+
+		time.Sleep(30 * time.Second)
+		validateMachineCreation(workspaceObj, numOfNode)
+		validateResourceStatus(workspaceObj)
+
+		time.Sleep(30 * time.Second)
+
+		validateInferenceResource(workspaceObj, int32(numOfNode), false)
+
+		validateWorkspaceReadiness(workspaceObj)
+	})
+
 })