refactored to check for specific machine names

Author: Fiona-Waters

test/e2e/instascale_test.go

@@ -1,7 +1,6 @@
 package e2e
 
 import (
-	"sync"
 	"testing"
 	"time"
 
@@ -14,13 +13,6 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
-var (
-	machinePoolsExist     bool
-	numInitialNodePools   int
-	numInitialMachineSets int
-	wg                    = &sync.WaitGroup{}
-)
-
 func TestInstascale(t *testing.T) {
 
 	test := With(t)
@@ -65,25 +57,25 @@ func TestInstascale(t *testing.T) {
 	}
 	defer connection.Close()
 
-	machinePoolsExist = true
 	// check existing cluster resources
-	numInitialMachinePools, err := MachinePoolsCount(connection)
-	if err != nil {
-		test.T().Errorf("Unable to count machine pools - Error : %v", err)
-	}
+	machinePoolsExist, err := MachinePoolsExist(connection)
+	test.Expect(err).NotTo(HaveOccurred())
+	nodePoolsExist, err := NodePoolsExist(connection)
+	test.Expect(err).NotTo(HaveOccurred())
 
-	if numInitialMachinePools == 0 {
-		machinePoolsExist = false
-		numInitialNodePools, err = NodePoolsCount(connection)
-		if err != nil {
-			test.T().Errorf("Unable to count node pools - Error : %v", err)
-		}
-		if numInitialNodePools == 0 {
-			numInitialMachineSets, err = MachineSetsCount()
-			if err != nil {
-				test.T().Errorf("Unable to count machine sets - Error : %v", err)
-			}
-		}
+	if machinePoolsExist {
+		// look for machine pool with aw name - expect not to find it
+		foundMachinePool, err := CheckMachinePools(connection, "test-instascale")
+		test.Expect(err).NotTo(HaveOccurred())
+		test.Expect(foundMachinePool).To(BeFalse())
+	} else if nodePoolsExist {
+		// look for node pool with aw name - expect not to find it
+		foundNodePool, err := CheckNodePools(connection, "test-instascale")
+		test.Expect(err).NotTo(HaveOccurred())
+		test.Expect(foundNodePool).To(BeFalse())
+	} else {
+		// TODO update to foundMachineSet
+		
 	}
 
 	// Batch Job
@@ -103,13 +95,13 @@ func TestInstascale(t *testing.T) {
 				Spec: corev1.PodSpec{
 					Containers: []corev1.Container{
 						{
-							Name:    "job",
-							Image:   GetPyTorchImage(),
+							Name:  "job",
+							Image: GetPyTorchImage(),
 							Env: []corev1.EnvVar{
 								corev1.EnvVar{Name: "PYTHONUSERBASE", Value: "/test2"},
 							},
 							Command: []string{"/bin/sh", "-c", "pip install -r /test/requirements.txt && torchrun /test/mnist.py"},
-							Args: []string{"$PYTHONUSERBASE"},
+							Args:    []string{"$PYTHONUSERBASE"},
 							VolumeMounts: []corev1.VolumeMount{
 								{
 									Name:      "test",
@@ -184,46 +176,41 @@ func TestInstascale(t *testing.T) {
 								},
 							},
 						},
-						GenericTemplate: Raw(test, job),
+						GenericTemplate:  Raw(test, job),
+						CompletionStatus: "Complete",
 					},
 				},
 			},
 		},
 	}
 
-		_, err = test.Client().MCAD().WorkloadV1beta1().AppWrappers(namespace.Name).Create(test.Ctx(), aw, metav1.CreateOptions{})
+	_, err = test.Client().MCAD().WorkloadV1beta1().AppWrappers(namespace.Name).Create(test.Ctx(), aw, metav1.CreateOptions{})
 	test.Expect(err).NotTo(HaveOccurred())
 	test.T().Logf("AppWrapper created successfully %s/%s", aw.Namespace, aw.Name)
 
 	test.Eventually(AppWrapper(test, namespace, aw.Name), TestTimeoutShort).
 		Should(WithTransform(AppWrapperState, Equal(mcadv1beta1.AppWrapperStateActive)))
 
-	// wait for required resources to be created before checking them again
-	time.Sleep(TestTimeoutShort)
-	if !machinePoolsExist {
-		numNodePools, err := NodePoolsCount(connection)
-		if err != nil {
-			test.T().Errorf("Unable to count node pools - Error : %v", err)
-		}
-		test.Expect(numNodePools).To(BeNumerically(">", numInitialNodePools))
-		test.T().Logf("number of node pools increased from %d to %d", numInitialNodePools, numNodePools)
-
-	} else if machinePoolsExist {
-		numMachinePools, err := MachinePoolsCount(connection)
-		if err != nil {
-			test.T().Errorf("Unable to count machine pools - Error : %v", err)
-		}
-		test.Expect(numMachinePools).To(BeNumerically(">", numInitialMachinePools))
-		test.T().Logf("number of machine pools increased from %d to %d", numInitialMachinePools, numMachinePools)
+	// time.Sleep is used twice throughout the test, each for 30 seconds. Can look into using sync package waitGroup instead if that makes more sense
+	// wait for required resources to scale up before checking them again
+	time.Sleep(TestTimeoutThirtySeconds)
+
+	if machinePoolsExist {
+		// look for machine pool with aw name - expect to find it
+		foundMachinePool, err := CheckMachinePools(connection, "test-instascale")
+		test.Expect(err).NotTo(HaveOccurred())
+		test.Expect(foundMachinePool).To(BeTrue())
+	} else if nodePoolsExist {
+		// look for node pool with aw name - expect to find it
+		foundNodePool, err := CheckNodePools(connection, "test-instascale")
+		test.Expect(err).NotTo(HaveOccurred())
+		test.Expect(foundNodePool).To(BeTrue())
 	} else {
-		numMachineSets, err := MachineSetsCount()
-		if err != nil {
-			test.T().Errorf("Unable to count machine sets - Error : %v", err)
-		}
-		test.Expect(numMachineSets).To(BeNumerically(">", numInitialMachineSets))
-		test.T().Logf("number of machine sets increased from %d to %d", numInitialMachineSets, numMachineSets)
+		// TODO update to foundMachineSet
+		
 	}
-	
+
+	// Assert that the job has completed
 	test.T().Logf("Waiting for Job %s/%s to complete", job.Namespace, job.Name)
 	test.Eventually(Job(test, job.Namespace, job.Name), TestTimeoutLong).Should(
 		Or(
@@ -235,30 +222,24 @@ func TestInstascale(t *testing.T) {
 	test.Expect(GetJob(test, job.Namespace, job.Name)).
 		To(WithTransform(ConditionStatus(batchv1.JobComplete), Equal(corev1.ConditionTrue)))
 
-	// AppWrapper not being updated to complete once job is finished
-
-	time.Sleep(TestTimeoutMedium)
-	if !machinePoolsExist {
-		numNodePoolsFinal, err := NodePoolsCount(connection)
-		if err != nil {
-			test.T().Errorf("Unable to count node pools - Error : %v", err)
-		}
-		test.Expect(numNodePoolsFinal).To(BeNumerically("==", numInitialNodePools))
-		test.T().Logf("number of machine pools decreased")
-
-	} else if machinePoolsExist {
-		numMachinePoolsFinal, err := MachinePoolsCount(connection)
-		if err != nil {
-			test.T().Errorf("Unable to count machine pools - Error : %v", err)
-		}
-		test.Expect(numMachinePoolsFinal).To(BeNumerically("==", numInitialMachinePools))
-		test.T().Logf("number of machine pools decreased")
+	test.Eventually(AppWrapper(test, namespace, aw.Name), TestTimeoutShort).
+		Should(WithTransform(AppWrapperState, Equal(mcadv1beta1.AppWrapperStateCompleted)))
+
+	// allow time for the resources to scale down before checking them again
+	time.Sleep(TestTimeoutThirtySeconds)
+
+	if machinePoolsExist {
+		// look for machine pool with aw name - expect to find it
+		foundMachinePool, err := CheckMachinePools(connection, "test-instascale")
+		test.Expect(err).NotTo(HaveOccurred())
+		test.Expect(foundMachinePool).To(BeFalse())
+	} else if nodePoolsExist {
+		// look for node pool with aw name - expect to find it
+		foundNodePool, err := CheckNodePools(connection, "test-instascale")
+		test.Expect(err).NotTo(HaveOccurred())
+		test.Expect(foundNodePool).To(BeFalse())
 	} else {
-		numMachineSetsFinal, err := MachineSetsCount()
-		if err != nil {
-			test.T().Errorf("Unable to count machine sets - Error : %v", err)
-		}
-		test.Expect(numMachineSetsFinal).To(BeNumerically("==", numInitialMachineSets))
-		test.T().Logf("number of machine sets decreased")
+		// TODO update to foundMachineSet
+		
 	}
 }

test/support/clusterpools.go

@@ -4,15 +4,17 @@ import (
 	"context"
 	"fmt"
 	"os"
+	"strings"
 
 	ocmsdk "github.com/openshift-online/ocm-sdk-go"
+	cmv1 "github.com/openshift-online/ocm-sdk-go/clustersmgmt/v1"
 	mapiclientset "github.com/openshift/client-go/machine/clientset/versioned"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 var (
-	ClusterID         string = os.Getenv("CLUSTERID")
-	machineClient     mapiclientset.Interface
+	ClusterID     string = os.Getenv("CLUSTERID")
+	machineClient mapiclientset.Interface
 )
 
 const (
@@ -39,30 +41,50 @@ func CreateOCMConnection(secret string) (*ocmsdk.Connection, error) {
 	return connection, nil
 }
 
+func MachinePoolsExist(connection *ocmsdk.Connection) (bool, error) {
+	machinePools := connection.ClustersMgmt().V1().Clusters().Cluster(ClusterID).MachinePools()
+	return machinePools != nil, nil
+}
+
+func NodePoolsExist(connection *ocmsdk.Connection) (bool, error) {
+	nodePools := connection.ClustersMgmt().V1().Clusters().Cluster(ClusterID).NodePools()
+	return nodePools != nil, nil
+}
 
-func MachinePoolsCount(connection *ocmsdk.Connection) (numMachinePools int, err error) {
+func CheckMachinePools(connection *ocmsdk.Connection, awName string) (foundMachinePool bool, err error) {
 	machinePoolsConnection := connection.ClustersMgmt().V1().Clusters().Cluster(ClusterID).MachinePools().List()
 	machinePoolsListResponse, err := machinePoolsConnection.Send()
 	if err != nil {
-		return 0, fmt.Errorf("unable to send request, error: %v", err)
+		return false, fmt.Errorf("unable to send request, error: %v", err)
 	}
 	machinePoolsList := machinePoolsListResponse.Items()
-	numMachinePools = machinePoolsList.Len()
+	machinePoolsList.Range(func(index int, item *cmv1.MachinePool) bool {
+		instanceName, _ := item.GetID()
+		if strings.Contains(instanceName, awName) {
+			foundMachinePool = true
+		}
+		return true
+	})
 
-	return numMachinePools, nil
+	return foundMachinePool, err
 }
 
-func NodePoolsCount(connection *ocmsdk.Connection) (numNodePools int, err error) {
+func CheckNodePools(connection *ocmsdk.Connection, awName string) (foundNodePool bool, err error) {
 	nodePoolsConnection := connection.ClustersMgmt().V1().Clusters().Cluster(ClusterID).NodePools().List()
 	nodePoolsListResponse, err := nodePoolsConnection.SendContext(context.Background())
 	if err != nil {
-		return 0, fmt.Errorf("unable to send request, error: %v", err)
+		return false, fmt.Errorf("unable to send request, error: %v", err)
 	}
 	nodePoolsList := nodePoolsListResponse.Items()
-	numNodePools = nodePoolsList.Len()
-	fmt.Println(numNodePools)
+	nodePoolsList.Range(func(index int, item *cmv1.NodePool) bool {
+		instanceName, _ := item.GetID()
+		if strings.Contains(instanceName, awName) {
+			foundNodePool = true
+		}
+		return true
+	})
 
-	return numNodePools, nil
+	return foundNodePool, err
 }
 
 func MachineSetsCount() (numMachineSets int, err error) {

test/support/support.go

@@ -31,6 +31,7 @@ var (
 	ApplyOptions = metav1.ApplyOptions{FieldManager: "codeflare-test", Force: true}
 
 	TestTimeoutShort  = 1 * time.Minute
+	TestTimeoutThirtySeconds = 30 * time.Second
 	TestTimeoutMedium = 2 * time.Minute
 	TestTimeoutLong   = 5 * time.Minute
 )