From f0e3c32fe560ece3aece6403e6303a9d966f8c8e Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Thu, 6 Jul 2023 14:33:01 +0200
Subject: [PATCH] Move CDI annotation code to utils package

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 pkg/kubelet/cm/dra/claiminfo.go         |  3 +-
 pkg/kubelet/cm/{dra => util/cdi}/cdi.go |  6 +--
 pkg/kubelet/cm/util/cdi/cdi_test.go     | 52 +++++++++++++++++++++++++
 3 files changed, 57 insertions(+), 4 deletions(-)
 rename pkg/kubelet/cm/{dra => util/cdi}/cdi.go (98%)
 create mode 100644 pkg/kubelet/cm/util/cdi/cdi_test.go

diff --git a/pkg/kubelet/cm/dra/claiminfo.go b/pkg/kubelet/cm/dra/claiminfo.go
index d1711a0771d..ae4c12a67b6 100644
--- a/pkg/kubelet/cm/dra/claiminfo.go
+++ b/pkg/kubelet/cm/dra/claiminfo.go
@@ -23,6 +23,7 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/kubernetes/pkg/kubelet/cm/dra/state"
+	"k8s.io/kubernetes/pkg/kubelet/cm/util/cdi"
 	kubecontainer "k8s.io/kubernetes/pkg/kubelet/container"
 )
 
@@ -57,7 +58,7 @@ func (info *ClaimInfo) addCDIDevices(pluginName string, cdiDevices []string) err
 	// NOTE: Passing CDI device names as annotations is a temporary solution
 	// It will be removed after all runtimes are updated
 	// to get CDI device names from the ContainerConfig.CDIDevices field
-	annotations, err := generateCDIAnnotations(info.ClaimUID, info.DriverName, cdiDevices)
+	annotations, err := cdi.GenerateAnnotations(info.ClaimUID, info.DriverName, cdiDevices)
 	if err != nil {
 		return fmt.Errorf("failed to generate container annotations, err: %+v", err)
 	}
diff --git a/pkg/kubelet/cm/dra/cdi.go b/pkg/kubelet/cm/util/cdi/cdi.go
similarity index 98%
rename from pkg/kubelet/cm/dra/cdi.go
rename to pkg/kubelet/cm/util/cdi/cdi.go
index b6118337817..6cd9a1b8127 100644
--- a/pkg/kubelet/cm/dra/cdi.go
+++ b/pkg/kubelet/cm/util/cdi/cdi.go
@@ -23,7 +23,7 @@ limitations under the License.
 // Long term it would be good to avoid this duplication:
 // https://github.com/container-orchestrated-devices/container-device-interface/issues/97
 
-package dra
+package cdi
 
 import (
 	"errors"
@@ -39,8 +39,8 @@ const (
 	annotationPrefix = "cdi.k8s.io/"
 )
 
-// generate container annotations using CDI UpdateAnnotations API.
-func generateCDIAnnotations(
+// GenerateAnnotations generate container annotations using CDI UpdateAnnotations API.
+func GenerateAnnotations(
 	claimUID types.UID,
 	driverName string,
 	cdiDevices []string,
diff --git a/pkg/kubelet/cm/util/cdi/cdi_test.go b/pkg/kubelet/cm/util/cdi/cdi_test.go
new file mode 100644
index 00000000000..ce84daadc0e
--- /dev/null
+++ b/pkg/kubelet/cm/util/cdi/cdi_test.go
@@ -0,0 +1,52 @@
+/*
+Copyright 2023 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package cdi
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"k8s.io/kubernetes/pkg/kubelet/container"
+)
+
+func TestGenerateAnnotations(t *testing.T) {
+	testCases := []struct {
+		description         string
+		deviceIDs           []string
+		expecteError        error
+		expectedAnnotations []container.Annotation
+	}{
+		{
+			description: "no devices",
+			deviceIDs:   []string{},
+		},
+		{
+			description:         "one device",
+			deviceIDs:           []string{"vendor.com/class=device1"},
+			expectedAnnotations: []container.Annotation{{Name: "cdi.k8s.io/test-driver-name_test-claim-uid", Value: "vendor.com/class=device1"}},
+		},
+	}
+
+	as := assert.New(t)
+	for _, tc := range testCases {
+		t.Run(tc.description, func(t *testing.T) {
+			annotations, err := GenerateAnnotations("test-claim-uid", "test-driver-name", tc.deviceIDs)
+			as.ErrorIs(err, tc.expecteError)
+			as.Equal(tc.expectedAnnotations, annotations)
+		})
+	}
+}