From 0406fbd8f47aa3514aa8b04978aebace1ff5d906 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Thu, 21 Mar 2024 16:41:28 +0530
Subject: [PATCH 01/20] Identify 'kanister-job' prefix pods with jobID

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index fa95fd4025..4100abf220 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -30,6 +30,7 @@ import (
 	"github.com/kanisterio/kanister/pkg/output"
 	"github.com/kanisterio/kanister/pkg/param"
 	"github.com/kanisterio/kanister/pkg/progress"
+	"strings"
 )
 
 const (
@@ -64,6 +65,17 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
+	// Mark labels to pods with prefix `jobPrefix`. Add the jobID as reference to the origin for the pod.
+	fields := field.FromContext(ctx)
+	for _, f := range fields.Fields() {
+		if strings.HasPrefix(f.Key(), consts.LabelPrefix) {
+			jobID := f.Value().(string)
+			if options.Labels == nil {
+				options.Labels = make(map[string]string)
+			}
+			options.Labels[consts.LabelPrefix+"JobID"] = jobID
+		}
+	}
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()

From 26357f75ab0f3a427e328ceb74a95ea2eb663d69 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Thu, 21 Mar 2024 19:38:52 +0530
Subject: [PATCH 02/20] Nil pointer fix

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index 4100abf220..e2dab1c6c0 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -66,14 +66,15 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		PodOverride:  podOverride,
 	}
 	// Mark labels to pods with prefix `jobPrefix`. Add the jobID as reference to the origin for the pod.
-	fields := field.FromContext(ctx)
-	for _, f := range fields.Fields() {
-		if strings.HasPrefix(f.Key(), consts.LabelPrefix) {
-			jobID := f.Value().(string)
-			if options.Labels == nil {
-				options.Labels = make(map[string]string)
+	if fields := field.FromContext(ctx); fields != nil {
+		for _, f := range fields.Fields() {
+			if strings.HasPrefix(f.Key(), consts.LabelPrefix) {
+				jobID := f.Value().(string)
+				if options.Labels == nil {
+					options.Labels = make(map[string]string)
+				}
+				options.Labels[consts.LabelPrefix+"JobID"] = jobID
 			}
-			options.Labels[consts.LabelPrefix+"JobID"] = jobID
 		}
 	}
 

From f0b94474dbd2d2523cbae1c124bb6c3569e32ed2 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Fri, 22 Mar 2024 09:04:38 +0530
Subject: [PATCH 03/20] Addressed review comment

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index e2dab1c6c0..3fbe447e9d 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -16,6 +16,7 @@ package function
 
 import (
 	"context"
+	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -30,7 +31,6 @@ import (
 	"github.com/kanisterio/kanister/pkg/output"
 	"github.com/kanisterio/kanister/pkg/param"
 	"github.com/kanisterio/kanister/pkg/progress"
-	"strings"
 )
 
 const (
@@ -65,7 +65,7 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
-	// Mark labels to pods with prefix `jobPrefix`. Add the jobID as reference to the origin for the pod.
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
 	if fields := field.FromContext(ctx); fields != nil {
 		for _, f := range fields.Fields() {
 			if strings.HasPrefix(f.Key(), consts.LabelPrefix) {

From a616e27ace3391ff5279df26e3df082bfa438572 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Fri, 22 Mar 2024 20:58:14 +0530
Subject: [PATCH 04/20] Moved debug label addition logic to common place

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go | 14 ++------------
 pkg/kube/utils.go         | 19 +++++++++++++++++++
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index 3fbe447e9d..da9064084f 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -16,7 +16,6 @@ package function
 
 import (
 	"context"
-	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -65,18 +64,9 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
+
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	if fields := field.FromContext(ctx); fields != nil {
-		for _, f := range fields.Fields() {
-			if strings.HasPrefix(f.Key(), consts.LabelPrefix) {
-				jobID := f.Value().(string)
-				if options.Labels == nil {
-					options.Labels = make(map[string]string)
-				}
-				options.Labels[consts.LabelPrefix+"JobID"] = jobID
-			}
-		}
-	}
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index 98112cbcab..21e7f7b8be 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -17,7 +17,10 @@ package kube
 import (
 	"context"
 	"fmt"
+	"strings"
 
+	"github.com/kanisterio/kanister/pkg/consts"
+	"github.com/kanisterio/kanister/pkg/field"
 	osversioned "github.com/openshift/client-go/apps/clientset/versioned"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
@@ -171,3 +174,19 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 
 	return false
 }
+
+// AddDebugLabelsToPodOptions adds additional label selector to `PodOptions`,
+// provided the context has a key starting with `keyPrefix`.
+func AddDebugLabelsToPodOptions(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
+	if fields := field.FromContext(ctx); fields != nil {
+		for _, f := range fields.Fields() {
+			if strings.HasPrefix(f.Key(), keyPrefix) {
+				value := f.Value().(string)
+				if options.Labels == nil {
+					options.Labels = make(map[string]string)
+				}
+				options.Labels[consts.LabelPrefix+keySuffix] = value
+			}
+		}
+	}
+}

From 5860d1de75e235a0e749cb41993a8642004d6f46 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Fri, 22 Mar 2024 21:09:17 +0530
Subject: [PATCH 05/20] Adjusted imports

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/kube/utils.go | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index 21e7f7b8be..affbb8f3a7 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -19,11 +19,12 @@ import (
 	"fmt"
 	"strings"
 
-	"github.com/kanisterio/kanister/pkg/consts"
-	"github.com/kanisterio/kanister/pkg/field"
 	osversioned "github.com/openshift/client-go/apps/clientset/versioned"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
+
+	"github.com/kanisterio/kanister/pkg/consts"
+	"github.com/kanisterio/kanister/pkg/field"
 )
 
 const (

From f50929f1ef07f768c126c39ed20e6634cefef545 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Mon, 25 Mar 2024 21:19:50 +0530
Subject: [PATCH 06/20] Added debug label to additional ephemral pods

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/backup_data_stats.go               | 3 +++
 pkg/function/checkRepository.go                 | 3 +++
 pkg/function/copy_volume_data.go                | 3 +++
 pkg/function/delete_data.go                     | 3 +++
 pkg/function/delete_data_using_kopia_server.go  | 4 ++++
 pkg/function/kube_task.go                       | 1 -
 pkg/function/prepare_data.go                    | 7 +++++--
 pkg/function/restore_data.go                    | 4 ++++
 pkg/function/restore_data_using_kopia_server.go | 3 +++
 pkg/kube/utils.go                               | 6 +++---
 10 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/pkg/function/backup_data_stats.go b/pkg/function/backup_data_stats.go
index 60e5a2a14c..cb0e1eee65 100644
--- a/pkg/function/backup_data_stats.go
+++ b/pkg/function/backup_data_stats.go
@@ -75,6 +75,9 @@ func backupDataStats(ctx context.Context, cli kubernetes.Interface, tp param.Tem
 		Command:      []string{"sh", "-c", "tail -f /dev/null"},
 		PodOverride:  podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := backupDataStatsPodFunc(tp, encryptionKey, backupArtifactPrefix, backupID, mode)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/checkRepository.go b/pkg/function/checkRepository.go
index 029b4573e6..c097793575 100644
--- a/pkg/function/checkRepository.go
+++ b/pkg/function/checkRepository.go
@@ -58,6 +58,9 @@ func CheckRepository(ctx context.Context, cli kubernetes.Interface, tp param.Tem
 		Command:      []string{"sh", "-c", "tail -f /dev/null"},
 		PodOverride:  podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := CheckRepositoryPodFunc(cli, tp, encryptionKey, targetPaths, insecureTLS)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/copy_volume_data.go b/pkg/function/copy_volume_data.go
index 9e7dd631b3..332e3f2bca 100644
--- a/pkg/function/copy_volume_data.go
+++ b/pkg/function/copy_volume_data.go
@@ -99,6 +99,9 @@ func copyVolumeData(
 		}},
 		PodOverride: podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := copyVolumeDataPodFunc(cli, tp, mountPoint, targetPath, encryptionKey, insecureTLS)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/delete_data.go b/pkg/function/delete_data.go
index 61df99cd32..2c7cac53c3 100644
--- a/pkg/function/delete_data.go
+++ b/pkg/function/delete_data.go
@@ -97,6 +97,9 @@ func deleteData(
 		Command:      []string{"sh", "-c", "tail -f /dev/null"},
 		PodOverride:  podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := deleteDataPodFunc(tp, reclaimSpace, encryptionKey, insecureTLS, targetPaths, deleteTags, deleteIdentifiers)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/delete_data_using_kopia_server.go b/pkg/function/delete_data_using_kopia_server.go
index 1b5db53308..b6a3a88594 100644
--- a/pkg/function/delete_data_using_kopia_server.go
+++ b/pkg/function/delete_data_using_kopia_server.go
@@ -25,6 +25,7 @@ import (
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
+	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/format"
 	kankopia "github.com/kanisterio/kanister/pkg/kopia"
 	kopiacmd "github.com/kanisterio/kanister/pkg/kopia/command"
@@ -158,6 +159,9 @@ func deleteDataFromServer(
 		Image:        image,
 		Command:      []string{"bash", "-c", "tail -f /dev/null"},
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := deleteDataFromServerPodFunc(
 		hostname,
diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index da9064084f..f2f4b869cb 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -64,7 +64,6 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
-
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
 	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
 
diff --git a/pkg/function/prepare_data.go b/pkg/function/prepare_data.go
index 240e68de91..4164af5fc5 100644
--- a/pkg/function/prepare_data.go
+++ b/pkg/function/prepare_data.go
@@ -20,14 +20,14 @@ import (
 	"io"
 	"time"
 
-	"github.com/kanisterio/kanister/pkg/consts"
-	"github.com/kanisterio/kanister/pkg/field"
 	"github.com/pkg/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
+	"github.com/kanisterio/kanister/pkg/consts"
+	"github.com/kanisterio/kanister/pkg/field"
 	"github.com/kanisterio/kanister/pkg/format"
 	"github.com/kanisterio/kanister/pkg/kube"
 	"github.com/kanisterio/kanister/pkg/param"
@@ -107,6 +107,9 @@ func prepareData(ctx context.Context, cli kubernetes.Interface, namespace, servi
 		ServiceAccountName: serviceAccount,
 		PodOverride:        podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := prepareDataPodFunc(cli)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/restore_data.go b/pkg/function/restore_data.go
index 564776ab09..77b03e9495 100644
--- a/pkg/function/restore_data.go
+++ b/pkg/function/restore_data.go
@@ -25,6 +25,7 @@ import (
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
+	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/format"
 	"github.com/kanisterio/kanister/pkg/kube"
 	"github.com/kanisterio/kanister/pkg/param"
@@ -140,6 +141,9 @@ func restoreData(ctx context.Context, cli kubernetes.Interface, tp param.Templat
 		Volumes:      validatedVols,
 		PodOverride:  podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := restoreDataPodFunc(tp, encryptionKey, backupArtifactPrefix, restorePath, backupTag, backupID, insecureTLS)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/restore_data_using_kopia_server.go b/pkg/function/restore_data_using_kopia_server.go
index 1ecec0e396..5a23920886 100644
--- a/pkg/function/restore_data_using_kopia_server.go
+++ b/pkg/function/restore_data_using_kopia_server.go
@@ -26,6 +26,7 @@ import (
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
+	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/format"
 	kankopia "github.com/kanisterio/kanister/pkg/kopia"
 	kopiacmd "github.com/kanisterio/kanister/pkg/kopia/command"
@@ -206,6 +207,8 @@ func restoreDataFromServer(
 		Volumes:      validatedVols,
 		PodOverride:  podOverride,
 	}
+	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := restoreDataFromServerPodFunc(
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index affbb8f3a7..ded358a760 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -23,7 +23,6 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
 
-	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/field"
 )
 
@@ -179,14 +178,15 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 // AddDebugLabelsToPodOptions adds additional label selector to `PodOptions`,
 // provided the context has a key starting with `keyPrefix`.
 func AddDebugLabelsToPodOptions(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
-	if fields := field.FromContext(ctx); fields != nil {
+	fields := field.FromContext(ctx)
+	if fields != nil {
 		for _, f := range fields.Fields() {
 			if strings.HasPrefix(f.Key(), keyPrefix) {
 				value := f.Value().(string)
 				if options.Labels == nil {
 					options.Labels = make(map[string]string)
 				}
-				options.Labels[consts.LabelPrefix+keySuffix] = value
+				options.Labels[keyPrefix+keySuffix] = value
 			}
 		}
 	}

From 67289a3d8f260fc2e4c316de012be56ecbb07c54 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Tue, 26 Mar 2024 16:21:20 +0530
Subject: [PATCH 07/20] Revert "Added debug label to additional ephemral pods"

This reverts commit 444786f899ae6548578c9358d56c1b808ab56c00.
---
 pkg/function/backup_data_stats.go               | 3 ---
 pkg/function/checkRepository.go                 | 3 ---
 pkg/function/copy_volume_data.go                | 3 ---
 pkg/function/delete_data.go                     | 3 ---
 pkg/function/delete_data_using_kopia_server.go  | 4 ----
 pkg/function/kube_task.go                       | 1 +
 pkg/function/prepare_data.go                    | 7 ++-----
 pkg/function/restore_data.go                    | 4 ----
 pkg/function/restore_data_using_kopia_server.go | 3 ---
 pkg/kube/utils.go                               | 6 +++---
 10 files changed, 6 insertions(+), 31 deletions(-)

diff --git a/pkg/function/backup_data_stats.go b/pkg/function/backup_data_stats.go
index cb0e1eee65..60e5a2a14c 100644
--- a/pkg/function/backup_data_stats.go
+++ b/pkg/function/backup_data_stats.go
@@ -75,9 +75,6 @@ func backupDataStats(ctx context.Context, cli kubernetes.Interface, tp param.Tem
 		Command:      []string{"sh", "-c", "tail -f /dev/null"},
 		PodOverride:  podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
-
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := backupDataStatsPodFunc(tp, encryptionKey, backupArtifactPrefix, backupID, mode)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/checkRepository.go b/pkg/function/checkRepository.go
index c097793575..029b4573e6 100644
--- a/pkg/function/checkRepository.go
+++ b/pkg/function/checkRepository.go
@@ -58,9 +58,6 @@ func CheckRepository(ctx context.Context, cli kubernetes.Interface, tp param.Tem
 		Command:      []string{"sh", "-c", "tail -f /dev/null"},
 		PodOverride:  podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
-
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := CheckRepositoryPodFunc(cli, tp, encryptionKey, targetPaths, insecureTLS)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/copy_volume_data.go b/pkg/function/copy_volume_data.go
index 332e3f2bca..9e7dd631b3 100644
--- a/pkg/function/copy_volume_data.go
+++ b/pkg/function/copy_volume_data.go
@@ -99,9 +99,6 @@ func copyVolumeData(
 		}},
 		PodOverride: podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
-
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := copyVolumeDataPodFunc(cli, tp, mountPoint, targetPath, encryptionKey, insecureTLS)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/delete_data.go b/pkg/function/delete_data.go
index 2c7cac53c3..61df99cd32 100644
--- a/pkg/function/delete_data.go
+++ b/pkg/function/delete_data.go
@@ -97,9 +97,6 @@ func deleteData(
 		Command:      []string{"sh", "-c", "tail -f /dev/null"},
 		PodOverride:  podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
-
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := deleteDataPodFunc(tp, reclaimSpace, encryptionKey, insecureTLS, targetPaths, deleteTags, deleteIdentifiers)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/delete_data_using_kopia_server.go b/pkg/function/delete_data_using_kopia_server.go
index b6a3a88594..1b5db53308 100644
--- a/pkg/function/delete_data_using_kopia_server.go
+++ b/pkg/function/delete_data_using_kopia_server.go
@@ -25,7 +25,6 @@ import (
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
-	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/format"
 	kankopia "github.com/kanisterio/kanister/pkg/kopia"
 	kopiacmd "github.com/kanisterio/kanister/pkg/kopia/command"
@@ -159,9 +158,6 @@ func deleteDataFromServer(
 		Image:        image,
 		Command:      []string{"bash", "-c", "tail -f /dev/null"},
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
-
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := deleteDataFromServerPodFunc(
 		hostname,
diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index f2f4b869cb..da9064084f 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -64,6 +64,7 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
+
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
 	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
 
diff --git a/pkg/function/prepare_data.go b/pkg/function/prepare_data.go
index 4164af5fc5..240e68de91 100644
--- a/pkg/function/prepare_data.go
+++ b/pkg/function/prepare_data.go
@@ -20,14 +20,14 @@ import (
 	"io"
 	"time"
 
+	"github.com/kanisterio/kanister/pkg/consts"
+	"github.com/kanisterio/kanister/pkg/field"
 	"github.com/pkg/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/client-go/kubernetes"
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
-	"github.com/kanisterio/kanister/pkg/consts"
-	"github.com/kanisterio/kanister/pkg/field"
 	"github.com/kanisterio/kanister/pkg/format"
 	"github.com/kanisterio/kanister/pkg/kube"
 	"github.com/kanisterio/kanister/pkg/param"
@@ -107,9 +107,6 @@ func prepareData(ctx context.Context, cli kubernetes.Interface, namespace, servi
 		ServiceAccountName: serviceAccount,
 		PodOverride:        podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
-
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := prepareDataPodFunc(cli)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/restore_data.go b/pkg/function/restore_data.go
index 77b03e9495..564776ab09 100644
--- a/pkg/function/restore_data.go
+++ b/pkg/function/restore_data.go
@@ -25,7 +25,6 @@ import (
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
-	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/format"
 	"github.com/kanisterio/kanister/pkg/kube"
 	"github.com/kanisterio/kanister/pkg/param"
@@ -141,9 +140,6 @@ func restoreData(ctx context.Context, cli kubernetes.Interface, tp param.Templat
 		Volumes:      validatedVols,
 		PodOverride:  podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
-
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := restoreDataPodFunc(tp, encryptionKey, backupArtifactPrefix, restorePath, backupTag, backupID, insecureTLS)
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/function/restore_data_using_kopia_server.go b/pkg/function/restore_data_using_kopia_server.go
index 5a23920886..1ecec0e396 100644
--- a/pkg/function/restore_data_using_kopia_server.go
+++ b/pkg/function/restore_data_using_kopia_server.go
@@ -26,7 +26,6 @@ import (
 
 	kanister "github.com/kanisterio/kanister/pkg"
 	crv1alpha1 "github.com/kanisterio/kanister/pkg/apis/cr/v1alpha1"
-	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/format"
 	kankopia "github.com/kanisterio/kanister/pkg/kopia"
 	kopiacmd "github.com/kanisterio/kanister/pkg/kopia/command"
@@ -207,8 +206,6 @@ func restoreDataFromServer(
 		Volumes:      validatedVols,
 		PodOverride:  podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := restoreDataFromServerPodFunc(
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index ded358a760..affbb8f3a7 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -23,6 +23,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
 
+	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/field"
 )
 
@@ -178,15 +179,14 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 // AddDebugLabelsToPodOptions adds additional label selector to `PodOptions`,
 // provided the context has a key starting with `keyPrefix`.
 func AddDebugLabelsToPodOptions(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
-	fields := field.FromContext(ctx)
-	if fields != nil {
+	if fields := field.FromContext(ctx); fields != nil {
 		for _, f := range fields.Fields() {
 			if strings.HasPrefix(f.Key(), keyPrefix) {
 				value := f.Value().(string)
 				if options.Labels == nil {
 					options.Labels = make(map[string]string)
 				}
-				options.Labels[keyPrefix+keySuffix] = value
+				options.Labels[consts.LabelPrefix+keySuffix] = value
 			}
 		}
 	}

From 65f92eaecaacc68f12729ef36671c0cff76a92b1 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Tue, 26 Mar 2024 16:23:30 +0530
Subject: [PATCH 08/20] Addressed review comments

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go | 1 -
 pkg/kube/utils.go         | 6 +++---
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index da9064084f..f2f4b869cb 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -64,7 +64,6 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
-
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
 	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
 
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index affbb8f3a7..ded358a760 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -23,7 +23,6 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
 
-	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/field"
 )
 
@@ -179,14 +178,15 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 // AddDebugLabelsToPodOptions adds additional label selector to `PodOptions`,
 // provided the context has a key starting with `keyPrefix`.
 func AddDebugLabelsToPodOptions(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
-	if fields := field.FromContext(ctx); fields != nil {
+	fields := field.FromContext(ctx)
+	if fields != nil {
 		for _, f := range fields.Fields() {
 			if strings.HasPrefix(f.Key(), keyPrefix) {
 				value := f.Value().(string)
 				if options.Labels == nil {
 					options.Labels = make(map[string]string)
 				}
-				options.Labels[consts.LabelPrefix+keySuffix] = value
+				options.Labels[keyPrefix+keySuffix] = value
 			}
 		}
 	}

From e4e47a75a5e03b4e59c863c1ca275000af37c4d0 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Mon, 1 Apr 2024 19:40:52 +0530
Subject: [PATCH 09/20] Addressed review comments

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go |  2 +-
 pkg/kube/utils.go         | 22 ++++++++++++----------
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index f2f4b869cb..dfe7c81375 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -65,7 +65,7 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		PodOverride:  podOverride,
 	}
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddDebugLabelsToPodOptions(ctx, options, consts.LabelPrefix, "JobID")
+	kube.AddLabelsToPodOptionsFromContext(ctx, options, consts.LabelPrefix, "JobID")
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index ded358a760..1d3040f33d 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -175,19 +175,21 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 	return false
 }
 
-// AddDebugLabelsToPodOptions adds additional label selector to `PodOptions`,
+// AddLabelsToPodOptionsFromContext adds additional label selector to `PodOptions`,
 // provided the context has a key starting with `keyPrefix`.
-func AddDebugLabelsToPodOptions(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
+func AddLabelsToPodOptionsFromContext(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
 	fields := field.FromContext(ctx)
-	if fields != nil {
-		for _, f := range fields.Fields() {
-			if strings.HasPrefix(f.Key(), keyPrefix) {
-				value := f.Value().(string)
-				if options.Labels == nil {
-					options.Labels = make(map[string]string)
-				}
-				options.Labels[keyPrefix+keySuffix] = value
+	if fields == nil {
+		return
+	}
+	for _, f := range fields.Fields() {
+		if strings.HasPrefix(f.Key(), keyPrefix) {
+			value := f.Value().(string)
+			if options.Labels == nil {
+				options.Labels = make(map[string]string)
 			}
+			options.Labels[keyPrefix+keySuffix] = value
 		}
 	}
+
 }

From cc295a8df062fac5a4a3176bde3a716ed9a31479 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Tue, 2 Apr 2024 10:42:17 +0530
Subject: [PATCH 10/20] Added unit test for jobid debug label in ephemeral pod

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/kube/pod_runner_test.go | 60 +++++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index b072abb1ee..bc05d6311d 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -18,11 +18,15 @@ import (
 	"context"
 	"os"
 
+	"github.com/pkg/errors"
 	. "gopkg.in/check.v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/kubernetes/fake"
 	"k8s.io/client-go/testing"
+
+	"github.com/kanisterio/kanister/pkg/consts"
+	"github.com/kanisterio/kanister/pkg/field"
 )
 
 type PodRunnerTestSuite struct{}
@@ -109,9 +113,65 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 	cancel()
 }
 
+// TestPodRunnerWithJobIDDebugLabelForSuccessCase: This test adds a debug entry (kanister.io/JobID) into the context and verifies the
+// pod got created with corresponding label using the entry or not.
+func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	randomUUID := "xyz123"
+	ctx = field.Context(ctx, consts.LabelPrefix+"JobID", randomUUID)
+
+	cli := fake.NewSimpleClientset()
+	cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		return false, nil, nil
+	})
+	cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		p := &corev1.Pod{
+			Status: corev1.PodStatus{
+				Phase: corev1.PodRunning,
+			},
+		}
+		return true, p, nil
+	})
+	deleted := make(chan struct{})
+	cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		c.Log("Pod deleted due to Context Cancelled")
+		close(deleted)
+		return true, nil, nil
+	})
+	po := &PodOptions{
+		Namespace: podRunnerNS,
+		Name:      podName,
+		Command:   []string{"sh", "-c", "tail -f /dev/null"},
+	}
+	AddLabelsToPodOptionsFromContext(ctx, po, consts.LabelPrefix, "JobID")
+	pr := NewPodRunner(cli, po)
+	errorCh := make(chan error)
+	go func() {
+		_, err := pr.Run(ctx, makePodRunnerTestForKeyPresence(consts.LabelPrefix+"JobID", randomUUID, deleted))
+		errorCh <- err
+	}()
+	deleted <- struct{}{}
+	c.Assert(<-errorCh, IsNil)
+	cancel()
+}
+
 func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 		<-ch
 		return nil, nil
 	}
 }
+
+func makePodRunnerTestForKeyPresence(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+		<-ch
+		value, ok := pc.Pod().Labels[labelKey]
+		if !ok {
+			return nil, errors.New("Key not present")
+		}
+		if value != labelValue {
+			return nil, errors.New("Value mismatch")
+		}
+		return nil, nil
+	}
+}

From bdb79f75b798f0090ce8b6fecbb3aafd72741844 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Tue, 2 Apr 2024 11:35:42 +0530
Subject: [PATCH 11/20] Fix lint error and addressed missed review comment

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/kube/pod_runner_test.go |  1 +
 pkg/kube/utils.go           | 14 +++++++-------
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index bc05d6311d..57fc82dc93 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -119,6 +119,7 @@ func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C
 	ctx, cancel := context.WithCancel(context.Background())
 	randomUUID := "xyz123"
 	ctx = field.Context(ctx, consts.LabelPrefix+"JobID", randomUUID)
+	ctx = field.Context(ctx, "some-test-key", "some-test-value")
 
 	cli := fake.NewSimpleClientset()
 	cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index 1d3040f33d..c7536ed5d2 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -183,13 +183,13 @@ func AddLabelsToPodOptionsFromContext(ctx context.Context, options *PodOptions,
 		return
 	}
 	for _, f := range fields.Fields() {
-		if strings.HasPrefix(f.Key(), keyPrefix) {
-			value := f.Value().(string)
-			if options.Labels == nil {
-				options.Labels = make(map[string]string)
-			}
-			options.Labels[keyPrefix+keySuffix] = value
+		if !strings.HasPrefix(f.Key(), keyPrefix) {
+			continue
 		}
+		value := f.Value().(string)
+		if options.Labels == nil {
+			options.Labels = make(map[string]string)
+		}
+		options.Labels[keyPrefix+keySuffix] = value
 	}
-
 }

From aac86610f69a72acde9b56276df5da777df0fa88 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Tue, 2 Apr 2024 18:05:10 +0530
Subject: [PATCH 12/20] Refactored add labels to pod function

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go   |  21 +++++++-
 pkg/kube/pod_runner_test.go | 102 +++++++++++++++++++++++-------------
 pkg/kube/utils.go           |  29 +++++-----
 3 files changed, 99 insertions(+), 53 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index dfe7c81375..9df9b56cc7 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -30,6 +30,7 @@ import (
 	"github.com/kanisterio/kanister/pkg/output"
 	"github.com/kanisterio/kanister/pkg/param"
 	"github.com/kanisterio/kanister/pkg/progress"
+	"strings"
 )
 
 const (
@@ -65,7 +66,8 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		PodOverride:  podOverride,
 	}
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	kube.AddLabelsToPodOptionsFromContext(ctx, options, consts.LabelPrefix, "JobID")
+	validateFunc := validateLabelKeyIsPresentFunc(consts.LabelPrefix)
+	kube.AddLabelsToPodOptionsFromContext(ctx, options, consts.LabelPrefix+"JobID", validateFunc)
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()
@@ -95,6 +97,23 @@ func kubeTaskPodFunc() func(ctx context.Context, pc kube.PodController) (map[str
 	}
 }
 
+// validateLabelKeyIsPresentFunc: This is a helper validation function used by kubetask to validate the presence of
+// label key. Result of this is used to add target label selector to the pod
+func validateLabelKeyIsPresentFunc(keyPrefix string) func(ctx context.Context) (bool, string) {
+	return func(ctx context.Context) (bool, string) {
+		fields := field.FromContext(ctx)
+		if fields == nil {
+			return false, ""
+		}
+		for _, f := range fields.Fields() {
+			if strings.HasPrefix(f.Key(), keyPrefix) {
+				return true, f.Value().(string)
+			}
+		}
+		return false, ""
+	}
+}
+
 func (ktf *kubeTaskFunc) Exec(ctx context.Context, tp param.TemplateParams, args map[string]interface{}) (map[string]interface{}, error) {
 	// Set progress percent
 	ktf.progressPercent = progress.StartedPercent
diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index 57fc82dc93..1b85c106b7 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -116,44 +116,65 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 // TestPodRunnerWithJobIDDebugLabelForSuccessCase: This test adds a debug entry (kanister.io/JobID) into the context and verifies the
 // pod got created with corresponding label using the entry or not.
 func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C) {
-	ctx, cancel := context.WithCancel(context.Background())
 	randomUUID := "xyz123"
-	ctx = field.Context(ctx, consts.LabelPrefix+"JobID", randomUUID)
-	ctx = field.Context(ctx, "some-test-key", "some-test-value")
-
-	cli := fake.NewSimpleClientset()
-	cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-		return false, nil, nil
-	})
-	cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-		p := &corev1.Pod{
-			Status: corev1.PodStatus{
-				Phase: corev1.PodRunning,
+	for _, tc := range []struct {
+		name              string
+		validateFn        func(_ context.Context) (bool, string)
+		afterPodRunTestFn func(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error)
+	}{
+		{
+			name: "test key not present in context",
+			validateFn: func(_ context.Context) (bool, string) {
+				return false, ""
+			},
+			afterPodRunTestFn: afterPodRunTestKeyAbsentFunc,
+		},
+		{
+			name: "test key is present in context",
+			validateFn: func(_ context.Context) (bool, string) {
+				return true, randomUUID
 			},
+			afterPodRunTestFn: afterPodRunTestKeyPresentFunc,
+		},
+	} {
+		ctx, cancel := context.WithCancel(context.Background())
+		ctx = field.Context(ctx, consts.LabelPrefix+"JobID", randomUUID)
+		ctx = field.Context(ctx, "some-test-key", "some-test-value")
+
+		cli := fake.NewSimpleClientset()
+		cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+			return false, nil, nil
+		})
+		cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+			p := &corev1.Pod{
+				Status: corev1.PodStatus{
+					Phase: corev1.PodRunning,
+				},
+			}
+			return true, p, nil
+		})
+		po := &PodOptions{
+			Namespace: podRunnerNS,
+			Name:      podName,
+			Command:   []string{"sh", "-c", "tail -f /dev/null"},
 		}
-		return true, p, nil
-	})
-	deleted := make(chan struct{})
-	cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-		c.Log("Pod deleted due to Context Cancelled")
-		close(deleted)
-		return true, nil, nil
-	})
-	po := &PodOptions{
-		Namespace: podRunnerNS,
-		Name:      podName,
-		Command:   []string{"sh", "-c", "tail -f /dev/null"},
+		deleted := make(chan struct{})
+		cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+			c.Log("Pod deleted due to Context Cancelled")
+			close(deleted)
+			return true, nil, nil
+		})
+		AddLabelsToPodOptionsFromContext(ctx, po, consts.LabelPrefix+"JobID", tc.validateFn)
+		pr := NewPodRunner(cli, po)
+		errorCh := make(chan error)
+		go func() {
+			_, err := pr.Run(ctx, tc.afterPodRunTestFn(consts.LabelPrefix+"JobID", randomUUID, deleted))
+			errorCh <- err
+		}()
+		deleted <- struct{}{}
+		c.Assert(<-errorCh, IsNil)
+		cancel()
 	}
-	AddLabelsToPodOptionsFromContext(ctx, po, consts.LabelPrefix, "JobID")
-	pr := NewPodRunner(cli, po)
-	errorCh := make(chan error)
-	go func() {
-		_, err := pr.Run(ctx, makePodRunnerTestForKeyPresence(consts.LabelPrefix+"JobID", randomUUID, deleted))
-		errorCh <- err
-	}()
-	deleted <- struct{}{}
-	c.Assert(<-errorCh, IsNil)
-	cancel()
 }
 
 func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
@@ -163,7 +184,7 @@ func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodCon
 	}
 }
 
-func makePodRunnerTestForKeyPresence(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+func afterPodRunTestKeyPresentFunc(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 		<-ch
 		value, ok := pc.Pod().Labels[labelKey]
@@ -176,3 +197,14 @@ func makePodRunnerTestForKeyPresence(labelKey, labelValue string, ch chan struct
 		return nil, nil
 	}
 }
+
+func afterPodRunTestKeyAbsentFunc(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+		<-ch
+		_, present := pc.Pod().Labels[labelKey]
+		if present {
+			return nil, errors.New("Key should not be present")
+		}
+		return nil, nil
+	}
+}
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index c7536ed5d2..bd3786a9eb 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -17,13 +17,9 @@ package kube
 import (
 	"context"
 	"fmt"
-	"strings"
-
 	osversioned "github.com/openshift/client-go/apps/clientset/versioned"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
-
-	"github.com/kanisterio/kanister/pkg/field"
 )
 
 const (
@@ -176,20 +172,19 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 }
 
 // AddLabelsToPodOptionsFromContext adds additional label selector to `PodOptions`,
-// provided the context has a key starting with `keyPrefix`.
-func AddLabelsToPodOptionsFromContext(ctx context.Context, options *PodOptions, keyPrefix, keySuffix string) {
-	fields := field.FromContext(ctx)
-	if fields == nil {
+// provided the validationFunc passes successfully.
+func AddLabelsToPodOptionsFromContext(
+	ctx context.Context,
+	options *PodOptions,
+	targetKey string,
+	validateFn func(context.Context) (bool, string),
+) {
+	ok, value := validateFn(ctx)
+	if !ok {
 		return
 	}
-	for _, f := range fields.Fields() {
-		if !strings.HasPrefix(f.Key(), keyPrefix) {
-			continue
-		}
-		value := f.Value().(string)
-		if options.Labels == nil {
-			options.Labels = make(map[string]string)
-		}
-		options.Labels[keyPrefix+keySuffix] = value
+	if options.Labels == nil {
+		options.Labels = make(map[string]string)
 	}
+	options.Labels[targetKey] = value
 }

From abb5a3999493eee8db8e7061e4294b1e8e53f9aa Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Tue, 2 Apr 2024 18:07:16 +0530
Subject: [PATCH 13/20] Rearranged imports

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go | 2 +-
 pkg/kube/utils.go         | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index 9df9b56cc7..df6a751b16 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -16,6 +16,7 @@ package function
 
 import (
 	"context"
+	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -30,7 +31,6 @@ import (
 	"github.com/kanisterio/kanister/pkg/output"
 	"github.com/kanisterio/kanister/pkg/param"
 	"github.com/kanisterio/kanister/pkg/progress"
-	"strings"
 )
 
 const (
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index bd3786a9eb..e01276662d 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -17,6 +17,7 @@ package kube
 import (
 	"context"
 	"fmt"
+
 	osversioned "github.com/openshift/client-go/apps/clientset/versioned"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"

From 39dfe2e463a8b7d4aa0dbcbccdb7e8237c06aaa5 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Tue, 2 Apr 2024 18:07:16 +0530
Subject: [PATCH 14/20] Rearranged imports with minor refactor

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go   | 7 +++++--
 pkg/kube/pod_runner_test.go | 6 ++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index df6a751b16..28ba4c3ed1 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -16,6 +16,7 @@ package function
 
 import (
 	"context"
+	"path"
 	"strings"
 	"time"
 
@@ -34,7 +35,9 @@ import (
 )
 
 const (
-	jobPrefix = "kanister-job-"
+	jobPrefix   = "kanister-job-"
+	jobIDSuffix = "JobID"
+
 	// KubeTaskFuncName gives the function name
 	KubeTaskFuncName       = "KubeTask"
 	KubeTaskNamespaceArg   = "namespace"
@@ -67,7 +70,7 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 	}
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
 	validateFunc := validateLabelKeyIsPresentFunc(consts.LabelPrefix)
-	kube.AddLabelsToPodOptionsFromContext(ctx, options, consts.LabelPrefix+"JobID", validateFunc)
+	kube.AddLabelsToPodOptionsFromContext(ctx, options, path.Join(consts.LabelPrefix, jobIDSuffix), validateFunc)
 
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()
diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index 1b85c106b7..001b01873b 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -27,6 +27,7 @@ import (
 
 	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/field"
+	"path"
 )
 
 type PodRunnerTestSuite struct{}
@@ -164,11 +165,12 @@ func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C
 			close(deleted)
 			return true, nil, nil
 		})
-		AddLabelsToPodOptionsFromContext(ctx, po, consts.LabelPrefix+"JobID", tc.validateFn)
+		var targetKey = path.Join(consts.LabelPrefix, "JobID")
+		AddLabelsToPodOptionsFromContext(ctx, po, targetKey, tc.validateFn)
 		pr := NewPodRunner(cli, po)
 		errorCh := make(chan error)
 		go func() {
-			_, err := pr.Run(ctx, tc.afterPodRunTestFn(consts.LabelPrefix+"JobID", randomUUID, deleted))
+			_, err := pr.Run(ctx, tc.afterPodRunTestFn(targetKey, randomUUID, deleted))
 			errorCh <- err
 		}()
 		deleted <- struct{}{}

From 67eaff223449edbf7fb159bf2470aa5e44ad581a Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Wed, 3 Apr 2024 17:41:30 +0530
Subject: [PATCH 15/20] Simplify validateLabelKeyIsPresentFromContext and
 AddLabelToPodOptions funcs

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go   |  29 +++++------
 pkg/kube/pod_runner_test.go | 100 ++++++++++++------------------------
 pkg/kube/utils.go           |  16 ++----
 3 files changed, 52 insertions(+), 93 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index 28ba4c3ed1..03891b9405 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -69,9 +69,10 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		PodOverride:  podOverride,
 	}
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	validateFunc := validateLabelKeyIsPresentFunc(consts.LabelPrefix)
-	kube.AddLabelsToPodOptionsFromContext(ctx, options, path.Join(consts.LabelPrefix, jobIDSuffix), validateFunc)
-
+	ok, val := validateLabelKeyIsPresentFromContext(ctx, consts.LabelPrefix)
+	if ok {
+		kube.AddLabelsToPodOptions(options, path.Join(consts.LabelPrefix, jobIDSuffix), val)
+	}
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()
 	return pr.Run(ctx, podFunc)
@@ -100,21 +101,19 @@ func kubeTaskPodFunc() func(ctx context.Context, pc kube.PodController) (map[str
 	}
 }
 
-// validateLabelKeyIsPresentFunc: This is a helper validation function used by kubetask to validate the presence of
+// validateLabelKeyIsPresent: This is a helper validation function used by kubetask to validate the presence of
 // label key. Result of this is used to add target label selector to the pod
-func validateLabelKeyIsPresentFunc(keyPrefix string) func(ctx context.Context) (bool, string) {
-	return func(ctx context.Context) (bool, string) {
-		fields := field.FromContext(ctx)
-		if fields == nil {
-			return false, ""
-		}
-		for _, f := range fields.Fields() {
-			if strings.HasPrefix(f.Key(), keyPrefix) {
-				return true, f.Value().(string)
-			}
-		}
+func validateLabelKeyIsPresentFromContext(ctx context.Context, keyPrefix string) (bool, string) {
+	fields := field.FromContext(ctx)
+	if fields == nil {
 		return false, ""
 	}
+	for _, f := range fields.Fields() {
+		if strings.HasPrefix(f.Key(), keyPrefix) {
+			return true, f.Value().(string)
+		}
+	}
+	return false, ""
 }
 
 func (ktf *kubeTaskFunc) Exec(ctx context.Context, tp param.TemplateParams, args map[string]interface{}) (map[string]interface{}, error) {
diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index 001b01873b..4bebf53c2a 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -118,65 +118,42 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 // pod got created with corresponding label using the entry or not.
 func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C) {
 	randomUUID := "xyz123"
-	for _, tc := range []struct {
-		name              string
-		validateFn        func(_ context.Context) (bool, string)
-		afterPodRunTestFn func(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error)
-	}{
-		{
-			name: "test key not present in context",
-			validateFn: func(_ context.Context) (bool, string) {
-				return false, ""
-			},
-			afterPodRunTestFn: afterPodRunTestKeyAbsentFunc,
-		},
-		{
-			name: "test key is present in context",
-			validateFn: func(_ context.Context) (bool, string) {
-				return true, randomUUID
+	ctx, cancel := context.WithCancel(context.Background())
+	ctx = field.Context(ctx, path.Join(consts.LabelPrefix, "JobID"), randomUUID)
+	cli := fake.NewSimpleClientset()
+	cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		return false, nil, nil
+	})
+	cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		p := &corev1.Pod{
+			Status: corev1.PodStatus{
+				Phase: corev1.PodRunning,
 			},
-			afterPodRunTestFn: afterPodRunTestKeyPresentFunc,
-		},
-	} {
-		ctx, cancel := context.WithCancel(context.Background())
-		ctx = field.Context(ctx, consts.LabelPrefix+"JobID", randomUUID)
-		ctx = field.Context(ctx, "some-test-key", "some-test-value")
-
-		cli := fake.NewSimpleClientset()
-		cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-			return false, nil, nil
-		})
-		cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-			p := &corev1.Pod{
-				Status: corev1.PodStatus{
-					Phase: corev1.PodRunning,
-				},
-			}
-			return true, p, nil
-		})
-		po := &PodOptions{
-			Namespace: podRunnerNS,
-			Name:      podName,
-			Command:   []string{"sh", "-c", "tail -f /dev/null"},
 		}
-		deleted := make(chan struct{})
-		cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-			c.Log("Pod deleted due to Context Cancelled")
-			close(deleted)
-			return true, nil, nil
-		})
-		var targetKey = path.Join(consts.LabelPrefix, "JobID")
-		AddLabelsToPodOptionsFromContext(ctx, po, targetKey, tc.validateFn)
-		pr := NewPodRunner(cli, po)
-		errorCh := make(chan error)
-		go func() {
-			_, err := pr.Run(ctx, tc.afterPodRunTestFn(targetKey, randomUUID, deleted))
-			errorCh <- err
-		}()
-		deleted <- struct{}{}
-		c.Assert(<-errorCh, IsNil)
-		cancel()
+		return true, p, nil
+	})
+	po := &PodOptions{
+		Namespace: podRunnerNS,
+		Name:      podName,
+		Command:   []string{"sh", "-c", "tail -f /dev/null"},
 	}
+	deleted := make(chan struct{})
+	cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+		c.Log("Pod deleted due to Context Cancelled")
+		close(deleted)
+		return true, nil, nil
+	})
+	var targetKey = path.Join(consts.LabelPrefix, "JobID")
+	AddLabelsToPodOptions(po, targetKey, randomUUID)
+	pr := NewPodRunner(cli, po)
+	errorCh := make(chan error)
+	go func() {
+		_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(targetKey, randomUUID, deleted))
+		errorCh <- err
+	}()
+	deleted <- struct{}{}
+	c.Assert(<-errorCh, IsNil)
+	cancel()
 }
 
 func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
@@ -199,14 +176,3 @@ func afterPodRunTestKeyPresentFunc(labelKey, labelValue string, ch chan struct{}
 		return nil, nil
 	}
 }
-
-func afterPodRunTestKeyAbsentFunc(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
-	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
-		<-ch
-		_, present := pc.Pod().Labels[labelKey]
-		if present {
-			return nil, errors.New("Key should not be present")
-		}
-		return nil, nil
-	}
-}
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index e01276662d..daf80b75e1 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -172,20 +172,14 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 	return false
 }
 
-// AddLabelsToPodOptionsFromContext adds additional label selector to `PodOptions`,
-// provided the validationFunc passes successfully.
-func AddLabelsToPodOptionsFromContext(
-	ctx context.Context,
+// AddLabelsToPodOptions adds additional label selector to `PodOptions`,
+func AddLabelsToPodOptions(
 	options *PodOptions,
-	targetKey string,
-	validateFn func(context.Context) (bool, string),
+	targetKey,
+	targetValue string,
 ) {
-	ok, value := validateFn(ctx)
-	if !ok {
-		return
-	}
 	if options.Labels == nil {
 		options.Labels = make(map[string]string)
 	}
-	options.Labels[targetKey] = value
+	options.Labels[targetKey] = targetValue
 }

From 4324063a574f134caf493a82c3ee877ca13dc69f Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Wed, 3 Apr 2024 18:00:10 +0530
Subject: [PATCH 16/20] Rearrange utility func and imports

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go   | 18 +-----------------
 pkg/kube/pod_runner_test.go |  2 +-
 pkg/kube/utils.go           | 20 +++++++++++++++++++-
 3 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index 03891b9405..b81f2148b2 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -17,7 +17,6 @@ package function
 import (
 	"context"
 	"path"
-	"strings"
 	"time"
 
 	"github.com/pkg/errors"
@@ -69,7 +68,7 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		PodOverride:  podOverride,
 	}
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	ok, val := validateLabelKeyIsPresentFromContext(ctx, consts.LabelPrefix)
+	ok, val := kube.ValidateLabelKeyIsPresentFromContext(ctx, consts.LabelPrefix)
 	if ok {
 		kube.AddLabelsToPodOptions(options, path.Join(consts.LabelPrefix, jobIDSuffix), val)
 	}
@@ -101,21 +100,6 @@ func kubeTaskPodFunc() func(ctx context.Context, pc kube.PodController) (map[str
 	}
 }
 
-// validateLabelKeyIsPresent: This is a helper validation function used by kubetask to validate the presence of
-// label key. Result of this is used to add target label selector to the pod
-func validateLabelKeyIsPresentFromContext(ctx context.Context, keyPrefix string) (bool, string) {
-	fields := field.FromContext(ctx)
-	if fields == nil {
-		return false, ""
-	}
-	for _, f := range fields.Fields() {
-		if strings.HasPrefix(f.Key(), keyPrefix) {
-			return true, f.Value().(string)
-		}
-	}
-	return false, ""
-}
-
 func (ktf *kubeTaskFunc) Exec(ctx context.Context, tp param.TemplateParams, args map[string]interface{}) (map[string]interface{}, error) {
 	// Set progress percent
 	ktf.progressPercent = progress.StartedPercent
diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index 4bebf53c2a..5842256202 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -17,6 +17,7 @@ package kube
 import (
 	"context"
 	"os"
+	"path"
 
 	"github.com/pkg/errors"
 	. "gopkg.in/check.v1"
@@ -27,7 +28,6 @@ import (
 
 	"github.com/kanisterio/kanister/pkg/consts"
 	"github.com/kanisterio/kanister/pkg/field"
-	"path"
 )
 
 type PodRunnerTestSuite struct{}
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index daf80b75e1..c94a163885 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -17,10 +17,13 @@ package kube
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	osversioned "github.com/openshift/client-go/apps/clientset/versioned"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/client-go/kubernetes"
+
+	"github.com/kanisterio/kanister/pkg/field"
 )
 
 const (
@@ -172,7 +175,7 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 	return false
 }
 
-// AddLabelsToPodOptions adds additional label selector to `PodOptions`,
+// AddLabelsToPodOptions adds additional label selector to `PodOptions`.
 func AddLabelsToPodOptions(
 	options *PodOptions,
 	targetKey,
@@ -183,3 +186,18 @@ func AddLabelsToPodOptions(
 	}
 	options.Labels[targetKey] = targetValue
 }
+
+// ValidateLabelKeyIsPresentFromContext: This is a helper validation function used to validate the presence of a
+// label key. Result of this is used to add target label selector to the pod.
+func ValidateLabelKeyIsPresentFromContext(ctx context.Context, keyPrefix string) (bool, string) {
+	fields := field.FromContext(ctx)
+	if fields == nil {
+		return false, ""
+	}
+	for _, f := range fields.Fields() {
+		if strings.HasPrefix(f.Key(), keyPrefix) {
+			return true, f.Value().(string)
+		}
+	}
+	return false, ""
+}

From b980202b5812f9d055c38e4cdf57482484194cf5 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Wed, 3 Apr 2024 19:37:46 +0530
Subject: [PATCH 17/20] Addressed review comment

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/function/kube_task.go   |  5 +-
 pkg/kube/pod_runner_test.go | 93 +++++++++++++++++++++++--------------
 pkg/kube/utils.go           | 35 +++++++-------
 3 files changed, 76 insertions(+), 57 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index b81f2148b2..7c2ab487e4 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -68,10 +68,7 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		PodOverride:  podOverride,
 	}
 	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
-	ok, val := kube.ValidateLabelKeyIsPresentFromContext(ctx, consts.LabelPrefix)
-	if ok {
-		kube.AddLabelsToPodOptions(options, path.Join(consts.LabelPrefix, jobIDSuffix), val)
-	}
+	kube.AddLabelsToPodOptionsFromContext(ctx, options, path.Join(consts.LabelPrefix, jobIDSuffix))
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()
 	return pr.Run(ctx, podFunc)
diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index 5842256202..1d3c5df243 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -118,42 +118,62 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 // pod got created with corresponding label using the entry or not.
 func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C) {
 	randomUUID := "xyz123"
-	ctx, cancel := context.WithCancel(context.Background())
-	ctx = field.Context(ctx, path.Join(consts.LabelPrefix, "JobID"), randomUUID)
-	cli := fake.NewSimpleClientset()
-	cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-		return false, nil, nil
-	})
-	cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-		p := &corev1.Pod{
-			Status: corev1.PodStatus{
-				Phase: corev1.PodRunning,
-			},
+	for _, tc := range []struct {
+		name        string
+		targetKey   string
+		targetValue string
+		hasError    bool
+	}{
+		{
+			name:        "target key present",
+			targetKey:   path.Join(consts.LabelPrefix, "JobID"),
+			targetValue: randomUUID,
+			hasError:    false,
+		},
+		{
+			name:        "target key not present",
+			targetKey:   path.Join(consts.LabelPrefix, "NonJobID"),
+			targetValue: "some-other-value",
+			hasError:    true,
+		},
+	} {
+		ctx, cancel := context.WithCancel(context.Background())
+		ctx = field.Context(ctx, tc.targetKey, randomUUID)
+		cli := fake.NewSimpleClientset()
+		cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+			return false, nil, nil
+		})
+		cli.PrependReactor("get", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+			p := &corev1.Pod{
+				Status: corev1.PodStatus{
+					Phase: corev1.PodRunning,
+				},
+			}
+			return true, p, nil
+		})
+		po := &PodOptions{
+			Namespace: podRunnerNS,
+			Name:      podName,
+			Command:   []string{"sh", "-c", "tail -f /dev/null"},
 		}
-		return true, p, nil
-	})
-	po := &PodOptions{
-		Namespace: podRunnerNS,
-		Name:      podName,
-		Command:   []string{"sh", "-c", "tail -f /dev/null"},
+		deleted := make(chan struct{})
+		cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
+			c.Log("Pod deleted due to Context Cancelled")
+			close(deleted)
+			return true, nil, nil
+		})
+		var targetKey = path.Join(consts.LabelPrefix, "JobID")
+		AddLabelsToPodOptionsFromContext(ctx, po, targetKey)
+		pr := NewPodRunner(cli, po)
+		errorCh := make(chan error)
+		go func() {
+			_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(targetKey, randomUUID, tc.hasError, deleted))
+			errorCh <- err
+		}()
+		deleted <- struct{}{}
+		c.Assert(<-errorCh, IsNil)
+		cancel()
 	}
-	deleted := make(chan struct{})
-	cli.PrependReactor("delete", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
-		c.Log("Pod deleted due to Context Cancelled")
-		close(deleted)
-		return true, nil, nil
-	})
-	var targetKey = path.Join(consts.LabelPrefix, "JobID")
-	AddLabelsToPodOptions(po, targetKey, randomUUID)
-	pr := NewPodRunner(cli, po)
-	errorCh := make(chan error)
-	go func() {
-		_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(targetKey, randomUUID, deleted))
-		errorCh <- err
-	}()
-	deleted <- struct{}{}
-	c.Assert(<-errorCh, IsNil)
-	cancel()
 }
 
 func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
@@ -163,9 +183,12 @@ func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodCon
 	}
 }
 
-func afterPodRunTestKeyPresentFunc(labelKey, labelValue string, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+func afterPodRunTestKeyPresentFunc(labelKey, labelValue string, ignoreError bool, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 		<-ch
+		if ignoreError {
+			return nil, nil
+		}
 		value, ok := pc.Pod().Labels[labelKey]
 		if !ok {
 			return nil, errors.New("Key not present")
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index c94a163885..84738e815a 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -17,7 +17,6 @@ package kube
 import (
 	"context"
 	"fmt"
-	"strings"
 
 	osversioned "github.com/openshift/client-go/apps/clientset/versioned"
 	corev1 "k8s.io/api/core/v1"
@@ -175,29 +174,29 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 	return false
 }
 
-// AddLabelsToPodOptions adds additional label selector to `PodOptions`.
-func AddLabelsToPodOptions(
+// AddLabelsToPodOptions adds additional label selector to `PodOptions`. It extracts the value from the context
+// if targetkey is present and assigns to the podoptions.
+func AddLabelsToPodOptionsFromContext(
+	ctx context.Context,
 	options *PodOptions,
-	targetKey,
-	targetValue string,
+	targetKey string,
 ) {
-	if options.Labels == nil {
-		options.Labels = make(map[string]string)
-	}
-	options.Labels[targetKey] = targetValue
-}
-
-// ValidateLabelKeyIsPresentFromContext: This is a helper validation function used to validate the presence of a
-// label key. Result of this is used to add target label selector to the pod.
-func ValidateLabelKeyIsPresentFromContext(ctx context.Context, keyPrefix string) (bool, string) {
 	fields := field.FromContext(ctx)
 	if fields == nil {
-		return false, ""
+		return
 	}
+	var targetValue = ""
 	for _, f := range fields.Fields() {
-		if strings.HasPrefix(f.Key(), keyPrefix) {
-			return true, f.Value().(string)
+		if f.Key() == targetKey {
+			targetValue = f.Value().(string)
 		}
 	}
-	return false, ""
+	// Not found
+	if targetValue == "" {
+		return
+	}
+	if options.Labels == nil {
+		options.Labels = make(map[string]string)
+	}
+	options.Labels[targetKey] = targetValue
 }

From cdf0f1c799310ca771d84dc44177edbc2568f9c9 Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Thu, 4 Apr 2024 12:30:04 +0530
Subject: [PATCH 18/20] Refactor test cases

Signed-off-by: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
---
 pkg/kube/pod_runner_test.go | 58 ++++++++++++++++++-------------------
 pkg/kube/utils.go           | 19 +++++-------
 2 files changed, 36 insertions(+), 41 deletions(-)

diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index 1d3c5df243..c45b986d8e 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -114,31 +114,37 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 	cancel()
 }
 
-// TestPodRunnerWithJobIDDebugLabelForSuccessCase: This test adds a debug entry (kanister.io/JobID) into the context and verifies the
+// TestPodRunnerWithDebugLabelForSuccessCase: This test adds a debug entry into the context and verifies the
 // pod got created with corresponding label using the entry or not.
-func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C) {
-	randomUUID := "xyz123"
+func (s *PodRunnerTestSuite) TestPodRunnerWithDebugLabelForSuccessCase(c *C) {
+	jobIDSuffix := "JobID"
 	for _, tc := range []struct {
-		name        string
-		targetKey   string
-		targetValue string
-		hasError    bool
+		name               string
+		targetKey          string
+		contextKey         string
+		contextValue       string
+		expectedLabel      bool
+		validationErrorMsg string
 	}{
 		{
-			name:        "target key present",
-			targetKey:   path.Join(consts.LabelPrefix, "JobID"),
-			targetValue: randomUUID,
-			hasError:    false,
+			name:               "target key (kanister.io/JobID) present in pod labels",
+			targetKey:          path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextKey:         path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextValue:       "xyz123",
+			expectedLabel:      true,
+			validationErrorMsg: "Expected label to be set",
 		},
 		{
-			name:        "target key not present",
-			targetKey:   path.Join(consts.LabelPrefix, "NonJobID"),
-			targetValue: "some-other-value",
-			hasError:    true,
+			name:               "target key (kanister.io/JobID) not present in pod labels",
+			targetKey:          path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextKey:         path.Join(consts.LabelPrefix, "NonJobID"),
+			contextValue:       "some-other-value",
+			expectedLabel:      false,
+			validationErrorMsg: "Expected label to be not set",
 		},
 	} {
 		ctx, cancel := context.WithCancel(context.Background())
-		ctx = field.Context(ctx, tc.targetKey, randomUUID)
+		ctx = field.Context(ctx, tc.contextKey, tc.contextValue)
 		cli := fake.NewSimpleClientset()
 		cli.PrependReactor("create", "pods", func(action testing.Action) (handled bool, ret runtime.Object, err error) {
 			return false, nil, nil
@@ -162,12 +168,11 @@ func (s *PodRunnerTestSuite) TestPodRunnerWithJobIDDebugLabelForSuccessCase(c *C
 			close(deleted)
 			return true, nil, nil
 		})
-		var targetKey = path.Join(consts.LabelPrefix, "JobID")
-		AddLabelsToPodOptionsFromContext(ctx, po, targetKey)
+		AddLabelsToPodOptionsFromContext(ctx, po, tc.targetKey)
 		pr := NewPodRunner(cli, po)
 		errorCh := make(chan error)
 		go func() {
-			_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(targetKey, randomUUID, tc.hasError, deleted))
+			_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(tc.targetKey, tc.validationErrorMsg, tc.expectedLabel, deleted))
 			errorCh <- err
 		}()
 		deleted <- struct{}{}
@@ -183,18 +188,13 @@ func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodCon
 	}
 }
 
-func afterPodRunTestKeyPresentFunc(labelKey, labelValue string, ignoreError bool, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+func afterPodRunTestKeyPresentFunc(labelKey, validationErrorMsg string, expectedLabel bool, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 		<-ch
-		if ignoreError {
-			return nil, nil
-		}
-		value, ok := pc.Pod().Labels[labelKey]
-		if !ok {
-			return nil, errors.New("Key not present")
-		}
-		if value != labelValue {
-			return nil, errors.New("Value mismatch")
+
+		_, got := pc.Pod().Labels[labelKey]
+		if got != expectedLabel {
+			return nil, errors.New(validationErrorMsg)
 		}
 		return nil, nil
 	}
diff --git a/pkg/kube/utils.go b/pkg/kube/utils.go
index 84738e815a..854e16683c 100644
--- a/pkg/kube/utils.go
+++ b/pkg/kube/utils.go
@@ -174,8 +174,8 @@ func PVCContainsReadOnlyAccessMode(pvc *corev1.PersistentVolumeClaim) bool {
 	return false
 }
 
-// AddLabelsToPodOptions adds additional label selector to `PodOptions`. It extracts the value from the context
-// if targetkey is present and assigns to the podoptions.
+// AddLabelsToPodOptionsFromContext adds a label to `PodOptions`. It extracts the value from the context
+// if targetKey is present and assigns to the options.
 func AddLabelsToPodOptionsFromContext(
 	ctx context.Context,
 	options *PodOptions,
@@ -185,18 +185,13 @@ func AddLabelsToPodOptionsFromContext(
 	if fields == nil {
 		return
 	}
-	var targetValue = ""
+	if options.Labels == nil {
+		options.Labels = make(map[string]string)
+	}
 	for _, f := range fields.Fields() {
 		if f.Key() == targetKey {
-			targetValue = f.Value().(string)
+			options.Labels[targetKey] = f.Value().(string)
+			return
 		}
 	}
-	// Not found
-	if targetValue == "" {
-		return
-	}
-	if options.Labels == nil {
-		options.Labels = make(map[string]string)
-	}
-	options.Labels[targetKey] = targetValue
 }

From befc0e300294f7787b1dabec6e3b356c5c0cbe9c Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Thu, 4 Apr 2024 16:33:48 +0530
Subject: [PATCH 19/20] Addressed review comments w.r.t formatting

---
 pkg/function/kube_task.go   |  2 +-
 pkg/kube/pod_runner_test.go | 42 +++++++++++++++++--------------------
 2 files changed, 20 insertions(+), 24 deletions(-)

diff --git a/pkg/function/kube_task.go b/pkg/function/kube_task.go
index 7c2ab487e4..74b10c4cf8 100644
--- a/pkg/function/kube_task.go
+++ b/pkg/function/kube_task.go
@@ -67,7 +67,7 @@ func kubeTask(ctx context.Context, cli kubernetes.Interface, namespace, image st
 		Command:      command,
 		PodOverride:  podOverride,
 	}
-	// Mark labels to pods with prefix `kanister.io`. Add the jobID as reference to the origin for the pod.
+	// Mark pod with label having key `kanister.io/JobID`, the value of which is a reference to the origin of the pod.
 	kube.AddLabelsToPodOptionsFromContext(ctx, options, path.Join(consts.LabelPrefix, jobIDSuffix))
 	pr := kube.NewPodRunner(cli, options)
 	podFunc := kubeTaskPodFunc()
diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index c45b986d8e..12b9190d69 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -114,33 +114,30 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 	cancel()
 }
 
-// TestPodRunnerWithDebugLabelForSuccessCase: This test adds a debug entry into the context and verifies the
+// TestPodRunnerWithDebugLabelForSuccessCase adds a debug entry into the context and verifies the
 // pod got created with corresponding label using the entry or not.
 func (s *PodRunnerTestSuite) TestPodRunnerWithDebugLabelForSuccessCase(c *C) {
 	jobIDSuffix := "JobID"
 	for _, tc := range []struct {
-		name               string
-		targetKey          string
-		contextKey         string
-		contextValue       string
-		expectedLabel      bool
-		validationErrorMsg string
+		name          string
+		targetKey     string
+		contextKey    string
+		contextValue  string
+		expectedLabel bool
 	}{
 		{
-			name:               "target key (kanister.io/JobID) present in pod labels",
-			targetKey:          path.Join(consts.LabelPrefix, jobIDSuffix),
-			contextKey:         path.Join(consts.LabelPrefix, jobIDSuffix),
-			contextValue:       "xyz123",
-			expectedLabel:      true,
-			validationErrorMsg: "Expected label to be set",
+			name:          "target key (kanister.io/JobID) present in pod labels",
+			targetKey:     path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextKey:    path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextValue:  "xyz123",
+			expectedLabel: true,
 		},
 		{
-			name:               "target key (kanister.io/JobID) not present in pod labels",
-			targetKey:          path.Join(consts.LabelPrefix, jobIDSuffix),
-			contextKey:         path.Join(consts.LabelPrefix, "NonJobID"),
-			contextValue:       "some-other-value",
-			expectedLabel:      false,
-			validationErrorMsg: "Expected label to be not set",
+			name:          "target key (kanister.io/JobID) not present in pod labels",
+			targetKey:     path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextKey:    path.Join(consts.LabelPrefix, "NonJobID"),
+			contextValue:  "some-other-value",
+			expectedLabel: false,
 		},
 	} {
 		ctx, cancel := context.WithCancel(context.Background())
@@ -172,7 +169,7 @@ func (s *PodRunnerTestSuite) TestPodRunnerWithDebugLabelForSuccessCase(c *C) {
 		pr := NewPodRunner(cli, po)
 		errorCh := make(chan error)
 		go func() {
-			_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(tc.targetKey, tc.validationErrorMsg, tc.expectedLabel, deleted))
+			_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(tc.targetKey, tc.expectedLabel, deleted))
 			errorCh <- err
 		}()
 		deleted <- struct{}{}
@@ -188,13 +185,12 @@ func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodCon
 	}
 }
 
-func afterPodRunTestKeyPresentFunc(labelKey, validationErrorMsg string, expectedLabel bool, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+func afterPodRunTestKeyPresentFunc(labelKey string, expectedLabel bool, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 		<-ch
-
 		_, got := pc.Pod().Labels[labelKey]
 		if got != expectedLabel {
-			return nil, errors.New(validationErrorMsg)
+			return nil, errors.New("Got different label than expected")
 		}
 		return nil, nil
 	}

From d07314b2ae9cf1990fc5b55cbca16614507f27ee Mon Sep 17 00:00:00 2001
From: Abhijit Mukherjee <abhijit.mukherjee@infracloud.io>
Date: Fri, 5 Apr 2024 11:55:37 +0530
Subject: [PATCH 20/20] Addressed review comment adding additional validation

---
 pkg/kube/pod_runner_test.go | 41 ++++++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 19 deletions(-)

diff --git a/pkg/kube/pod_runner_test.go b/pkg/kube/pod_runner_test.go
index 12b9190d69..85d250de7a 100644
--- a/pkg/kube/pod_runner_test.go
+++ b/pkg/kube/pod_runner_test.go
@@ -119,25 +119,25 @@ func (s *PodRunnerTestSuite) TestPodRunnerForSuccessCase(c *C) {
 func (s *PodRunnerTestSuite) TestPodRunnerWithDebugLabelForSuccessCase(c *C) {
 	jobIDSuffix := "JobID"
 	for _, tc := range []struct {
-		name          string
-		targetKey     string
-		contextKey    string
-		contextValue  string
-		expectedLabel bool
+		name            string
+		targetKey       string
+		contextKey      string
+		contextValue    string
+		isLabelExpected bool
 	}{
 		{
-			name:          "target key (kanister.io/JobID) present in pod labels",
-			targetKey:     path.Join(consts.LabelPrefix, jobIDSuffix),
-			contextKey:    path.Join(consts.LabelPrefix, jobIDSuffix),
-			contextValue:  "xyz123",
-			expectedLabel: true,
+			name:            "target key (kanister.io/JobID) present in pod labels",
+			targetKey:       path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextKey:      path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextValue:    "xyz123",
+			isLabelExpected: true,
 		},
 		{
-			name:          "target key (kanister.io/JobID) not present in pod labels",
-			targetKey:     path.Join(consts.LabelPrefix, jobIDSuffix),
-			contextKey:    path.Join(consts.LabelPrefix, "NonJobID"),
-			contextValue:  "some-other-value",
-			expectedLabel: false,
+			name:            "target key (kanister.io/JobID) not present in pod labels",
+			targetKey:       path.Join(consts.LabelPrefix, jobIDSuffix),
+			contextKey:      path.Join(consts.LabelPrefix, "NonJobID"),
+			contextValue:    "some-other-value",
+			isLabelExpected: false,
 		},
 	} {
 		ctx, cancel := context.WithCancel(context.Background())
@@ -169,7 +169,7 @@ func (s *PodRunnerTestSuite) TestPodRunnerWithDebugLabelForSuccessCase(c *C) {
 		pr := NewPodRunner(cli, po)
 		errorCh := make(chan error)
 		go func() {
-			_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(tc.targetKey, tc.expectedLabel, deleted))
+			_, err := pr.Run(ctx, afterPodRunTestKeyPresentFunc(tc.targetKey, tc.contextValue, tc.isLabelExpected, deleted))
 			errorCh <- err
 		}()
 		deleted <- struct{}{}
@@ -185,13 +185,16 @@ func makePodRunnerTestFunc(ch chan struct{}) func(ctx context.Context, pc PodCon
 	}
 }
 
-func afterPodRunTestKeyPresentFunc(labelKey string, expectedLabel bool, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
+func afterPodRunTestKeyPresentFunc(labelKey, expectedLabelValue string, isLabelExpected bool, ch chan struct{}) func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 	return func(ctx context.Context, pc PodController) (map[string]interface{}, error) {
 		<-ch
-		_, got := pc.Pod().Labels[labelKey]
-		if got != expectedLabel {
+		labelValue, found := pc.Pod().Labels[labelKey]
+		if found != isLabelExpected {
 			return nil, errors.New("Got different label than expected")
 		}
+		if isLabelExpected && labelValue != expectedLabelValue {
+			return nil, errors.New("Found label doesn't match with expected label")
+		}
 		return nil, nil
 	}
 }