[YUNIKORN-1281] Add e2e tests (#460)

Closes: #460

Signed-off-by: Peter Bacsko <pbacsko@cloudera.com>
diff --git a/test/e2e/framework/helpers/yunikorn/yk_utils.go b/test/e2e/framework/helpers/yunikorn/yk_utils.go
index 6be2044..acb01f8 100644
--- a/test/e2e/framework/helpers/yunikorn/yk_utils.go
+++ b/test/e2e/framework/helpers/yunikorn/yk_utils.go
@@ -25,21 +25,15 @@
 )
 
 type ResourceUsage struct {
-	memory int64
-	vCPU   int64
+	resource map[string]int64
 }
 
 func (r *ResourceUsage) ParseResourceUsage(resource map[string]int64) {
-	r.memory = resource["memory"]
-	r.vCPU = resource["vcore"]
+	r.resource = resource
 }
 
-func (r *ResourceUsage) GetMemory() int64 {
-	return r.memory
-}
-
-func (r *ResourceUsage) GetVCPU() int64 {
-	return r.vCPU
+func (r *ResourceUsage) GetResourceValue(resourceName string) int64 {
+	return r.resource[resourceName]
 }
 
 func GetYKUrl() string {
diff --git a/test/e2e/queue_quota_mgmt/queue_quota_mgmt_test.go b/test/e2e/queue_quota_mgmt/queue_quota_mgmt_test.go
index d244049..55e5f37 100644
--- a/test/e2e/queue_quota_mgmt/queue_quota_mgmt_test.go
+++ b/test/e2e/queue_quota_mgmt/queue_quota_mgmt_test.go
@@ -21,13 +21,13 @@
 import (
 	"fmt"
 	"math"
-	"strconv"
 	"time"
 
 	"github.com/apache/yunikorn-core/pkg/webservice/dao"
 	"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/common"
 	"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/k8s"
 	"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/yunikorn"
+	siCommon "github.com/apache/yunikorn-scheduler-interface/lib/go/common"
 
 	v1 "k8s.io/api/core/v1"
 )
@@ -37,11 +37,18 @@
 	var restClient yunikorn.RClient
 	var err error
 	var sleepRespPod *v1.Pod
-	var maxCPU int64
-	var maxMem int64
 	var maxCPUAnnotation = "yunikorn.apache.org/namespace.max.cpu"
 	var maxMemAnnotation = "yunikorn.apache.org/namespace.max.memory"
-	var annotations map[string]string
+	var quotaAnnotation = "yunikorn.apache.org/namespace.quota"
+	var oldAnnotations map[string]string
+	var expectedOldAnnotations map[string]int64
+	var newAnnotations map[string]string
+	var expectedNewAnnotations map[string]int64
+	var oldAndNewAnnotations map[string]string
+	var expectedOldAndNewAnnotations map[string]int64
+	var emptyOldAnnotations map[string]string
+	var wrongOldAnnotations map[string]string
+	var wrongNewAnnotations map[string]string
 	var ns string
 	var queueInfo *dao.PartitionQueueDAOInfo
 	var maxResource yunikorn.ResourceUsage
@@ -55,23 +62,53 @@
 		Ω(kClient.SetClient()).To(BeNil())
 		// Initializing rest client
 		restClient = yunikorn.RClient{}
-		annotations = map[string]string{}
 		pods = []string{}
+		oldAnnotations = map[string]string{}
+		newAnnotations = map[string]string{}
+		oldAndNewAnnotations = map[string]string{}
+		emptyOldAnnotations = map[string]string{}
+		wrongOldAnnotations = map[string]string{}
+		wrongNewAnnotations = map[string]string{}
+		expectedOldAnnotations = map[string]int64{}
+		expectedNewAnnotations = map[string]int64{}
+		expectedOldAndNewAnnotations = map[string]int64{}
+		oldAnnotations[maxCPUAnnotation] = "300m"
+		oldAnnotations[maxMemAnnotation] = "300M"
+		newAnnotations[quotaAnnotation] = "{\"cpu\": \"400m\", \"memory\": \"400M\", \"nvidia.com/gpu\": \"1\"}"
+		oldAndNewAnnotations[maxCPUAnnotation] = "300m"
+		oldAndNewAnnotations[maxMemAnnotation] = "300M"
+		oldAndNewAnnotations[quotaAnnotation] = "{\"cpu\": \"600m\", \"memory\": \"600M\", \"nvidia.com/gpu\": \"2\"}"
+		emptyOldAnnotations[maxCPUAnnotation] = "nil"
+		emptyOldAnnotations[maxMemAnnotation] = "nil"
+		wrongOldAnnotations[maxCPUAnnotation] = "a"
+		wrongOldAnnotations[maxMemAnnotation] = "b"
+		wrongNewAnnotations[quotaAnnotation] = "{\"cpu\": \"error\", \"memory\": \"error\"}"
+		expectedOldAnnotations[siCommon.CPU] = 300
+		expectedOldAnnotations[siCommon.Memory] = 300
+		expectedNewAnnotations[siCommon.CPU] = 400
+		expectedNewAnnotations[siCommon.Memory] = 400
+		expectedNewAnnotations["nvidia.com/gpu"] = 1
+		expectedOldAndNewAnnotations[siCommon.CPU] = 600
+		expectedOldAndNewAnnotations[siCommon.Memory] = 600
+		expectedOldAndNewAnnotations["nvidia.com/gpu"] = 2
 	})
 
 	/*
-		1. Create a namespace with annotations set to max CPU and max Mem
+		1. Create a namespace with different types of annotations set to max CPU, max Mem, max GPU etc
 		2. Submit pod with CPU and Mem requests
 		3. Verify the CPU and memory usage reported by Yunikorn is accurate.
 		4. Submit another pod with resource request that exceeds the limits
 		5. Verify the pod is in un-schedulable state
 		6. Wait until pod-1 completes and verify that pod-2 is scheduled now.
 	*/
-	It("Verify_Queue_Quota_Allocation", func() {
-		maxMem = 300
-		maxCPU = 300
-		annotations[maxCPUAnnotation] = strconv.FormatInt(maxCPU, 10) + "m"
-		annotations[maxMemAnnotation] = strconv.FormatInt(maxMem, 10) + "M"
+	verifyQuotaAllocationUsingAnnotation := func(annotations map[string]string, resources map[string]int64, podCount int64) {
+		maxCPU := resources[siCommon.CPU]
+		maxMem := resources[siCommon.Memory]
+		maxGPU := int64(0)
+		ok := false
+		if maxGPU, ok = resources["nvidia.com/gpu"]; ok {
+			maxGPU = resources["nvidia.com/gpu"]
+		}
 
 		ns = "ns-" + common.RandSeq(10)
 		By(fmt.Sprintf("create %s namespace with maxCPU: %dm and maxMem: %dM", ns, maxCPU, maxMem))
@@ -84,7 +121,8 @@
 		var reqMem int64 = 100
 
 		var iter int64
-		for iter = 1; iter <= 3; iter++ {
+		nextPod := podCount + 1
+		for iter = 1; iter <= podCount; iter++ {
 			sleepPodConfigs := k8s.SleepPodConfig{NS: ns, Time: 60, CPU: reqCPU, Mem: reqMem}
 			sleepObj, podErr := k8s.InitSleepPod(sleepPodConfigs)
 			Ω(podErr).NotTo(HaveOccurred())
@@ -108,38 +146,39 @@
 			usedPerctResource.ParseResourceUsage(queueInfo.AbsUsedCapacity)
 
 			By(fmt.Sprintf("App-%d: Verify max capacity on the queue is accurate", iter))
-			Ω(maxResource.GetVCPU()).Should(Equal(maxCPU))
-			Ω(maxResource.GetMemory()).Should(Equal(maxMem * 1000 * 1000))
+			Ω(maxResource.GetResourceValue(siCommon.CPU)).Should(Equal(maxCPU))
+			Ω(maxResource.GetResourceValue(siCommon.Memory)).Should(Equal(maxMem * 1000 * 1000))
+			Ω(maxResource.GetResourceValue("nvidia.com/gpu")).Should(Equal(maxGPU))
 
 			By(fmt.Sprintf("App-%d: Verify used capacity on the queue is accurate after 1st pod deployment", iter))
-			Ω(usedResource.GetVCPU()).Should(Equal(reqCPU * iter))
-			Ω(usedResource.GetMemory()).Should(Equal(reqMem * iter * 1000 * 1000))
+			Ω(usedResource.GetResourceValue(siCommon.CPU)).Should(Equal(reqCPU * iter))
+			Ω(usedResource.GetResourceValue(siCommon.Memory)).Should(Equal(reqMem * iter * 1000 * 1000))
 
 			var perctCPU = int64(math.Floor((float64(reqCPU*iter) / float64(maxCPU)) * 100))
 			var perctMem = int64(math.Floor((float64(reqMem*iter) / float64(maxMem)) * 100))
 			By(fmt.Sprintf("App-%d: Verify used percentage capacity on the queue is accurate after 1st pod deployment", iter))
-			Ω(usedPerctResource.GetVCPU()).Should(Equal(perctCPU))
-			Ω(usedPerctResource.GetMemory()).Should(Equal(perctMem))
+			Ω(usedPerctResource.GetResourceValue(siCommon.CPU)).Should(Equal(perctCPU))
+			Ω(usedPerctResource.GetResourceValue(siCommon.Memory)).Should(Equal(perctMem))
 		}
 
-		By("App-4: Submit another app which exceeds the queue quota limitation")
+		By(fmt.Sprintf("App-%d: Submit another app which exceeds the queue quota limitation", nextPod))
 		sleepPodConfigs := k8s.SleepPodConfig{NS: ns, Time: 60, CPU: reqCPU, Mem: reqMem}
 		sleepObj, app4PodErr := k8s.InitSleepPod(sleepPodConfigs)
 		Ω(app4PodErr).NotTo(HaveOccurred())
 
-		By(fmt.Sprintf("App-4: Deploy the sleep app:%s to %s namespace", sleepObj.Name, ns))
+		By(fmt.Sprintf("App-%d: Deploy the sleep app:%s to %s namespace", nextPod, sleepObj.Name, ns))
 		sleepRespPod, err = kClient.CreatePod(sleepObj, ns)
 		Ω(err).NotTo(HaveOccurred())
 		pods = append(pods, sleepObj.Name)
 
-		By(fmt.Sprintf("App-4: Verify app:%s in accepted state", sleepObj.Name))
+		By(fmt.Sprintf("App-%d: Verify app:%s in accepted state", nextPod, sleepObj.Name))
 		// Wait for pod to move to accepted state
 		err = restClient.WaitForAppStateTransition("default", "root."+ns, sleepRespPod.ObjectMeta.Labels["applicationId"],
 			yunikorn.States().Application.Accepted,
-			120)
+			240)
 		Ω(err).NotTo(HaveOccurred())
 
-		By(fmt.Sprintf("Pod-4: Verify pod:%s is in pending state", sleepObj.Name))
+		By(fmt.Sprintf("Pod-%d: Verify pod:%s is in pending state", nextPod, sleepObj.Name))
 		err = kClient.WaitForPodPending(ns, sleepObj.Name, time.Duration(60)*time.Second)
 		Ω(err).NotTo(HaveOccurred())
 
@@ -148,10 +187,22 @@
 		err = kClient.WaitForPodSucceeded(ns, pods[0], time.Duration(360)*time.Second)
 		Ω(err).NotTo(HaveOccurred())
 
-		By(fmt.Sprintf("Pod-4: Verify Pod:%s moved to running state", sleepObj.Name))
+		By(fmt.Sprintf("Pod-%d: Verify Pod:%s moved to running state", nextPod, sleepObj.Name))
 		Ω(kClient.WaitForPodRunning(sleepRespPod.Namespace, sleepRespPod.Name, time.Duration(60)*time.Second)).NotTo(HaveOccurred())
 
-	}, 360)
+	}
+
+	It("Verify_Queue_Quota_Allocation_Using_Old_Annotations", func() {
+		verifyQuotaAllocationUsingAnnotation(oldAnnotations, expectedOldAnnotations, 3)
+	})
+
+	It("Verify_Queue_Quota_Allocation_Using_New_Annotations", func() {
+		verifyQuotaAllocationUsingAnnotation(newAnnotations, expectedNewAnnotations, 4)
+	})
+
+	It("Verify_Queue_Quota_Allocation_Using_Old_New_Annotations", func() {
+		verifyQuotaAllocationUsingAnnotation(oldAndNewAnnotations, expectedOldAndNewAnnotations, 6)
+	})
 
 	It("Verify_NS_Without_Quota", func() {
 		ns = "ns-" + common.RandSeq(10)
@@ -161,9 +212,7 @@
 		Ω(ns1.Status.Phase).To(Equal(v1.NamespaceActive))
 	})
 
-	DescribeTable("", func(x string, y string) {
-		annotations[maxCPUAnnotation] = x
-		annotations[maxMemAnnotation] = y
+	DescribeTable("", func(annotations map[string]string) {
 		ns = "ns-" + common.RandSeq(10)
 		By(fmt.Sprintf("create %s namespace with empty annotations", ns))
 		ns1, err1 := kClient.CreateNamespace(ns, annotations)
@@ -181,8 +230,8 @@
 		Ω(err).NotTo(HaveOccurred())
 		Ω(kClient.WaitForPodRunning(sleepRespPod.Namespace, sleepRespPod.Name, time.Duration(60)*time.Second)).NotTo(HaveOccurred())
 	},
-		Entry("Verify_NS_With_Empty_Quota_Annotations", nil, nil),
-		Entry("Verify_NS_With_Wrong_Quota_Annotations", "a", "b"),
+		Entry("Verify_NS_With_Empty_Quota_Annotations", emptyOldAnnotations),
+		Entry("Verify_NS_With_Wrong_Quota_Annotations", wrongOldAnnotations),
 	)
 
 	// Hierarchical Queues - Quota enforcement