Update the Handler when watch on SetContainerImageFromRegistry functi…

components/odh-notebook-controller/controllers/notebook_controller_test.go

@@ -242,6 +242,39 @@ var _ = Describe("The Openshift Notebook controller", func() {
 
 	})
 
+	// New test case for notebook update
+	When("Updating a Notebook", func() {
+		const (
+			Name      = "test-notebook-update"
+			Namespace = "default"
+		)
+
+		notebook := createNotebook(Name, Namespace)
+
+		It("Should update the Notebook specification", func() {
+			ctx := context.Background()
+
+			By("By creating a new Notebook")
+			Expect(cli.Create(ctx, notebook)).Should(Succeed())
+			time.Sleep(interval)
+
+			By("By updating the Notebook's image")
+			key := types.NamespacedName{Name: Name, Namespace: Namespace}
+			Expect(cli.Get(ctx, key, notebook)).Should(Succeed())
+
+			updatedImage := "registry.redhat.io/ubi8/ubi:updated"
+			notebook.Spec.Template.Spec.Containers[0].Image = updatedImage
+			Expect(cli.Update(ctx, notebook)).Should(Succeed())
+			time.Sleep(interval)
+
+			By("By checking that the Notebook's image is updated")
+			Eventually(func() string {
+				Expect(cli.Get(ctx, key, notebook)).Should(Succeed())
+				return notebook.Spec.Template.Spec.Containers[0].Image
+			}, duration, interval).Should(Equal(updatedImage))
+		})
+	})
+
 	When("Creating a Notebook, test Networkpolicies", func() {
 		const (
 			Name      = "test-notebook-np"

components/odh-notebook-controller/controllers/notebook_webhook.go

@@ -252,7 +252,10 @@ func (w *NotebookWebhook) Handle(ctx context.Context, req admission.Request) adm
 		if err != nil {
 			return admission.Errored(http.StatusInternalServerError, err)
 		}
+	}
 
+	// Check Imagestream Info both on create and update operations
+	if req.Operation == admissionv1.Create || req.Operation == admissionv1.Update {
 		// Check Imagestream Info
 		err = SetContainerImageFromRegistry(ctx, w.Config, notebook, log)
 		if err != nil {
@@ -456,100 +459,109 @@ func InjectCertConfig(notebook *nbv1.Notebook, configMapName string) error {
 // Otherwise, it checks the last-image-selection annotation to find the image stream and fetches the image from status.dockerImageReference,
 // assigning it to the container.image value.
 func SetContainerImageFromRegistry(ctx context.Context, config *rest.Config, notebook *nbv1.Notebook, log logr.Logger) error {
-    // Create a dynamic client
-    dynamicClient, err := dynamic.NewForConfig(config)
-    if err != nil {
-        log.Error(err, "Error creating dynamic client")
-        return err
-    }
-    // Specify the GroupVersionResource for imagestreams
-    ims := schema.GroupVersionResource{
-        Group:    "image.openshift.io",
-        Version:  "v1",
-        Resource: "imagestreams",
-    }
-
-    annotations := notebook.GetAnnotations()
-    if annotations != nil {
-        if imageSelection, exists := annotations["notebooks.opendatahub.io/last-image-selection"]; exists {
-            // Check if the image selection has an internal registry, if so  will pickup this. This value constructed on the initialization of the Notebook CR.
-            if strings.Contains(notebook.Spec.Template.Spec.Containers[0].Image, "image-registry.openshift-image-registry.svc:5000") {
-                log.Info("Internal registry found. Will pickup the default value from image field.")
-                return nil
-            } else {
-                // Split the imageSelection to imagestream and tag
-                parts := strings.Split(imageSelection, ":")
-                if len(parts) != 2 {
-                    log.Error(nil, "Invalid image selection format")
-                    return fmt.Errorf("invalid image selection format")
-                }
-
-                imagestreamName := parts[0]
-                tag := parts[1]
-
-                // Specify the namespaces to search in
-                namespaces := []string{"opendatahub", "redhat-ods-applications"}
-
-                imagestreamFound := false
-
-                for _, namespace := range namespaces {
-                    // List imagestreams in the specified namespace
-                    imagestreams, err := dynamicClient.Resource(ims).Namespace(namespace).List(ctx, metav1.ListOptions{})
-                    if err != nil {
-                        log.Error(err, "Cannot list imagestreams", "namespace", namespace)
-                        continue
-                    }
-
-                    // Iterate through the imagestreams to find matches
-                    for _, item := range imagestreams.Items {
-                        metadata := item.Object["metadata"].(map[string]interface{})
-                        name := metadata["name"].(string)
-
-                        if name == imagestreamName {
-                            status := item.Object["status"].(map[string]interface{})
-
-                            log.Info("No Internal registry found, pick up imageHash from status.tag.dockerImageReference")
-
-                            tags := status["tags"].([]interface{})
-                            for _, t := range tags {
-                                tagMap := t.(map[string]interface{})
-                                tagName := tagMap["tag"].(string)
-                                if tagName == tag {
-                                    items := tagMap["items"].([]interface{})
-                                    if len(items) > 0 {
-                                        // Sort items by creationTimestamp to get the most recent one
-                                        sort.Slice(items, func(i, j int) bool {
-                                            iTime := items[i].(map[string]interface{})["created"].(string)
-                                            jTime := items[j].(map[string]interface{})["created"].(string)
-                                            return iTime > jTime // Lexicographical comparison of RFC3339 timestamps
-                                        })
-                                        imageHash := items[0].(map[string]interface{})["dockerImageReference"].(string)
-                                        notebook.Spec.Template.Spec.Containers[0].Image = imageHash
-                                        // Update the JUPYTER_IMAGE environment variable
-                                        for i, envVar := range notebook.Spec.Template.Spec.Containers[0].Env {
-                                            if envVar.Name == "JUPYTER_IMAGE" {
-                                                notebook.Spec.Template.Spec.Containers[0].Env[i].Value = imageHash
-                                                break
-                                            }
-                                        }
-                                        imagestreamFound = true
-                                        break
-                                    }
-                                }
-                            }
-                        }
-                    }
-                    if imagestreamFound {
-                        break
-                    }
-                }
-
-                if !imagestreamFound {
-                    log.Info("Imagestream not found in any of the specified namespaces", "imagestreamName", imagestreamName, "tag", tag)
-                }
-            }
-        }
-    }
-
-    return nil
+	// Create a dynamic client
+	dynamicClient, err := dynamic.NewForConfig(config)
+	if err != nil {
+		log.Error(err, "Error creating dynamic client")
+		return err
+	}
+	// Specify the GroupVersionResource for imagestreams
+	ims := schema.GroupVersionResource{
+		Group:    "image.openshift.io",
+		Version:  "v1",
+		Resource: "imagestreams",
+	}
+
+	annotations := notebook.GetAnnotations()
+	if annotations != nil {
+		if imageSelection, exists := annotations["notebooks.opendatahub.io/last-image-selection"]; exists {
+
+			containerFound := false
+			// Iterate over containers to find the one matching the notebook name
+			for i, container := range notebook.Spec.Template.Spec.Containers {
+				if container.Name == notebook.Name {
+					containerFound = true
+
+					// Check if the container.Image value has an internal registry, if so  will pickup this without extra checks.
+					// This value constructed on the initialization of the Notebook CR.
+					if strings.Contains(container.Image, "image-registry.openshift-image-registry.svc:5000") {
+						log.Info("Internal registry found. Will pick up the default value from image field.")
+						return nil
+					} else {
+						log.Info("No internal registry found, let's pick up image reference from relevant ImageStream 'status.tags[].tag.dockerImageReference'")
+
+						// Split the imageSelection to imagestream and tag
+						imageSelected := strings.Split(imageSelection, ":")
+						if len(imageSelected) != 2 {
+							log.Error(nil, "Invalid image selection format")
+							return fmt.Errorf("invalid image selection format")
+						}
+
+						// Specify the namespaces to search in
+						namespaces := []string{"opendatahub", "redhat-ods-applications"}
+						imagestreamFound := false
+						for _, namespace := range namespaces {
+							// List imagestreams in the specified namespace
+							imagestreams, err := dynamicClient.Resource(ims).Namespace(namespace).List(ctx, metav1.ListOptions{})
+							if err != nil {
+								log.Error(err, "Cannot list imagestreams", "namespace", namespace)
+								continue
+							}
+
+							// Iterate through the imagestreams to find matches
+							for _, item := range imagestreams.Items {
+								metadata := item.Object["metadata"].(map[string]interface{})
+								name := metadata["name"].(string)
+
+								// Match with the ImageStream name
+								if name == imageSelected[0] {
+									status := item.Object["status"].(map[string]interface{})
+
+									// Match to the corresponding tag of the image
+									tags := status["tags"].([]interface{})
+									for _, t := range tags {
+										tagMap := t.(map[string]interface{})
+										tagName := tagMap["tag"].(string)
+										if tagName == imageSelected[1] {
+											items := tagMap["items"].([]interface{})
+											if len(items) > 0 {
+												// Sort items by creationTimestamp to get the most recent one
+												sort.Slice(items, func(i, j int) bool {
+													iTime := items[i].(map[string]interface{})["created"].(string)
+													jTime := items[j].(map[string]interface{})["created"].(string)
+													return iTime > jTime // Lexicographical comparison of RFC3339 timestamps
+												})
+												imageHash := items[0].(map[string]interface{})["dockerImageReference"].(string)
+												// Update the Containers[i].Image value
+												notebook.Spec.Template.Spec.Containers[i].Image = imageHash
+												// Update the JUPYTER_IMAGE environment variable with the image selection for example "jupyter-datascience-notebook:2023.2"
+												for i, envVar := range container.Env {
+													if envVar.Name == "JUPYTER_IMAGE" {
+														container.Env[i].Value = imageSelection
+														break
+													}
+												}
+												imagestreamFound = true
+												break
+											}
+										}
+									}
+								}
+							}
+							if imagestreamFound {
+								break
+							}
+						}
+						if !imagestreamFound {
+							log.Error(nil, "Imagestream not found in any of the specified namespaces", "imageSelected", imageSelected[0], "tag", imageSelected[1])
+						}
+					}
+				}
+			}
+			if !containerFound {
+				log.Error(nil, "No container found matching the notebook name", "notebookName", notebook.Name)
+			}
+		}
+	}
+	return nil
 }

components/odh-notebook-controller/e2e/notebook_controller_setup_test.go

@@ -4,11 +4,12 @@ import (
 	"context"
 	"flag"
 	"fmt"
-	netv1 "k8s.io/api/networking/v1"
 	"os"
 	"testing"
 	"time"
 
+	netv1 "k8s.io/api/networking/v1"
+
 	nbv1 "github.com/kubeflow/kubeflow/components/notebook-controller/api/v1"
 	routev1 "github.com/openshift/api/route/v1"
 	"github.com/pkg/errors"
@@ -135,8 +136,9 @@ func TestE2ENotebookController(t *testing.T) {
 	if !t.Run("validate controllers", testNotebookControllerValidation) {
 		return
 	}
-	// Run create and delete tests for all the test notebooks
+	// Run create, update and delete tests for all the test notebooks
 	t.Run("create", creationTestSuite)
+	t.Run("update", updateTestSuite)
 	if !skipDeletion {
 		t.Run("delete", deletionTestSuite)
 	}

components/odh-notebook-controller/e2e/notebook_update_test.go

@@ -0,0 +1,92 @@
+package e2e
+
+import (
+	"fmt"
+	"testing"
+
+	nbv1 "github.com/kubeflow/kubeflow/components/notebook-controller/api/v1"
+	"github.com/stretchr/testify/require"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/apimachinery/pkg/util/wait"
+)
+
+func updateTestSuite(t *testing.T) {
+	testCtx, err := NewTestContext()
+	require.NoError(t, err)
+	notebooksForSelectedDeploymentMode := notebooksForScenario(testCtx.testNotebooks, deploymentMode)
+	for _, nbContext := range notebooksForSelectedDeploymentMode {
+		// prepend Notebook name to every subtest
+		t.Run(nbContext.nbObjectMeta.Name, func(t *testing.T) {
+			t.Run("Update Notebook instance", func(t *testing.T) {
+				err = testCtx.testNotebookUpdate(nbContext)
+				require.NoError(t, err, "error updating Notebook object ")
+			})
+			t.Run("Notebook Route Validation After Update", func(t *testing.T) {
+				if deploymentMode == ServiceMesh {
+					t.Skipf("Skipping as it's not relevant for Service Mesh scenario")
+				}
+				err = testCtx.testNotebookRouteCreation(nbContext.nbObjectMeta)
+				require.NoError(t, err, "error testing Route for Notebook after update ")
+			})
+
+			t.Run("Notebook Network Policies Validation After Update", func(t *testing.T) {
+				err = testCtx.testNetworkPolicyCreation(nbContext.nbObjectMeta)
+				require.NoError(t, err, "error testing Network Policies for Notebook after update ")
+			})
+
+			t.Run("Notebook Statefulset Validation After Update", func(t *testing.T) {
+				err = testCtx.testNotebookValidation(nbContext.nbObjectMeta)
+				require.NoError(t, err, "error testing StatefulSet for Notebook after update ")
+			})
+
+			t.Run("Notebook OAuth sidecar Validation After Update", func(t *testing.T) {
+				if deploymentMode == ServiceMesh {
+					t.Skipf("Skipping as it's not relevant for Service Mesh scenario")
+				}
+				err = testCtx.testNotebookOAuthSidecar(nbContext.nbObjectMeta)
+				require.NoError(t, err, "error testing sidecar for Notebook after update ")
+			})
+
+			t.Run("Verify Notebook Traffic After Update", func(t *testing.T) {
+				err = testCtx.testNotebookTraffic(nbContext.nbObjectMeta)
+				require.NoError(t, err, "error testing Notebook traffic after update ")
+			})
+		})
+	}
+}
+
+func (tc *testContext) testNotebookUpdate(nbContext notebookContext) error {
+	notebookLookupKey := types.NamespacedName{Name: nbContext.nbObjectMeta.Name, Namespace: nbContext.nbObjectMeta.Namespace}
+	updatedNotebook := &nbv1.Notebook{}
+
+	err := tc.customClient.Get(tc.ctx, notebookLookupKey, updatedNotebook)
+	if err != nil {
+		return fmt.Errorf("error getting Notebook %s: %v", nbContext.nbObjectMeta.Name, err)
+	}
+
+	// Example update: Change the Notebook image
+	updatedNotebook.Spec.Template.Spec.Containers[0].Image = "new-image:latest"
+
+	err = tc.customClient.Update(tc.ctx, updatedNotebook)
+	if err != nil {
+		return fmt.Errorf("error updating Notebook %s: %v", updatedNotebook.Name, err)
+	}
+
+	// Wait for the update to be applied
+	err = wait.Poll(tc.resourceRetryInterval, tc.resourceCreationTimeout, func() (done bool, err error) {
+		note := &nbv1.Notebook{}
+		err = tc.customClient.Get(tc.ctx, notebookLookupKey, note)
+		if err != nil {
+			return false, err
+		}
+		if note.Spec.Template.Spec.Containers[0].Image == "new-image:latest" {
+			return true, nil
+		}
+		return false, nil
+	})
+
+	if err != nil {
+		return fmt.Errorf("error validating notebook update: %s", err)
+	}
+	return nil
+}