From 7f7fbfb21b504f55e8291398b5e3636f1a29295c Mon Sep 17 00:00:00 2001 From: Laurentiu Bradin <109964136+z103cb@users.noreply.github.com> Date: Mon, 29 May 2023 15:22:57 +0300 Subject: [PATCH 1/3] Added e2e tests to emulate the condition. Minor fixes in the e2e script. --- hack/run-e2e-kind.sh | 10 ++++- test/e2e-kuttl/quota-errors/00-assert.yaml | 25 +++++++++++ test/e2e-kuttl/quota-errors/01-assert.yaml | 24 +++++++++++ test/e2e-kuttl/quota-errors/02-assert.yaml | 5 +++ test/e2e-kuttl/quota-errors/02-install.yaml | 4 ++ test/e2e-kuttl/quota-errors/03-assert.yaml | 27 ++++++++++++ test/e2e-kuttl/quota-errors/03-install.yaml | 42 +++++++++++++++++++ .../quota-errors/04-delete-app-wrapper.yaml | 9 ++++ .../quota-errors/05-reapply-app-wrapper.yaml | 10 +++++ test/e2e-kuttl/quota-errors/99-cleanup.yaml | 9 ++++ test/e2e-kuttl/quota-forest/99-cleanup.yaml | 9 ++++ 11 files changed, 173 insertions(+), 1 deletion(-) create mode 100644 test/e2e-kuttl/quota-errors/00-assert.yaml create mode 100644 test/e2e-kuttl/quota-errors/01-assert.yaml create mode 100644 test/e2e-kuttl/quota-errors/02-assert.yaml create mode 100644 test/e2e-kuttl/quota-errors/02-install.yaml create mode 100644 test/e2e-kuttl/quota-errors/03-assert.yaml create mode 100644 test/e2e-kuttl/quota-errors/03-install.yaml create mode 100644 test/e2e-kuttl/quota-errors/04-delete-app-wrapper.yaml create mode 100644 test/e2e-kuttl/quota-errors/05-reapply-app-wrapper.yaml create mode 100644 test/e2e-kuttl/quota-errors/99-cleanup.yaml create mode 100644 test/e2e-kuttl/quota-forest/99-cleanup.yaml diff --git a/hack/run-e2e-kind.sh b/hack/run-e2e-kind.sh index 513d068b7..8d5596329 100755 --- a/hack/run-e2e-kind.sh +++ b/hack/run-e2e-kind.sh @@ -33,6 +33,7 @@ export CLUSTER_CONTEXT="--name test" # Using older image due to older version of kubernetes cluster" export IMAGE_ECHOSERVER="kicbase/echo-server:1.0" export IMAGE_UBUNTU_LATEST="ubuntu:latest" +export IMAGE_UBI_LATEST="registry.access.redhat.com/ubi8/ubi:latest" export KIND_OPT=${KIND_OPT:=" --config ${ROOT_DIR}/hack/e2e-kind-config.yaml"} export KA_BIN=_output/bin export WAIT_TIME="20s" @@ -220,6 +221,13 @@ function kind-up-cluster { exit 1 fi + docker pull ${IMAGE_UBI_LATEST} + if [ $? -ne 0 ] + then + echo "Failed to pull ${IMAGE_UBI_LATEST}" + exit 1 + fi + if [[ "$MCAD_IMAGE_PULL_POLICY" = "Always" ]] then docker pull ${IMAGE_MCAD} @@ -236,7 +244,7 @@ function kind-up-cluster { fi docker images - for image in ${IMAGE_ECHOSERVER} ${IMAGE_UBUNTU_LATEST} ${IMAGE_MCAD} + for image in ${IMAGE_ECHOSERVER} ${IMAGE_UBUNTU_LATEST} ${IMAGE_MCAD} ${IMAGE_UBI_LATEST} do kind load docker-image ${image} ${CLUSTER_CONTEXT} if [ $? -ne 0 ] diff --git a/test/e2e-kuttl/quota-errors/00-assert.yaml b/test/e2e-kuttl/quota-errors/00-assert.yaml new file mode 100644 index 000000000..853505926 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/00-assert.yaml @@ -0,0 +1,25 @@ +# Verify CRDs existence +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + name: appwrappers.mcad.ibm.com +status: + acceptedNames: + kind: AppWrapper + listKind: AppWrapperList + plural: appwrappers + singular: appwrapper + storedVersions: + - v1beta1 +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + name: quotasubtrees.ibm.com +status: + acceptedNames: + kind: QuotaSubtree + singular: quotasubtree + plural: quotasubtrees + storedVersions: + - v1 diff --git a/test/e2e-kuttl/quota-errors/01-assert.yaml b/test/e2e-kuttl/quota-errors/01-assert.yaml new file mode 100644 index 000000000..13ef4b424 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/01-assert.yaml @@ -0,0 +1,24 @@ +# Verify subtree creations +apiVersion: ibm.com/v1 +kind: QuotaSubtree +metadata: + name: context-root + namespace: kube-system + labels: + tree: quota_context +--- +apiVersion: ibm.com/v1 +kind: QuotaSubtree +metadata: + name: service-root + namespace: kube-system + labels: + tree: quota_service +--- +apiVersion: ibm.com/v1 +kind: QuotaSubtree +metadata: + name: context-root-children + namespace: kube-system + labels: + tree: quota_context diff --git a/test/e2e-kuttl/quota-errors/02-assert.yaml b/test/e2e-kuttl/quota-errors/02-assert.yaml new file mode 100644 index 000000000..217845bd5 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/02-assert.yaml @@ -0,0 +1,5 @@ +# Verify test namespace existence +apiVersion: v1 +kind: Namespace +metadata: + name: quota-errors \ No newline at end of file diff --git a/test/e2e-kuttl/quota-errors/02-install.yaml b/test/e2e-kuttl/quota-errors/02-install.yaml new file mode 100644 index 000000000..4daa24099 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/02-install.yaml @@ -0,0 +1,4 @@ +apiVersion: v1 +kind: Namespace +metadata: + name: quota-errors \ No newline at end of file diff --git a/test/e2e-kuttl/quota-errors/03-assert.yaml b/test/e2e-kuttl/quota-errors/03-assert.yaml new file mode 100644 index 000000000..9fced0a93 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/03-assert.yaml @@ -0,0 +1,27 @@ +# Verify AppWrapper was dispatched and pod was created +apiVersion: mcad.ibm.com/v1beta1 +kind: AppWrapper +metadata: + name: deployment-silver-lo-pri-1replica + namespace: quota-errors + labels: + quota_context: "silver" + quota_service: "service-root" +status: + state: Running +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: deployment-silver-lo-pri-1replica + namespace: quota-errors + labels: + app: deployment-silver-lo-pri-1replica + appwrapper.mcad.ibm.com: deployment-silver-lo-pri-1replica + resourceName: deployment-silver-lo-pri-1replica +status: + availableReplicas: 1 + observedGeneration: 1 + readyReplicas: 1 + replicas: 1 + updatedReplicas: 1 \ No newline at end of file diff --git a/test/e2e-kuttl/quota-errors/03-install.yaml b/test/e2e-kuttl/quota-errors/03-install.yaml new file mode 100644 index 000000000..d93d43878 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/03-install.yaml @@ -0,0 +1,42 @@ +apiVersion: mcad.ibm.com/v1beta1 +kind: AppWrapper +metadata: + name: deployment-silver-lo-pri-1replica + namespace: quota-errors + labels: + quota_context: "silver" + quota_service: "service-root" +spec: + resources: + GenericItems: + - replicas: 1 + generictemplate: + apiVersion: apps/v1 + kind: Deployment + metadata: + name: deployment-silver-lo-pri-1replica + namespace: quota-errors + labels: + app: deployment-silver-lo-pri-1replica + spec: + selector: + matchLabels: + app: deployment-silver-lo-pri-1replica + replicas: 1 + template: + metadata: + labels: + app: deployment-silver-lo-pri-1replica + spec: + containers: + - name: deployment-silver-lo-pri-1replica + image: kicbase/echo-server:1.0 + ports: + - containerPort: 80 + resources: + requests: + cpu: 100m + memory: 32Mi + limits: + cpu: 100m + memory: 32Mi diff --git a/test/e2e-kuttl/quota-errors/04-delete-app-wrapper.yaml b/test/e2e-kuttl/quota-errors/04-delete-app-wrapper.yaml new file mode 100644 index 000000000..2dfe173fb --- /dev/null +++ b/test/e2e-kuttl/quota-errors/04-delete-app-wrapper.yaml @@ -0,0 +1,9 @@ +apiVersion: kuttl.dev/v1beta1 +kind: TestStep +apply: [] +assert: [] +error: [] +unitTest: false +delete: [] +commands: + - command: kubectl delete appwrapper -n quota-errors deployment-silver-lo-pri-1replica diff --git a/test/e2e-kuttl/quota-errors/05-reapply-app-wrapper.yaml b/test/e2e-kuttl/quota-errors/05-reapply-app-wrapper.yaml new file mode 100644 index 000000000..be87d1a25 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/05-reapply-app-wrapper.yaml @@ -0,0 +1,10 @@ +apiVersion: kuttl.dev/v1beta1 +kind: TestStep +apply: + - 03-install.yaml +assert: + - 03-assert.yaml +error: [] +unitTest: false +delete: [] +commands: [] \ No newline at end of file diff --git a/test/e2e-kuttl/quota-errors/99-cleanup.yaml b/test/e2e-kuttl/quota-errors/99-cleanup.yaml new file mode 100644 index 000000000..3bc2a2377 --- /dev/null +++ b/test/e2e-kuttl/quota-errors/99-cleanup.yaml @@ -0,0 +1,9 @@ +apiVersion: kuttl.dev/v1beta1 +kind: TestStep +apply: [] +assert: [] +error: [] +unitTest: false +delete: [] +commands: + - command: kubectl delete namespace quota-errors \ No newline at end of file diff --git a/test/e2e-kuttl/quota-forest/99-cleanup.yaml b/test/e2e-kuttl/quota-forest/99-cleanup.yaml new file mode 100644 index 000000000..3998eb375 --- /dev/null +++ b/test/e2e-kuttl/quota-forest/99-cleanup.yaml @@ -0,0 +1,9 @@ +apiVersion: kuttl.dev/v1beta1 +kind: TestStep +apply: [] +assert: [] +error: [] +unitTest: false +delete: [] +commands: + - command: kubectl delete namespace test \ No newline at end of file From 24f0c6a118943546bbd2b3b5d3f173819594e722 Mon Sep 17 00:00:00 2001 From: Laurentiu Bradin <109964136+z103cb@users.noreply.github.com> Date: Tue, 30 May 2023 13:22:48 +0300 Subject: [PATCH 2/3] Incorporated feedback and tests from #398 --- .../qm_lib_backend_with_quotasubt_mgr.go | 24 +++- test/e2e-kuttl/quota-forest/09-assert.yaml | 17 +++ test/e2e-kuttl/quota-forest/09-install.yaml | 133 ++++++++++++++++++ 3 files changed, 167 insertions(+), 7 deletions(-) create mode 100644 test/e2e-kuttl/quota-forest/09-assert.yaml create mode 100644 test/e2e-kuttl/quota-forest/09-install.yaml diff --git a/pkg/controller/quota/quotaforestmanager/qm_lib_backend_with_quotasubt_mgr.go b/pkg/controller/quota/quotaforestmanager/qm_lib_backend_with_quotasubt_mgr.go index 1143186ee..95a727ea3 100644 --- a/pkg/controller/quota/quotaforestmanager/qm_lib_backend_with_quotasubt_mgr.go +++ b/pkg/controller/quota/quotaforestmanager/qm_lib_backend_with_quotasubt_mgr.go @@ -530,20 +530,19 @@ func (qm *QuotaManager) Fits(aw *arbv1.AppWrapper, awResDemands *clusterstateapi klog.V(4).Infof("[Fits] Sending quota allocation request: %#v ", consumerInfo) allocResponse, err := qm.quotaManagerBackend.AllocateForest(QuotaManagerForestName, consumerID) if err != nil { + qm.removeConsumer(consumerID) klog.Errorf("[Fits] Error allocating consumer: %s/%s, err=%#v.", aw.Namespace, aw.Name, err) return false, nil, err.Error() } - if allocResponse != nil && len(strings.TrimSpace(allocResponse.GetMessage())) > 0 { - klog.Errorf("[Fits] Error allocating consumer: %s/%s, msg=%s, err=%#v.", - aw.Namespace, aw.Name, allocResponse.GetMessage(), err) - return false, nil, allocResponse.GetMessage() - } klog.V(4).Infof("[Fits] allocation response received. Quota Allocated: %t, Message: '%s', Preempted app wrappers count: %d", allocResponse.IsAllocated(), strings.TrimSpace(allocResponse.GetMessage()), len(allocResponse.GetPreemptedIds())) doesFit := allocResponse.IsAllocated() + if !doesFit { + qm.removeConsumer(consumerID) + return doesFit, preemptIds, strings.TrimSpace(allocResponse.GetMessage()) + } preemptIds = qm.getAppWrappers(allocResponse.GetPreemptedIds()) - - return doesFit, preemptIds, "" + return doesFit, preemptIds, strings.TrimSpace(allocResponse.GetMessage()) } func (qm *QuotaManager) getAppWrappers(preemptIds []string) []*arbv1.AppWrapper { @@ -614,3 +613,14 @@ func (qm *QuotaManager) Release(aw *arbv1.AppWrapper) bool { return released } +func (qm *QuotaManager) removeConsumer(consumerID string) { + // removing the consumer to allow for the consumer to be added if and when + // the function is called for the same app wrapper + removed, err := qm.quotaManagerBackend.RemoveConsumer(consumerID) + if err != nil { + klog.Warningf("Failed to remove consumer %s, %#v", consumerID, err) + } + if !removed { + klog.Warningf("Failed to remove consumer %s", consumerID) + } +} diff --git a/test/e2e-kuttl/quota-forest/09-assert.yaml b/test/e2e-kuttl/quota-forest/09-assert.yaml new file mode 100644 index 000000000..fc13e7ed6 --- /dev/null +++ b/test/e2e-kuttl/quota-forest/09-assert.yaml @@ -0,0 +1,17 @@ +#Verify AppWrappers finished successfully +--- +apiVersion: mcad.ibm.com/v1beta1 +kind: AppWrapper +metadata: + name: my-job-1 + namespace: test +status: + state: Completed +--- +apiVersion: mcad.ibm.com/v1beta1 +kind: AppWrapper +metadata: + name: my-job-2 + namespace: test +status: + state: Completed diff --git a/test/e2e-kuttl/quota-forest/09-install.yaml b/test/e2e-kuttl/quota-forest/09-install.yaml new file mode 100644 index 000000000..a8f437be2 --- /dev/null +++ b/test/e2e-kuttl/quota-forest/09-install.yaml @@ -0,0 +1,133 @@ +--- +apiVersion: kuttl.dev/v1beta1 +kind: TestStep +commands: + - command: kubectl -n test delete appwrappers,jobs --all +--- +apiVersion: mcad.ibm.com/v1beta1 +kind: AppWrapper +metadata: + name: my-job-1 + namespace: test + labels: + quota_context: bronze + quota_service: service-root +spec: + schedulingSpec: + minAvailable: 1 + resources: + GenericItems: + - replicas: 1 + completionstatus: Complete + custompodresources: + - replicas: 1 + requests: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi + limits: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi + generictemplate: + apiVersion: batch/v1 + kind: Job + metadata: + name: my-job-1 + namespace: test + labels: + appwrapper.mcad.ibm.com: my-job-1 + spec: + parallelism: 1 + completions: 1 + template: + metadata: + name: my-job-1 + namespace: test + labels: + appwrapper.mcad.ibm.com: my-job-1 + spec: + terminationGracePeriodSeconds: 1 + restartPolicy: Never + containers: + - name: ubuntu + image: ubuntu:latest + imagePullPolicy: IfNotPresent + command: + - sh + - -c + - | + sleep 30 + resources: + requests: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi + limits: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi +--- +apiVersion: mcad.ibm.com/v1beta1 +kind: AppWrapper +metadata: + name: my-job-2 + namespace: test + labels: + quota_context: bronze + quota_service: service-root +spec: + schedulingSpec: + minAvailable: 1 + resources: + GenericItems: + - replicas: 1 + completionstatus: Complete + custompodresources: + - replicas: 1 + requests: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi + limits: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi + generictemplate: + apiVersion: batch/v1 + kind: Job + metadata: + name: my-job-2 + namespace: test + labels: + appwrapper.mcad.ibm.com: my-job-2 + spec: + parallelism: 1 + completions: 1 + template: + metadata: + name: my-job-2 + namespace: test + labels: + appwrapper.mcad.ibm.com: my-job-2 + spec: + terminationGracePeriodSeconds: 1 + restartPolicy: Never + containers: + - name: ubuntu + image: ubuntu:latest + imagePullPolicy: IfNotPresent + command: + - sh + - -c + - | + sleep 30 + resources: + requests: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi + limits: + cpu: 900m + nvidia.com/gpu: 0 + memory: 300Mi From f0173af1acf17cd48e61dd7de80ae46e942713d6 Mon Sep 17 00:00:00 2001 From: Laurentiu Bradin <109964136+z103cb@users.noreply.github.com> Date: Tue, 30 May 2023 14:59:11 +0300 Subject: [PATCH 3/3] Added unit tests to emulate a long running consumer --- go.mod | 1 + go.sum | 2 + .../quota-manager/quota/quotamanager_test.go | 168 ++++++++++++++++++ 3 files changed, 171 insertions(+) diff --git a/go.mod b/go.mod index 461b64425..6a19ad0b1 100644 --- a/go.mod +++ b/go.mod @@ -3,6 +3,7 @@ module github.com/project-codeflare/multi-cluster-app-dispatcher go 1.18 require ( + github.com/eapache/go-resiliency v1.3.0 github.com/emicklei/go-restful v2.16.0+incompatible github.com/golang/protobuf v1.4.3 github.com/hashicorp/go-multierror v1.1.1 diff --git a/go.sum b/go.sum index 18055623a..84487dcf2 100644 --- a/go.sum +++ b/go.sum @@ -96,6 +96,8 @@ github.com/docopt/docopt-go v0.0.0-20180111231733-ee0de3bc6815/go.mod h1:WwZ+bS3 github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo= github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= +github.com/eapache/go-resiliency v1.3.0 h1:RRL0nge+cWGlxXbUzJ7yMcq6w2XBEr19dCN6HECGaT0= +github.com/eapache/go-resiliency v1.3.0/go.mod h1:5yPzW0MIvSe0JDsv0v+DvcjEv2FyD6iZYSs1ZI+iQho= github.com/elazarl/goproxy v0.0.0-20180725130230-947c36da3153/go.mod h1:/Zj4wYkgs4iZTTu3o/KG3Itv/qCCa8VVMlb3i9OVuzc= github.com/emicklei/go-restful v0.0.0-20170410110728-ff4f55a20633/go.mod h1:otzb+WCGbkyDHkqmQmT5YD2WR4BBwUdeQoFo8l/7tVs= github.com/emicklei/go-restful v2.9.5+incompatible/go.mod h1:otzb+WCGbkyDHkqmQmT5YD2WR4BBwUdeQoFo8l/7tVs= diff --git a/pkg/quotaplugins/quota-forest/quota-manager/quota/quotamanager_test.go b/pkg/quotaplugins/quota-forest/quota-manager/quota/quotamanager_test.go index 692f54691..0a22255ce 100644 --- a/pkg/quotaplugins/quota-forest/quota-manager/quota/quotamanager_test.go +++ b/pkg/quotaplugins/quota-forest/quota-manager/quota/quotamanager_test.go @@ -19,7 +19,9 @@ package quota_test import ( "strings" "testing" + "time" + "github.com/eapache/go-resiliency/retrier" "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/quotaplugins/quota-forest/quota-manager/quota" "github.com/project-codeflare/multi-cluster-app-dispatcher/pkg/quotaplugins/quota-forest/quota-manager/quota/utils" "github.com/stretchr/testify/assert" @@ -225,3 +227,169 @@ func TestNewQuotaManagerConsumerAllocationRelease(t *testing.T) { }) } } +func TestQuotaManagerQuotaUsedLongRunningConsumers(t *testing.T) { + forestName := "unit-test-1" + qmManagerUnderTest := quota.NewManager() + + err := qmManagerUnderTest.AddForest(forestName) + assert.NoError(t, err, "No error expected when adding a forest") + testTreeName, err := qmManagerUnderTest.AddTreeFromString( + `{ + "kind": "QuotaTree", + "metadata": { + "name": "test-tree" + }, + "spec": { + "resourceNames": [ + "cpu", + "memory" + ], + "nodes": { + "root": { + "parent": "nil", + "hard": "true", + "quota": { + "cpu": "10", + "memory": "256" + } + }, + "gold": { + "parent": "root", + "hard": "true", + "quota": { + "cpu": "10", + "memory": "256" + } + } + } + } + }`) + if err != nil { + assert.Fail(t, "No error expected when adding a tree to forest") + } + err = qmManagerUnderTest.AddTreeToForest(forestName, testTreeName) + assert.NoError(t, err, "No error expected when adding a tree from forest") + modeSet := qmManagerUnderTest.SetMode(quota.Normal) + assert.True(t, modeSet, "Setting the mode should not fail.") + + // Define the test table + var tests = []struct { + name string + consumer utils.JConsumer + }{ + // the table itself + {"Gold consumer 1", + utils.JConsumer{ + Kind: "Consumer", + MetaData: utils.JMetaData{ + Name: "gpld-consumer-data", + }, + Spec: utils.JConsumerSpec{ + ID: "gold-consumer-1", + Trees: []utils.JConsumerTreeSpec{ + { + TreeName: testTreeName, + GroupID: "gold", + Request: map[string]int{ + "cpu": 10, + "memory": 4, + "gpu": 0, + }, + Priority: 0, + CType: 0, + UnPreemptable: false, + }, + }, + }, + }, + }, + // the table itself + {"Gold consumer 2", + utils.JConsumer{ + Kind: "Consumer", + MetaData: utils.JMetaData{ + Name: "gpld-consumer-data", + }, + Spec: utils.JConsumerSpec{ + ID: "gold-consumer-2", + Trees: []utils.JConsumerTreeSpec{ + { + TreeName: testTreeName, + GroupID: "gold", + Request: map[string]int{ + "cpu": 10, + "memory": 4, + "gpu": 0, + }, + Priority: 0, + CType: 0, + UnPreemptable: false, + }, + }, + }, + }, + }, + } + // Execute tests in parallel + for _, tc := range tests { + tc := tc // capture range variable + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + // Get list of quota management tree IDs + qmTreeIDs := qmManagerUnderTest.GetTreeNames() + + consumerInfo, err := quota.NewConsumerInfo(tc.consumer) + assert.NoError(t, err, "No error expected when building consumer") + assert.Contains(t, qmTreeIDs, tc.consumer.Spec.Trees[0].TreeName) + + retryAllocation := retrier.New(retrier.LimitedExponentialBackoff(10, 10*time.Millisecond, 1*time.Second), + &AllocationClassifier{}) + + err = retryAllocation.Run(func() error { + added, err2 := qmManagerUnderTest.AddConsumer(consumerInfo) + assert.NoError(t, err2, "No error expected when adding consumer") + assert.True(t, added, "Consumer is expected to be added") + + response, err2 := qmManagerUnderTest.AllocateForest(forestName, consumerInfo.GetID()) + if err2 == nil { + assert.Equal(t, 0, len(strings.TrimSpace(response.GetMessage())), "A empty response is expected") + assert.True(t, response.IsAllocated(), "The allocation should succeed") + } else { + removed, err3 := qmManagerUnderTest.RemoveConsumer(consumerInfo.GetID()) + assert.NoError(t, err3, "No Error expected when removing consumer") + assert.True(t, removed, "Removal of consumer should succeed") + } + return err2 + + }) + if err != nil { + assert.Failf(t, "Allocation of quota should have succeed: '%s'", err.Error()) + } + + //simulate a long running consumer that has quota allocated + time.Sleep(10 * time.Millisecond) + + deAllocated := qmManagerUnderTest.DeAllocateForest(forestName, consumerInfo.GetID()) + assert.True(t, deAllocated, "De-allocation expected to succeed") + + removed, err := qmManagerUnderTest.RemoveConsumer(consumerInfo.GetID()) + assert.NoError(t, err, "No Error expected when removing consumer") + assert.True(t, removed, "Removal of consumer should succeed") + + }) + } + +} + +type AllocationClassifier struct { +} + +func (c *AllocationClassifier) Classify(err error) retrier.Action { + if err == nil { + return retrier.Succeed + } + if strings.TrimSpace(err.Error()) == "Failed to allocate quota on quota designation 'test-tree'" { + return retrier.Retry + } + return retrier.Fail +}