package k8s

import (
	"testing"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
)

func TestContainerAllocationCountsVendorGPUMemoryKey(t *testing.T) {
	container := corev1.Container{
		Resources: corev1.ResourceRequirements{
			Requests: corev1.ResourceList{
				corev1.ResourceName("nvidia.com/gpumem"): resource.MustParse("10000"),
			},
			Limits: corev1.ResourceList{
				corev1.ResourceName("nvidia.com/gpumem"): resource.MustParse("12000"),
			},
		},
	}

	allocation := containerAllocation(container)
	if allocation.GPUMemoryRequestsMB != 10000 {
		t.Fatalf("expected GPU memory requests 10000 MB, got %d", allocation.GPUMemoryRequestsMB)
	}
	if allocation.GPUMemoryLimitsMB != 12000 {
		t.Fatalf("expected GPU memory limits 12000 MB, got %d", allocation.GPUMemoryLimitsMB)
	}
}