From 3afa0f26e4fa17e89f0e540c8ce11229a53c0790 Mon Sep 17 00:00:00 2001
From: Raul Sevilla <rsevilla@redhat.com>
Date: Fri, 4 Sep 2020 00:51:39 +0200
Subject: [PATCH] Add WaitFor

Signed-off-by: Raul Sevilla <rsevilla@redhat.com>
---
 README.md            |  1 +
 pkg/burner/burner.go | 46 +++++++++++++++++++++++++++++---------------
 pkg/config/types.go  |  2 ++
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/README.md b/README.md
index a148c5617..f063dcbda 100644
--- a/README.md
+++ b/README.md
@@ -124,6 +124,7 @@ All the magic `kube-burner` does is described in the configuration file. This fi
 | cleanup              | Cleanup clean up old namespaces                                                  | Boolean | true     | true    |
 | podWait              | Wait for all pods to be running before moving forward to the next job iteration  | Boolean | true     | true    |
 | waitWhenFinished     | Wait for all pods to be running when all iterations are completed                | Boolean | true     | false   |
+| waitFor              | List containing the objects Kind wait for. Wait for all if empty                 | List    | ["Deployments", "Build"]| []      |
 | jobIterationDelay    | How many milliseconds to wait between each job iteration                         | Integer | 2000     | false   |
 | jobPause             | How many milliseconds to pause after finishing the job                           | Integer | 10000    | 0       |
 | qps                  | Limit object creation queries per second                                         | Integer | 25       | 0       |
diff --git a/pkg/burner/burner.go b/pkg/burner/burner.go
index aa584ef53..19e0e7e1e 100644
--- a/pkg/burner/burner.go
+++ b/pkg/burner/burner.go
@@ -157,21 +157,37 @@ func getExecutor(jobConfig config.Job) Executor {
 			unstructured: uns,
 			inputVars:    o.InputVars,
 		}
-		switch kind := obj.unstructured.GetKind(); kind {
-		case "Deployment":
-			obj.waitFunc = waitForDeployments
-		case "ReplicaSet":
-			obj.waitFunc = waitForRS
-		case "ReplicationController":
-			obj.waitFunc = waitForRC
-		case "DaemonSet":
-			obj.waitFunc = waitForDS
-		case "Pod":
-			obj.waitFunc = waitForPod
-		case "Build":
-			obj.waitFunc = waitForBuild
-		default:
-			log.Infof("Resource of kind %s has no wait function", kind)
+		if jobConfig.PodWait || jobConfig.WaitWhenFinished {
+			waitFor := true
+			if len(jobConfig.WaitFor) > 0 {
+				waitFor = false
+				for _, kind := range jobConfig.WaitFor {
+					if obj.unstructured.GetKind() == kind {
+						waitFor = true
+						break
+					}
+				}
+			}
+			if waitFor {
+				kind := obj.unstructured.GetKind()
+				switch kind {
+				case "Deployment":
+					obj.waitFunc = waitForDeployments
+				case "ReplicaSet":
+					obj.waitFunc = waitForRS
+				case "ReplicationController":
+					obj.waitFunc = waitForRC
+				case "DaemonSet":
+					obj.waitFunc = waitForDS
+				case "Pod":
+					obj.waitFunc = waitForPod
+				case "Build":
+					obj.waitFunc = waitForBuild
+				}
+				if obj.waitFunc != nil {
+					log.Debugf("Added wait function for %s", kind)
+				}
+			}
 		}
 		log.Infof("Job %s: %d iterations with %d %s replicas", jobConfig.Name, jobConfig.JobIterations, obj.replicas, restMapping.GroupVersionKind.Kind)
 		ex.objects = append(ex.objects, obj)
diff --git a/pkg/config/types.go b/pkg/config/types.go
index 9cd860bb0..29250b979 100644
--- a/pkg/config/types.go
+++ b/pkg/config/types.go
@@ -97,6 +97,8 @@ type Job struct {
 	PodWait bool `yaml:"podWait"`
 	// WaitWhenFinished Wait for pods to be running when all iterations are completed
 	WaitWhenFinished bool `yaml:"waitWhenFinished"`
+	// WaitFor list of objects to wait for, if not specified wait for all
+	WaitFor []string `yaml:"waitFor"`
 	// Cleanup clean up old namespaces
 	Cleanup bool `yaml:"cleanup"`
 	// whether to create namespaces or not with each iteration