From 52bef67324021d22fa591898f9c0a7101b3aef04 Mon Sep 17 00:00:00 2001
From: Peter Bacsko <bacskop@gmail.com>
Date: Fri, 8 Nov 2024 16:05:41 -0600
Subject: [PATCH] [YUNIKORN-2966] Not all tags are created for foreign
 allocations (#936)

Closes: #936

Signed-off-by: Craig Condit <ccondit@apache.org>
---
 pkg/cache/context_test.go    |  9 ++++++---
 pkg/common/si_helper.go      |  9 ++++-----
 pkg/common/si_helper_test.go | 12 +++++++++---
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/pkg/cache/context_test.go b/pkg/cache/context_test.go
index efeeffd78..91f48e60a 100644
--- a/pkg/cache/context_test.go
+++ b/pkg/cache/context_test.go
@@ -610,8 +610,10 @@ func assertAddForeignPod(t *testing.T, podName, host string, allocRequest *si.Al
 	t.Helper()
 	assert.Equal(t, 1, len(allocRequest.Allocations))
 	tags := allocRequest.Allocations[0].AllocationTags
-	assert.Equal(t, 2, len(tags))
+	assert.Equal(t, 4, len(tags))
 	assert.Equal(t, siCommon.AllocTypeDefault, tags[siCommon.Foreign])
+	assert.Equal(t, tags["kubernetes.io/meta/namespace"], "testNamespace")
+	assert.Equal(t, tags["kubernetes.io/meta/podName"], podName)
 	assert.Equal(t, podName, allocRequest.Allocations[0].AllocationKey)
 	assert.Equal(t, host, allocRequest.Allocations[0].NodeID)
 }
@@ -2390,8 +2392,9 @@ func foreignPod(podName, memory, cpu string) *v1.Pod {
 			APIVersion: "v1",
 		},
 		ObjectMeta: apis.ObjectMeta{
-			Name: podName,
-			UID:  types.UID(podName),
+			Name:      podName,
+			UID:       types.UID(podName),
+			Namespace: "testNamespace",
 		},
 		Spec: v1.PodSpec{
 			Containers: containers,
diff --git a/pkg/common/si_helper.go b/pkg/common/si_helper.go
index 272e18a18..00be116ef 100644
--- a/pkg/common/si_helper.go
+++ b/pkg/common/si_helper.go
@@ -123,18 +123,17 @@ func CreateAllocationForForeignPod(pod *v1.Pod) *si.AllocationRequest {
 		}
 	}
 
+	tags := CreateTagsForTask(pod)
+	tags[common.Foreign] = podType
+	tags[common.CreationTime] = strconv.FormatInt(pod.CreationTimestamp.Unix(), 10)
 	allocation := si.Allocation{
-		AllocationTags: map[string]string{
-			common.Foreign: podType,
-		},
+		AllocationTags:   tags,
 		AllocationKey:    string(pod.UID),
 		ResourcePerAlloc: GetPodResource(pod),
 		Priority:         CreatePriorityForTask(pod),
 		NodeID:           pod.Spec.NodeName,
 	}
 
-	allocation.AllocationTags[common.CreationTime] = strconv.FormatInt(pod.CreationTimestamp.Unix(), 10)
-
 	return &si.AllocationRequest{
 		Allocations: []*si.Allocation{&allocation},
 		RmID:        conf.GetSchedulerConf().ClusterID,
diff --git a/pkg/common/si_helper_test.go b/pkg/common/si_helper_test.go
index 9943f55cd..84a802e76 100644
--- a/pkg/common/si_helper_test.go
+++ b/pkg/common/si_helper_test.go
@@ -439,6 +439,10 @@ func TestCreateAllocationForForeignPod(t *testing.T) {
 			CreationTimestamp: apis.Time{
 				Time: time.Unix(1, 0),
 			},
+			Labels: map[string]string{
+				"testKey": "testValue",
+			},
+			Namespace: "testNamespace",
 		},
 		Spec: v1.PodSpec{
 			Containers: containers,
@@ -459,9 +463,11 @@ func TestCreateAllocationForForeignPod(t *testing.T) {
 	assert.Equal(t, int64(500000000), res.Resources["memory"].Value)
 	assert.Equal(t, int64(1000), res.Resources["vcore"].Value)
 	assert.Equal(t, int64(1), res.Resources["pods"].Value)
-	assert.Equal(t, 2, len(alloc.AllocationTags))
+	assert.Equal(t, 5, len(alloc.AllocationTags))
 	assert.Equal(t, "1", alloc.AllocationTags[common.CreationTime])
-	assert.Equal(t, common.AllocTypeDefault, alloc.AllocationTags[common.Foreign])
+	assert.Equal(t, alloc.AllocationTags["kubernetes.io/meta/namespace"], "testNamespace")
+	assert.Equal(t, alloc.AllocationTags["kubernetes.io/meta/podName"], "test")
+	assert.Equal(t, alloc.AllocationTags["kubernetes.io/label/testKey"], "testValue")
 
 	// set priority & change pod type to static
 	prio := int32(12)
@@ -472,7 +478,7 @@ func TestCreateAllocationForForeignPod(t *testing.T) {
 		},
 	}
 	allocReq = CreateAllocationForForeignPod(pod)
-	assert.Equal(t, 2, len(alloc.AllocationTags))
+	assert.Equal(t, 5, len(alloc.AllocationTags))
 	alloc = allocReq.Allocations[0]
 	assert.Equal(t, common.AllocTypeStatic, alloc.AllocationTags[common.Foreign])
 	assert.Equal(t, int32(12), alloc.Priority)