Update pytorch_mcad_test

tests/integration/pytorch_mcad_test.go

@@ -20,14 +20,11 @@ import (
 	"testing"
 
 	. "github.com/onsi/gomega"
+	"github.com/opendatahub-io/distributed-workloads/tests/integration/support"
 	. "github.com/project-codeflare/codeflare-common/support"
 	mcadv1beta1 "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/apis/controller/v1beta1"
 
-	corev1 "k8s.io/api/core/v1"
 	rbacv1 "k8s.io/api/rbac/v1"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-
-	"github.com/opendatahub-io/distributed-workloads/tests/integration/support"
 )
 
 func TestMnistPyTorchMCAD(t *testing.T) {
@@ -38,23 +35,10 @@ func TestMnistPyTorchMCAD(t *testing.T) {
 
 	// Test configuration
 	jupyterNotebookConfigMapFileName := "mnist_mcad_mini.ipynb"
-	config := &corev1.ConfigMap{
-		TypeMeta: metav1.TypeMeta{
-			APIVersion: corev1.SchemeGroupVersion.String(),
-			Kind:       "ConfigMap",
-		},
-		ObjectMeta: metav1.ObjectMeta{
-			Name: "notebooks-mcad",
-		},
-		BinaryData: map[string][]byte{
-			// MNIST MCAD Notebook
-			jupyterNotebookConfigMapFileName: ReadFile(test, "resources/mnist_mcad_mini.ipynb"),
-		},
-		Immutable: Ptr(true),
-	}
-	config, err := test.Client().Core().CoreV1().ConfigMaps(namespace.Name).Create(test.Ctx(), config, metav1.CreateOptions{})
-	test.Expect(err).NotTo(HaveOccurred())
-	test.T().Logf("Created ConfigMap %s/%s successfully", config.Namespace, config.Name)
+	config := CreateConfigMap(test, namespace.Name, map[string][]byte{
+		// MNIST MCAD Notebook
+		jupyterNotebookConfigMapFileName: ReadFile(test, "resources/mnist_mcad_mini.ipynb"),
+	})
 
 	// Create RBAC, retrieve token for user with limited rights
 	policyRules := []rbacv1.PolicyRule{

tests/integration/support/config.go

@@ -18,11 +18,15 @@ package support
 
 import (
 	"os"
+
+	. "github.com/project-codeflare/codeflare-common/support"
 )
 
 const (
 	// The environment variable for namespace where ODH is installed to.
 	odhNamespaceEnvVar = "ODH_NAMESPACE"
+	// The environment variable for ODH Notebook ImageStream name
+	notebookImageStreamName = "NOTEBOOK_IMAGE_STREAM_NAME"
 )
 
 func GetOpenDataHubNamespace() string {
@@ -35,3 +39,11 @@ func lookupEnvOrDefault(key, value string) string {
 	}
 	return value
 }
+
+func GetNotebookImageStreamName(t Test) string {
+	isName, ok := os.LookupEnv(notebookImageStreamName)
+	if !ok {
+		t.T().Fatalf("Expected environment variable %s not found, please use this environment variable to specify what ImageStream to use for Notebook.", notebookImageStreamName)
+	}
+	return isName
+}

tests/integration/support/notebook.go

@@ -21,20 +21,23 @@ import (
 	"embed"
 	"html/template"
 
-	. "github.com/onsi/gomega"
+	gomega "github.com/onsi/gomega"
 	. "github.com/project-codeflare/codeflare-common/support"
 
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 	"k8s.io/apimachinery/pkg/util/yaml"
+
+	imagev1 "github.com/openshift/api/image/v1"
 )
 
 //go:embed resources/*
 var files embed.FS
 
+const recommendedTagAnnotation = "opendatahub.io/workbench-image-recommended"
+
 var notebookResource = schema.GroupVersionResource{Group: "kubeflow.org", Version: "v1", Resource: "notebooks"}
 
 type NotebookProps struct {
@@ -43,36 +46,20 @@ type NotebookProps struct {
 	KubernetesBearerToken     string
 	Namespace                 string
 	OpenDataHubNamespace      string
-	CodeFlareImageStreamTag   string
+	ImageStreamName           string
+	ImageStreamTag            string
 	NotebookConfigMapName     string
 	NotebookConfigMapFileName string
 	NotebookPVC               string
 }
 
 func CreateNotebook(test Test, namespace *corev1.Namespace, notebookToken, jupyterNotebookConfigMapName, jupyterNotebookConfigMapFileName string) {
 	// Create PVC for Notebook
-	notebookPVC := &corev1.PersistentVolumeClaim{
-		TypeMeta: metav1.TypeMeta{
-			APIVersion: corev1.SchemeGroupVersion.String(),
-			Kind:       "PersistentVolumeClaim",
-		},
-		ObjectMeta: metav1.ObjectMeta{
-			Name: "jupyterhub-nb-kube-3aadmin-pvc",
-		},
-		Spec: corev1.PersistentVolumeClaimSpec{
-			AccessModes: []corev1.PersistentVolumeAccessMode{
-				corev1.ReadWriteOnce,
-			},
-			Resources: corev1.ResourceRequirements{
-				Requests: corev1.ResourceList{
-					corev1.ResourceStorage: resource.MustParse("10Gi"),
-				},
-			},
-		},
-	}
-	notebookPVC, err := test.Client().Core().CoreV1().PersistentVolumeClaims(namespace.Name).Create(test.Ctx(), notebookPVC, metav1.CreateOptions{})
-	test.Expect(err).NotTo(HaveOccurred())
-	test.T().Logf("Created PersistentVolumeClaim %s/%s successfully", notebookPVC.Namespace, notebookPVC.Name)
+	notebookPVC := CreatePersistentVolumeClaim(test, namespace.Name, "10Gi", corev1.ReadWriteOnce)
+
+	// Retrieve ImageStream tag for
+	is := GetImageStream(test, GetOpenDataHubNamespace(), GetNotebookImageStreamName(test))
+	recommendedTagName := getRecommendedImageStreamTag(test, is)
 
 	// Read the Notebook CR from resources and perform replacements for custom values using go template
 	notebookProps := NotebookProps{
@@ -81,25 +68,36 @@ func CreateNotebook(test Test, namespace *corev1.Namespace, notebookToken, jupyt
 		KubernetesBearerToken:     notebookToken,
 		Namespace:                 namespace.Name,
 		OpenDataHubNamespace:      GetOpenDataHubNamespace(),
-		CodeFlareImageStreamTag:   GetODHCodeFlareImageStreamTag(test),
+		ImageStreamName:           GetNotebookImageStreamName(test),
+		ImageStreamTag:            recommendedTagName,
 		NotebookConfigMapName:     jupyterNotebookConfigMapName,
 		NotebookConfigMapFileName: jupyterNotebookConfigMapFileName,
 		NotebookPVC:               notebookPVC.Name,
 	}
 	notebookTemplate, err := files.ReadFile("resources/custom-nb-small.yaml")
-	test.Expect(err).NotTo(HaveOccurred())
+	test.Expect(err).NotTo(gomega.HaveOccurred())
 	parsedNotebookTemplate, err := template.New("notebook").Parse(string(notebookTemplate))
-	test.Expect(err).NotTo(HaveOccurred())
+	test.Expect(err).NotTo(gomega.HaveOccurred())
 
 	// Filter template and store results to the buffer
 	notebookBuffer := new(bytes.Buffer)
 	err = parsedNotebookTemplate.Execute(notebookBuffer, notebookProps)
-	test.Expect(err).NotTo(HaveOccurred())
+	test.Expect(err).NotTo(gomega.HaveOccurred())
 
 	// Create Notebook CR
 	notebookCR := &unstructured.Unstructured{}
 	err = yaml.NewYAMLOrJSONDecoder(notebookBuffer, 8192).Decode(notebookCR)
-	test.Expect(err).NotTo(HaveOccurred())
+	test.Expect(err).NotTo(gomega.HaveOccurred())
 	_, err = test.Client().Dynamic().Resource(notebookResource).Namespace(namespace.Name).Create(test.Ctx(), notebookCR, metav1.CreateOptions{})
-	test.Expect(err).NotTo(HaveOccurred())
+	test.Expect(err).NotTo(gomega.HaveOccurred())
+}
+
+func getRecommendedImageStreamTag(test Test, is *imagev1.ImageStream) (tagName string) {
+	for _, tag := range is.Spec.Tags {
+		if tag.Annotations[recommendedTagAnnotation] == "true" {
+			return tag.Name
+		}
+	}
+	test.T().Fatalf("tag with annotation '%s' not found in ImageStream %s", recommendedTagAnnotation, is.Name)
+	return
 }