Update new template and add extract file

docs/generated/templates.md

@@ -759,6 +759,26 @@ KubeLinter supports the following templates:
 **Supported Objects**: DeploymentLike
 
 
+## StatefulSet VolumeClaimTemplate Annotation
+
+**Key**: `statefulset-volumeclaimtemplate-annotation`
+
+**Description**: Check if StatefulSet's VolumeClaimTemplate contains a specific annotation
+
+**Supported Objects**: DeploymentLike
+
+
+**Parameters**:
+
+```yaml
+- description: Annotation specifies the required annotation to match.
+  name: annotation
+  negationAllowed: true
+  regexAllowed: true
+  required: true
+  type: string
+```
+
 ## Target Port
 
 **Key**: `target-port`

pkg/extract/sts_spec.go

@@ -0,0 +1,32 @@
+package extract
+
+import (
+	"reflect"
+
+	"golang.stackrox.io/kube-linter/pkg/k8sutil"
+	appsV1 "k8s.io/api/apps/v1"
+)
+
+func StatefulSetSpec(obj k8sutil.Object) (appsV1.StatefulSetSpec, bool) {
+	switch obj := obj.(type) {
+	case *appsV1.StatefulSet:
+		return obj.Spec, true
+	default:
+
+		kind := obj.GetObjectKind().GroupVersionKind().Kind
+		if kind != "StatefulSet" {
+			return appsV1.StatefulSetSpec{}, false
+		}
+
+		objValue := reflect.Indirect(reflect.ValueOf(obj))
+		spec := objValue.FieldByName("Spec")
+		if !spec.IsValid() {
+			return appsV1.StatefulSetSpec{}, false
+		}
+		statefulSetSpec, ok := spec.Interface().(appsV1.StatefulSetSpec)
+		if ok {
+			return statefulSetSpec, true
+		}
+		return appsV1.StatefulSetSpec{}, false
+	}
+}

pkg/lintcontext/mocks/context.go

@@ -28,3 +28,8 @@ func (l *MockLintContext) InvalidObjects() []lintcontext.InvalidObject {
 func NewMockContext() *MockLintContext {
 	return &MockLintContext{objects: make(map[string]k8sutil.Object)}
 }
+
+// AddObject adds an object to the MockLintContext
+func (l *MockLintContext) AddObject(key string, obj k8sutil.Object) {
+	l.objects[key] = obj
+}

pkg/templates/volumeclaimtemplates/template.go

@@ -6,52 +6,47 @@ import (
 	"golang.stackrox.io/kube-linter/pkg/check"
 	"golang.stackrox.io/kube-linter/pkg/config"
 	"golang.stackrox.io/kube-linter/pkg/diagnostic"
+	"golang.stackrox.io/kube-linter/pkg/extract"
 	"golang.stackrox.io/kube-linter/pkg/lintcontext"
 	"golang.stackrox.io/kube-linter/pkg/objectkinds"
 	"golang.stackrox.io/kube-linter/pkg/templates"
 	"golang.stackrox.io/kube-linter/pkg/templates/volumeclaimtemplates/internal/params"
-	v1 "k8s.io/api/apps/v1"
-	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
-	"k8s.io/apimachinery/pkg/runtime"
 )
 
 const (
-	templateKey = "statefulset-volumeclaimtemplate-annotation"
+    templateKey = "statefulset-volumeclaimtemplate-annotation"
 )
 
 func init() {
-	templates.Register(check.Template{
-		HumanName:   "StatefulSet VolumeClaimTemplate Annotation",
-		Key:         templateKey,
-		Description: "Check if StatefulSet's VolumeClaimTemplate contains a specific annotation",
-		SupportedObjectKinds: config.ObjectKindsDesc{
-			ObjectKinds: []string{objectkinds.DeploymentLike},
-		},
-		Parameters: params.ParamDescs,
-		ParseAndValidateParams: params.ParseAndValidate,
-		Instantiate: params.WrapInstantiateFunc(func(p params.Params) (check.Func, error) {
-			return func(_ lintcontext.LintContext, object lintcontext.Object) []diagnostic.Diagnostic {
-				k8sObj, ok := object.K8sObject.(*unstructured.Unstructured)
-				if !ok {
-					return nil
-				}
-				if k8sObj.GetKind() != "StatefulSet" {
-					return nil
-				}
-				var statefulSet v1.StatefulSet
-				err := runtime.DefaultUnstructuredConverter.FromUnstructured(k8sObj.UnstructuredContent(), &statefulSet)
-				if err != nil {
-					return nil
-				}
-				for _, volumeClaimTemplate := range statefulSet.Spec.VolumeClaimTemplates {
-					if annotationValue, found := volumeClaimTemplate.Annotations[p.Annotation]; found {
-						return []diagnostic.Diagnostic{{
-							Message: fmt.Sprintf("found annotation %q with value %q in VolumeClaimTemplate", p.Annotation, annotationValue),
-						}}
-					}
-				}
-				return nil
-			}, nil
-	 } ),
-	})
+    templates.Register(check.Template{
+        HumanName:   "StatefulSet VolumeClaimTemplate Annotation",
+        Key:         templateKey,
+        Description: "Check if StatefulSet's VolumeClaimTemplate contains a specific annotation",
+        SupportedObjectKinds: config.ObjectKindsDesc{
+            ObjectKinds: []string{objectkinds.DeploymentLike},
+        },
+        Parameters:             params.ParamDescs,
+        ParseAndValidateParams: params.ParseAndValidate,
+        Instantiate: params.WrapInstantiateFunc(func(p params.Params) (check.Func, error) {
+            return func(_ lintcontext.LintContext, object lintcontext.Object) []diagnostic.Diagnostic {
+                sts, ok := extract.StatefulSetSpec(object.K8sObject)
+                if !ok {
+                    fmt.Println("failed to extract StatefulSet spec")
+                    return nil
+                }
+
+                var diagnostics []diagnostic.Diagnostic
+
+                for _, vct := range sts.VolumeClaimTemplates {
+                    if vct.Annotations == nil || vct.Annotations[p.Annotation] == "" {
+                        diagnostics = append(diagnostics, diagnostic.Diagnostic{
+                            Message: fmt.Sprintf("StatefulSet's VolumeClaimTemplate is missing required annotation: %s", p.Annotation),
+                        })
+                    }
+                }
+
+                return diagnostics
+            }, nil
+        }),
+    })
 }

pkg/templates/volumeclaimtemplates/template_test.go

@@ -0,0 +1,66 @@
+package volumeclaimtemplates
+
+import (
+	"testing"
+
+	"golang.stackrox.io/kube-linter/pkg/lintcontext/mocks"
+	"golang.stackrox.io/kube-linter/pkg/templates"
+	"golang.stackrox.io/kube-linter/pkg/templates/volumeclaimtemplates/internal/params"
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func TestStatefulSetVolumeClaimTemplateAnnotation(t *testing.T) {
+	tests := []struct {
+		name        string
+		annotation  string  // Adjusted to match the parameter name in Params
+		wantDiags   int
+	}{
+		{"WithAnnotation", "value", 0},
+		{"WithoutAnnotation", "", 1},  // Empty string or any value for negative case
+	}
+
+
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			sts := &appsv1.StatefulSet{
+				ObjectMeta: metav1.ObjectMeta{Name: "statefulset"},
+				Spec: appsv1.StatefulSetSpec{
+					VolumeClaimTemplates: []corev1.PersistentVolumeClaim{
+						{ObjectMeta: metav1.ObjectMeta{Annotations: tt.annotations}},
+					},
+				},
+			}
+
+			// Creating mock lint context
+			mockLintCtx := mocks.NewMockContext()
+			mockLintCtx.AddObject("statefulset", sts)
+
+			// Fetching template
+			template, found := templates.Get("statefulset-volumeclaimtemplate-annotation")
+			if !found {
+				t.Fatalf("failed to get template")
+			}
+
+			// Parsing and validating parameters
+			params, err := params.ParseAndValidate(map[string]interface{}{})
+			if err != nil {
+				t.Fatalf("failed to parse and validate params: %v", err)
+			}
+
+			// Instantiating check function
+			checkFunc, err := template.Instantiate(params)
+			if err != nil {
+				t.Fatalf("failed to instantiate check function: %v", err)
+			}
+
+			// Running the check function
+			diags := checkFunc(mockLintCtx, mockLintCtx.Objects()[0])
+			if len(diags) != tt.wantDiags {
+				t.Errorf("got %d diagnostics, want %d", len(diags), tt.wantDiags)
+			}
+		})
+	}
+}