kserve/pkg/controller/v1alpha2/llmisvc/managed_dra_test.go

/*
Copyright 2023 The KServe 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 llmisvc

import (
	"context"
	"errors"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	corev1 "k8s.io/api/core/v1"
	resourcev1 "k8s.io/api/resource/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/client/interceptor"
	"sigs.k8s.io/controller-runtime/pkg/controller/controllerutil"
	"sigs.k8s.io/yaml"

	"github.com/kserve/kserve/pkg/apis/serving/v1alpha2"
	"github.com/kserve/kserve/pkg/constants"
)

func newManagedDRATestLLMISVC(name string, annotations map[string]string) *v1alpha2.LLMInferenceService {
	return &v1alpha2.LLMInferenceService{
		TypeMeta: metav1.TypeMeta{
			APIVersion: "serving.kserve.io/v1alpha2",
			Kind:       "LLMInferenceService",
		},
		ObjectMeta: metav1.ObjectMeta{
			Name:        name,
			Namespace:   "default",
			UID:         "test-uid-dra",
			Annotations: annotations,
		},
	}
}

// TestYAMLBlockScalarRoundTripsToSelectors verifies that the documented YAML `|`
// block scalar correctly parses into newline-separated CEL selectors.
func TestYAMLBlockScalarRoundTripsToSelectors(t *testing.T) {
	manifest := `
apiVersion: serving.kserve.io/v1alpha2
kind: LLMInferenceService
metadata:
  name: my-llm
  namespace: default
  annotations:
    serving.kserve.io/exp-dra-device-class: gpu.example.com
    serving.kserve.io/exp-dra-device-count: "1"
    serving.kserve.io/exp-dra-cel-selector: |
      device.attributes['gpu.example.com'].model == 'LATEST-GPU-MODEL'
      device.capacity['gpu.example.com'].memory.compareTo(quantity('4Gi')) >= 0
`

	llmSvc := &v1alpha2.LLMInferenceService{}
	require.NoError(t, yaml.Unmarshal([]byte(manifest), llmSvc))

	// Verify basic annotations are parsed as strings.
	assert.Equal(t, "1", llmSvc.Annotations[constants.ManagedDRADeviceCountAnnotationKey])
	assert.Equal(t, "gpu.example.com", llmSvc.Annotations[constants.ManagedDRADeviceClassAnnotationKey])

	// Verify the block scalar preserves newlines.
	raw := llmSvc.Annotations[constants.ManagedDRACelSelectorAnnotationKey]
	assert.Contains(t, raw, "\n")

	selectors := llmSvc.ManagedDRACelSelectors()
	require.Len(t, selectors, 2)
	assert.Equal(t, "device.attributes['gpu.example.com'].model == 'LATEST-GPU-MODEL'", selectors[0])
	assert.Equal(t, "device.capacity['gpu.example.com'].memory.compareTo(quantity('4Gi')) >= 0", selectors[1])

	// Verify the assembled template applies all selectors.
	tmpl, err := expectedManagedDRATemplate(llmSvc)
	require.NoError(t, err)
	require.Len(t, tmpl.Spec.Spec.Devices.Requests, 1)
	require.Len(t, tmpl.Spec.Spec.Devices.Requests[0].Exactly.Selectors, 2)
}

func TestBuildDeviceRequests(t *testing.T) {
	t.Run("single device without CEL", func(t *testing.T) {
		reqs := buildDeviceRequests("gpu.nvidia.com", nil, 1)
		require.Len(t, reqs, 1)
		assert.Equal(t, "device", reqs[0].Name)
		assert.Equal(t, "gpu.nvidia.com", reqs[0].Exactly.DeviceClassName)
		assert.Empty(t, reqs[0].Exactly.Selectors)
	})

	t.Run("single device with single CEL", func(t *testing.T) {
		cel := "device.capacity['gpu.nvidia.com']['memory'].compareTo(quantity('40Gi')) > 0"
		reqs := buildDeviceRequests("gpu.nvidia.com", []string{cel}, 1)
		require.Len(t, reqs, 1)
		assert.Equal(t, "device", reqs[0].Name)
		require.Len(t, reqs[0].Exactly.Selectors, 1)
		assert.Equal(t, cel, reqs[0].Exactly.Selectors[0].CEL.Expression)
	})

	t.Run("single device with multiple CEL selectors", func(t *testing.T) {
		cels := []string{
			"device.attributes['gpu.nvidia.com']['type'] == 'A100'",
			"device.capacity['gpu.nvidia.com']['memory'].compareTo(quantity('40Gi')) > 0",
		}
		reqs := buildDeviceRequests("gpu.nvidia.com", cels, 1)
		require.Len(t, reqs, 1)
		require.Len(t, reqs[0].Exactly.Selectors, 2)
		assert.Equal(t, cels[0], reqs[0].Exactly.Selectors[0].CEL.Expression)
		assert.Equal(t, cels[1], reqs[0].Exactly.Selectors[1].CEL.Expression)
	})

	t.Run("multiple devices", func(t *testing.T) {
		reqs := buildDeviceRequests("gpu.nvidia.com", nil, 3)
		require.Len(t, reqs, 1)
		assert.Equal(t, "device", reqs[0].Name)
		assert.Equal(t, "gpu.nvidia.com", reqs[0].Exactly.DeviceClassName)
		assert.Equal(t, int64(3), reqs[0].Exactly.Count)
	})

	t.Run("multiple devices with multiple CEL selectors", func(t *testing.T) {
		cels := []string{
			"device.attributes['gpu.nvidia.com']['type'] == 'mig'",
			"device.attributes['gpu.nvidia.com']['profile'] == '3g.40gb'",
		}
		reqs := buildDeviceRequests("mig-3g.40gb", cels, 2)
		require.Len(t, reqs, 1)
		assert.Equal(t, "device", reqs[0].Name)
		assert.Equal(t, "mig-3g.40gb", reqs[0].Exactly.DeviceClassName)
		assert.Equal(t, int64(2), reqs[0].Exactly.Count)
		require.Len(t, reqs[0].Exactly.Selectors, 2)
		assert.Equal(t, cels[0], reqs[0].Exactly.Selectors[0].CEL.Expression)
		assert.Equal(t, cels[1], reqs[0].Exactly.Selectors[1].CEL.Expression)
	})
}

func TestExpectedManagedDRATemplate(t *testing.T) {
	t.Run("device class only", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("my-model", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
		})
		tmpl, err := expectedManagedDRATemplate(llmSvc)
		require.NoError(t, err)

		assert.Equal(t, "my-model-managed-dra", tmpl.Name)
		assert.Equal(t, "default", tmpl.Namespace)
		require.Len(t, tmpl.Spec.Spec.Devices.Requests, 1)
		assert.Equal(t, "device", tmpl.Spec.Spec.Devices.Requests[0].Name)
		assert.Equal(t, "gpu.nvidia.com", tmpl.Spec.Spec.Devices.Requests[0].Exactly.DeviceClassName)
		require.Len(t, tmpl.OwnerReferences, 1)
		assert.True(t, *tmpl.OwnerReferences[0].Controller)
	})

	t.Run("device class with whitespace is trimmed", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("my-model", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "  gpu.nvidia.com  ",
		})
		tmpl, err := expectedManagedDRATemplate(llmSvc)
		require.NoError(t, err)
		assert.Equal(t, "gpu.nvidia.com", tmpl.Spec.Spec.Devices.Requests[0].Exactly.DeviceClassName)
	})

	t.Run("invalid device count surfaces a parse error", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("my-model", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
			constants.ManagedDRADeviceCountAnnotationKey: "abc",
		})
		_, err := expectedManagedDRATemplate(llmSvc)
		require.Error(t, err)
	})
}

func TestInjectManagedDRA_NoAnnotation(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", nil)
	podSpec := &corev1.PodSpec{
		Containers: []corev1.Container{{Name: "main"}},
	}
	injectManagedDRA(llmSvc, podSpec)

	assert.Empty(t, podSpec.ResourceClaims)
	assert.Empty(t, podSpec.Containers[0].Resources.Claims)
}

func TestTargetContainerForDRA(t *testing.T) {
	annKey := constants.ManagedDRAContainerNameAnnotationKey

	tests := []struct {
		name       string
		containers []corev1.Container
		// nil means the annotation key is absent.
		annotation *string
		expected   int
	}{
		{
			name:       "no annotation, no containers -> -1",
			containers: nil,
			annotation: nil,
			expected:   -1,
		},
		{
			name:       "no annotation, single container -> first",
			containers: []corev1.Container{{Name: "main"}},
			annotation: nil,
			expected:   0,
		},
		{
			name:       "no annotation, multiple containers -> first",
			containers: []corev1.Container{{Name: "main"}, {Name: "sidecar"}, {Name: "monitor"}},
			annotation: nil,
			expected:   0,
		},
		{
			name:       "empty annotation value falls back to first container",
			containers: []corev1.Container{{Name: "main"}, {Name: "sidecar"}},
			annotation: ptr.To(""),
			expected:   0,
		},
		{
			name:       "whitespace-only annotation falls back to first container",
			containers: []corev1.Container{{Name: "main"}, {Name: "sidecar"}},
			annotation: ptr.To("   \t\n "),
			expected:   0,
		},
		{
			name:       "annotation matches first container",
			containers: []corev1.Container{{Name: "main"}, {Name: "sidecar"}},
			annotation: ptr.To("main"),
			expected:   0,
		},
		{
			name:       "annotation matches middle container",
			containers: []corev1.Container{{Name: "init"}, {Name: "worker"}, {Name: "sidecar"}},
			annotation: ptr.To("worker"),
			expected:   1,
		},
		{
			name:       "annotation matches last container",
			containers: []corev1.Container{{Name: "init"}, {Name: "worker"}, {Name: "sidecar"}},
			annotation: ptr.To("sidecar"),
			expected:   2,
		},
		{
			name:       "annotation with surrounding whitespace is trimmed before matching",
			containers: []corev1.Container{{Name: "main"}, {Name: "sidecar"}},
			annotation: ptr.To("  sidecar  "),
			expected:   1,
		},
		{
			name:       "annotation names a container that does not exist -> -1",
			containers: []corev1.Container{{Name: "main"}, {Name: "sidecar"}},
			annotation: ptr.To("ghost"),
			expected:   -1,
		},
		{
			name:       "annotation set but PodSpec has no containers -> -1",
			containers: nil,
			annotation: ptr.To("main"),
			expected:   -1,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			annotations := map[string]string{}
			if tt.annotation != nil {
				annotations[annKey] = *tt.annotation
			}
			llmSvc := newManagedDRATestLLMISVC("t", annotations)
			podSpec := &corev1.PodSpec{Containers: tt.containers}

			assert.Equal(t, tt.expected, targetContainerForDRA(llmSvc, podSpec))
		})
	}
}

func TestInjectManagedDRA_SingleContainer(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey: "gpu.example.com",
	})
	podSpec := &corev1.PodSpec{
		Containers: []corev1.Container{{Name: "main"}},
	}
	injectManagedDRA(llmSvc, podSpec)

	require.Len(t, podSpec.ResourceClaims, 1)
	assert.Equal(t, managedDRAClaimName, podSpec.ResourceClaims[0].Name)
	assert.NotNil(t, podSpec.ResourceClaims[0].ResourceClaimTemplateName)
	assert.Equal(t, "test-managed-dra", *podSpec.ResourceClaims[0].ResourceClaimTemplateName)
	assert.Nil(t, podSpec.ResourceClaims[0].ResourceClaimName)

	require.Len(t, podSpec.Containers[0].Resources.Claims, 1)
	assert.Equal(t, managedDRAClaimName, podSpec.Containers[0].Resources.Claims[0].Name)
}

func TestInjectManagedDRA_FirstContainerByDefault(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey: "gpu.example.com",
	})
	podSpec := &corev1.PodSpec{
		Containers: []corev1.Container{
			{Name: "primary"},
			{Name: "sidecar"},
			{Name: "monitor"},
		},
	}
	injectManagedDRA(llmSvc, podSpec)

	require.Len(t, podSpec.ResourceClaims, 1)

	require.Len(t, podSpec.Containers[0].Resources.Claims, 1, "first container should have the device claim")
	assert.Equal(t, managedDRAClaimName, podSpec.Containers[0].Resources.Claims[0].Name)

	for _, ctr := range podSpec.Containers[1:] {
		assert.Empty(t, ctr.Resources.Claims, "container %q must not receive the device claim", ctr.Name)
	}
}

func TestInjectManagedDRA_ExplicitContainerName(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey:   "gpu.example.com",
		constants.ManagedDRAContainerNameAnnotationKey: "worker",
	})
	podSpec := &corev1.PodSpec{
		Containers: []corev1.Container{
			{Name: "init"},
			{Name: "worker"},
			{Name: "monitor"},
		},
	}
	injectManagedDRA(llmSvc, podSpec)

	require.Len(t, podSpec.ResourceClaims, 1)
	assert.Empty(t, podSpec.Containers[0].Resources.Claims, "container %q must not receive the device claim", podSpec.Containers[0].Name)
	require.Len(t, podSpec.Containers[1].Resources.Claims, 1, "annotated container should have the device claim")
	assert.Equal(t, managedDRAClaimName, podSpec.Containers[1].Resources.Claims[0].Name)
	assert.Empty(t, podSpec.Containers[2].Resources.Claims, "container %q must not receive the device claim", podSpec.Containers[2].Name)
}

func TestInjectManagedDRA_AnnotatedContainerMissing(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey:   "gpu.example.com",
		constants.ManagedDRAContainerNameAnnotationKey: "ghost",
	})
	podSpec := &corev1.PodSpec{
		Containers: []corev1.Container{
			{Name: "sidecar"},
			{Name: "monitor"},
		},
	}
	injectManagedDRA(llmSvc, podSpec)

	require.Len(t, podSpec.ResourceClaims, 1)
	for _, ctr := range podSpec.Containers {
		assert.Empty(t, ctr.Resources.Claims, "container %q must not receive the device claim", ctr.Name)
	}
}

func TestInjectManagedDRA_NoContainers(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey: "gpu.example.com",
	})
	podSpec := &corev1.PodSpec{}
	injectManagedDRA(llmSvc, podSpec)

	require.Len(t, podSpec.ResourceClaims, 1)
	assert.Empty(t, podSpec.Containers)
}

func TestInjectManagedDRA_Idempotent(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey: "gpu.example.com",
	})
	podSpec := &corev1.PodSpec{
		Containers: []corev1.Container{{Name: "main"}, {Name: "sidecar"}},
	}
	injectManagedDRA(llmSvc, podSpec)
	injectManagedDRA(llmSvc, podSpec)

	assert.Len(t, podSpec.ResourceClaims, 1)
	assert.Len(t, podSpec.Containers[0].Resources.Claims, 1)
	assert.Empty(t, podSpec.Containers[1].Resources.Claims)
}

func TestInjectManagedDRA_PreservesExistingClaims(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey: "gpu.example.com",
	})
	podSpec := &corev1.PodSpec{
		ResourceClaims: []corev1.PodResourceClaim{
			{Name: "user-claim"},
		},
		Containers: []corev1.Container{
			{
				Name: "main",
				Resources: corev1.ResourceRequirements{
					Claims: []corev1.ResourceClaim{{Name: "user-claim"}},
				},
			},
		},
	}
	injectManagedDRA(llmSvc, podSpec)

	assert.Len(t, podSpec.ResourceClaims, 2)
	assert.Len(t, podSpec.Containers[0].Resources.Claims, 2)
}

func TestInjectManagedDRAIntoConfig(t *testing.T) {
	const deviceClass = "gpu.example.com"

	mainPod := func() *corev1.PodSpec {
		return &corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}}
	}

	tests := []struct {
		name            string
		annotated       bool
		cfg             *v1alpha2.LLMInferenceServiceConfig
		wantTemplateInj bool
		wantWorkerInj   bool
		wantPrefTmplInj bool
		wantPrefWrkrInj bool
	}{
		{
			name:      "no annotation: nothing injected even if all slots set",
			annotated: false,
			cfg: &v1alpha2.LLMInferenceServiceConfig{
				Spec: v1alpha2.LLMInferenceServiceSpec{
					WorkloadSpec: v1alpha2.WorkloadSpec{Template: mainPod(), Worker: mainPod()},
					Prefill:      &v1alpha2.WorkloadSpec{Template: mainPod(), Worker: mainPod()},
				},
			},
		},
		{
			name:      "annotated: all four slots populated -> all four injected",
			annotated: true,
			cfg: &v1alpha2.LLMInferenceServiceConfig{
				Spec: v1alpha2.LLMInferenceServiceSpec{
					WorkloadSpec: v1alpha2.WorkloadSpec{Template: mainPod(), Worker: mainPod()},
					Prefill:      &v1alpha2.WorkloadSpec{Template: mainPod(), Worker: mainPod()},
				},
			},
			wantTemplateInj: true,
			wantWorkerInj:   true,
			wantPrefTmplInj: true,
			wantPrefWrkrInj: true,
		},
		{
			name:      "annotated: only Template set (single-node main)",
			annotated: true,
			cfg: &v1alpha2.LLMInferenceServiceConfig{
				Spec: v1alpha2.LLMInferenceServiceSpec{
					WorkloadSpec: v1alpha2.WorkloadSpec{Template: mainPod()},
				},
			},
			wantTemplateInj: true,
		},
		{
			name:      "annotated: only Worker set (multi-node worker, no leader)",
			annotated: true,
			cfg: &v1alpha2.LLMInferenceServiceConfig{
				Spec: v1alpha2.LLMInferenceServiceSpec{
					WorkloadSpec: v1alpha2.WorkloadSpec{Worker: mainPod()},
				},
			},
			wantWorkerInj: true,
		},
		{
			name:      "annotated: only Prefill.Template set (single-node prefill)",
			annotated: true,
			cfg: &v1alpha2.LLMInferenceServiceConfig{
				Spec: v1alpha2.LLMInferenceServiceSpec{
					Prefill: &v1alpha2.WorkloadSpec{Template: mainPod()},
				},
			},
			wantPrefTmplInj: true,
		},
		{
			name:      "annotated: only Prefill.Worker set (multi-node prefill worker)",
			annotated: true,
			cfg: &v1alpha2.LLMInferenceServiceConfig{
				Spec: v1alpha2.LLMInferenceServiceSpec{
					Prefill: &v1alpha2.WorkloadSpec{Worker: mainPod()},
				},
			},
			wantPrefWrkrInj: true,
		},
		{
			name:      "annotated: Prefill is nil -> only decode slots considered, no nil deref",
			annotated: true,
			cfg: &v1alpha2.LLMInferenceServiceConfig{
				Spec: v1alpha2.LLMInferenceServiceSpec{
					WorkloadSpec: v1alpha2.WorkloadSpec{Template: mainPod(), Worker: mainPod()},
				},
			},
			wantTemplateInj: true,
			wantWorkerInj:   true,
		},
		{
			name:      "annotated: all slots nil -> no-op",
			annotated: true,
			cfg:       &v1alpha2.LLMInferenceServiceConfig{},
		},
	}

	hasClaim := func(t *testing.T, pod *corev1.PodSpec) bool {
		t.Helper()
		if pod == nil {
			return false
		}
		for _, c := range pod.ResourceClaims {
			if c.Name == managedDRAClaimName {
				return true
			}
		}
		return false
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var ann map[string]string
			if tt.annotated {
				ann = map[string]string{constants.ManagedDRADeviceClassAnnotationKey: deviceClass}
			}
			llmSvc := newManagedDRATestLLMISVC("test", ann)

			injectManagedDRAIntoConfig(llmSvc, tt.cfg)

			assert.Equal(t, tt.wantTemplateInj, hasClaim(t, tt.cfg.Spec.Template), "Template")
			assert.Equal(t, tt.wantWorkerInj, hasClaim(t, tt.cfg.Spec.Worker), "Worker")
			if tt.cfg.Spec.Prefill != nil {
				assert.Equal(t, tt.wantPrefTmplInj, hasClaim(t, tt.cfg.Spec.Prefill.Template), "Prefill.Template")
				assert.Equal(t, tt.wantPrefWrkrInj, hasClaim(t, tt.cfg.Spec.Prefill.Worker), "Prefill.Worker")
			}
		})
	}
}

func TestInjectManagedDRAIntoConfig_AnnotatedContainerAcrossSlots(t *testing.T) {
	mkPod := func() *corev1.PodSpec {
		return &corev1.PodSpec{Containers: []corev1.Container{
			{Name: "init"},
			{Name: "worker"},
			{Name: "sidecar"},
		}}
	}

	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey:   "gpu.example.com",
		constants.ManagedDRAContainerNameAnnotationKey: "worker",
	})
	cfg := &v1alpha2.LLMInferenceServiceConfig{
		Spec: v1alpha2.LLMInferenceServiceSpec{
			WorkloadSpec: v1alpha2.WorkloadSpec{Template: mkPod(), Worker: mkPod()},
			Prefill:      &v1alpha2.WorkloadSpec{Template: mkPod(), Worker: mkPod()},
		},
	}

	injectManagedDRAIntoConfig(llmSvc, cfg)

	for name, pod := range map[string]*corev1.PodSpec{
		"Template":         cfg.Spec.Template,
		"Worker":           cfg.Spec.Worker,
		"Prefill.Template": cfg.Spec.Prefill.Template,
		"Prefill.Worker":   cfg.Spec.Prefill.Worker,
	} {
		require.Len(t, pod.ResourceClaims, 1, "%s pod-level claim", name)
		assert.Empty(t, pod.Containers[0].Resources.Claims, "%s init must not receive the claim", name)
		require.Len(t, pod.Containers[1].Resources.Claims, 1, "%s worker should receive the claim", name)
		assert.Equal(t, managedDRAClaimName, pod.Containers[1].Resources.Claims[0].Name, "%s worker claim name", name)
		assert.Empty(t, pod.Containers[2].Resources.Claims, "%s sidecar must not receive the claim", name)
	}
}

func TestInjectManagedDRAIntoConfig_NilCfg(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey: "gpu.example.com",
	})
	assert.NotPanics(t, func() {
		injectManagedDRAIntoConfig(llmSvc, nil)
	})
}

func TestInjectManagedDRAIntoConfig_Idempotent(t *testing.T) {
	llmSvc := newManagedDRATestLLMISVC("test", map[string]string{
		constants.ManagedDRADeviceClassAnnotationKey: "gpu.example.com",
	})
	mkPod := func() *corev1.PodSpec {
		return &corev1.PodSpec{Containers: []corev1.Container{{Name: "main"}}}
	}
	cfg := &v1alpha2.LLMInferenceServiceConfig{
		Spec: v1alpha2.LLMInferenceServiceSpec{
			WorkloadSpec: v1alpha2.WorkloadSpec{Template: mkPod(), Worker: mkPod()},
			Prefill:      &v1alpha2.WorkloadSpec{Template: mkPod(), Worker: mkPod()},
		},
	}

	injectManagedDRAIntoConfig(llmSvc, cfg)
	injectManagedDRAIntoConfig(llmSvc, cfg)

	for name, pod := range map[string]*corev1.PodSpec{
		"Template":         cfg.Spec.Template,
		"Worker":           cfg.Spec.Worker,
		"Prefill.Template": cfg.Spec.Prefill.Template,
		"Prefill.Worker":   cfg.Spec.Prefill.Worker,
	} {
		assert.Len(t, pod.ResourceClaims, 1, "%s pod-level claim should not be duplicated", name)
		assert.Len(t, pod.Containers[0].Resources.Claims, 1, "%s main-container claim should not be duplicated", name)
	}
}

func TestSemanticResourceClaimTemplateIsEqual(t *testing.T) {
	makeTemplate := func(className string) *resourcev1.ResourceClaimTemplate {
		return &resourcev1.ResourceClaimTemplate{
			Spec: resourcev1.ResourceClaimTemplateSpec{
				Spec: resourcev1.ResourceClaimSpec{
					Devices: resourcev1.DeviceClaim{
						Requests: []resourcev1.DeviceRequest{
							{
								Name:    "gpu",
								Exactly: &resourcev1.ExactDeviceRequest{DeviceClassName: className},
							},
						},
					},
				},
			},
		}
	}

	assert.True(t, semanticResourceClaimTemplateIsEqual(makeTemplate("gpu.nvidia.com"), makeTemplate("gpu.nvidia.com")))
	assert.False(t, semanticResourceClaimTemplateIsEqual(makeTemplate("gpu.nvidia.com"), makeTemplate("gpu.amd.com")))
}

func TestReconcileManagedDRA(t *testing.T) {
	ctx := context.Background()
	scheme := runtime.NewScheme()
	require.NoError(t, corev1.AddToScheme(scheme))
	require.NoError(t, resourcev1.AddToScheme(scheme))
	require.NoError(t, v1alpha2.AddToScheme(scheme))

	t.Run("no managed DRA, nothing to clean up", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("test-svc", nil)
		fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		err := r.reconcileManagedDRA(ctx, llmSvc)
		require.NoError(t, err)
	})

	t.Run("invalid device count", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("test-svc", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
			constants.ManagedDRADeviceCountAnnotationKey: "invalid",
		})
		fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		err := r.reconcileManagedDRA(ctx, llmSvc)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "invalid syntax")
	})

	t.Run("template claim is created", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("test-svc", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
		})
		fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		err := r.reconcileManagedDRA(ctx, llmSvc)
		require.NoError(t, err)

		tmpl := &resourcev1.ResourceClaimTemplate{}
		err = fakeClient.Get(ctx, types.NamespacedName{Name: "test-svc-managed-dra", Namespace: "default"}, tmpl)
		require.NoError(t, err)
		assert.Equal(t, "device", tmpl.Spec.Spec.Devices.Requests[0].Name)
	})

	t.Run("template claim with multiple CEL selectors", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("test-svc-cel", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
			constants.ManagedDRACelSelectorAnnotationKey: "device.attributes['gpu.nvidia.com']['type'] == 'A100'\n" +
				"device.capacity['gpu.nvidia.com']['memory'].compareTo(quantity('40Gi')) > 0",
		})
		fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		require.NoError(t, r.reconcileManagedDRA(ctx, llmSvc))

		tmpl := &resourcev1.ResourceClaimTemplate{}
		require.NoError(t, fakeClient.Get(ctx, types.NamespacedName{Name: "test-svc-cel-managed-dra", Namespace: "default"}, tmpl))
		require.Len(t, tmpl.Spec.Spec.Devices.Requests, 1)
		require.Len(t, tmpl.Spec.Spec.Devices.Requests[0].Exactly.Selectors, 2)
	})

	t.Run("cleanup deletes orphaned template after DRA disabled", func(t *testing.T) {
		objNamespacedName := types.NamespacedName{Name: "test-svc-managed-dra", Namespace: "default"}

		llmSvcWithDRA := newManagedDRATestLLMISVC("test-svc", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
		})

		fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		require.NoError(t, r.reconcileManagedDRA(ctx, llmSvcWithDRA))
		require.NoError(t, fakeClient.Get(ctx, objNamespacedName, &resourcev1.ResourceClaimTemplate{}),
			"template should exist after first reconcile")

		llmSvcWithoutDRA := newManagedDRATestLLMISVC("test-svc", nil)

		require.NoError(t, r.reconcileManagedDRA(ctx, llmSvcWithoutDRA))

		err := fakeClient.Get(ctx, objNamespacedName, &resourcev1.ResourceClaimTemplate{})
		require.Error(t, err, "orphaned template should be cleaned up after DRA is disabled")
		assert.True(t, apierrors.IsNotFound(err), "expected NotFound, got %v", err)
	})

	t.Run("cleanup is a no-op when there is no orphaned template", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("never-had-dra", nil)
		fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		require.NoError(t, r.reconcileManagedDRA(ctx, llmSvc))
	})

	t.Run("cleanup deletes orphaned template when service is force-stopped", func(t *testing.T) {
		objNamespacedName := types.NamespacedName{Name: "test-svc-stopped-managed-dra", Namespace: "default"}

		llmSvcWithDRA := newManagedDRATestLLMISVC("test-svc-stopped", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
		})

		fakeClient := fake.NewClientBuilder().WithScheme(scheme).Build()
		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		require.NoError(t, r.reconcileManagedDRA(ctx, llmSvcWithDRA))
		require.NoError(t, fakeClient.Get(ctx, objNamespacedName, &resourcev1.ResourceClaimTemplate{}),
			"template should exist after first reconcile")

		llmSvcWithDRA.Annotations[constants.StopAnnotationKey] = "true"
		require.NoError(t, r.reconcileManagedDRA(ctx, llmSvcWithDRA))

		err := fakeClient.Get(ctx, objNamespacedName, &resourcev1.ResourceClaimTemplate{})
		require.Error(t, err, "orphaned template should be cleaned up when service is stopped")
		assert.True(t, apierrors.IsNotFound(err), "expected NotFound, got %v", err)
	})
}

func TestReconcileManagedDRA_CreateError(t *testing.T) {
	ctx := context.Background()
	scheme := runtime.NewScheme()
	require.NoError(t, corev1.AddToScheme(scheme))
	require.NoError(t, resourcev1.AddToScheme(scheme))
	require.NoError(t, v1alpha2.AddToScheme(scheme))

	t.Run("template claim error", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("test-svc", map[string]string{
			constants.ManagedDRADeviceClassAnnotationKey: "gpu.nvidia.com",
		})

		fakeClient := fake.NewClientBuilder().
			WithScheme(scheme).
			WithInterceptorFuncs(interceptor.Funcs{
				Create: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.CreateOption) error {
					return assert.AnError
				},
			}).
			Build()

		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		err := r.reconcileManagedDRA(ctx, llmSvc)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "failed to reconcile Managed DRA ResourceClaimTemplate")
	})

	t.Run("cleanup error", func(t *testing.T) {
		llmSvc := newManagedDRATestLLMISVC("test-svc", nil)

		// Create the existing object so Delete doesn't return nil early
		existingTemplate := &resourcev1.ResourceClaimTemplate{
			ObjectMeta: metav1.ObjectMeta{
				Name:      managedDRAResourceName(llmSvc),
				Namespace: llmSvc.GetNamespace(),
			},
		}
		require.NoError(t, controllerutil.SetControllerReference(llmSvc, existingTemplate, scheme))

		fakeClient := fake.NewClientBuilder().
			WithScheme(scheme).
			WithObjects(existingTemplate).
			WithInterceptorFuncs(interceptor.Funcs{
				Delete: func(ctx context.Context, client client.WithWatch, obj client.Object, opts ...client.DeleteOption) error {
					return errors.New("mock delete error")
				},
			}).
			Build()

		r := &LLMISVCReconciler{
			Client:        fakeClient,
			EventRecorder: record.NewFakeRecorder(10),
		}

		err := r.reconcileManagedDRA(ctx, llmSvc)
		require.Error(t, err)
		assert.Contains(t, err.Error(), "failed to cleanup Managed DRA ResourceClaimTemplate")
	})
}