ocdp-go/backend/internal/domain/service/quota_precheck_test.go

package service

import (
	"errors"
	"testing"

	"github.com/ocdp/cluster-service/internal/domain/entity"
	"github.com/ocdp/cluster-service/internal/domain/repository"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
)

func TestCompareWorkspaceQuotaReportsExceededRequests(t *testing.T) {
	t.Parallel()

	workspace := &entity.Workspace{
		QuotaCPU:    "2",
		QuotaMemory: "4Gi",
		QuotaGPU:    "1",
		QuotaGPUMem: "10000",
	}
	estimate := &repository.ResourceEstimate{
		Requests: repository.ResourceVector{
			CPU:         resource.MustParse("2500m"),
			Memory:      resource.MustParse("3Gi"),
			GPU:         1,
			GPUMemoryMB: 12000,
		},
	}

	result, err := CompareWorkspaceQuota(workspace, estimate)
	if !errors.Is(err, ErrQuotaExceeded) {
		t.Fatalf("expected ErrQuotaExceeded, got %v", err)
	}
	if result == nil || result.Allowed {
		t.Fatalf("expected denied result, got %#v", result)
	}
	if len(result.Exceeded) != 2 {
		t.Fatalf("expected 2 exceeded resources, got %#v", result.Exceeded)
	}
	if result.Exceeded[0].Name != "requests.cpu" {
		t.Fatalf("expected requests.cpu exceeded first, got %#v", result.Exceeded)
	}
	if result.Exceeded[1].Name != "requests.nvidia.com/gpumem" {
		t.Fatalf("expected requests.nvidia.com/gpumem exceeded second, got %#v", result.Exceeded)
	}
}

func TestCompareWorkspaceQuotaUsesLimitsAsEffectiveRequests(t *testing.T) {
	t.Parallel()

	workspace := &entity.Workspace{
		QuotaGPU:    "0",
		QuotaGPUMem: "9999",
	}
	estimate := &repository.ResourceEstimate{
		Limits: repository.ResourceVector{
			GPU:         1,
			GPUMemoryMB: 10000,
		},
	}

	result, err := CompareWorkspaceQuota(workspace, estimate)
	if !errors.Is(err, ErrQuotaExceeded) {
		t.Fatalf("expected ErrQuotaExceeded from limits-only GPU resources, got %v", err)
	}
	if result == nil || len(result.Exceeded) != 2 {
		t.Fatalf("expected gpu and gpumem to be exceeded, got %#v", result)
	}
}

func TestCompareBindingQuotaSubtractsCurrentReleaseFromUsedQuota(t *testing.T) {
	t.Parallel()

	binding := &entity.WorkspaceClusterBinding{
		QuotaCPU:    "1",
		QuotaMemory: "2Gi",
		QuotaGPU:    "1",
		QuotaGPUMem: "10000",
	}
	usage := &repository.ResourceQuotaUsage{
		Used: repository.ResourceVector{
			CPU:         resource.MustParse("1"),
			Memory:      resource.MustParse("2Gi"),
			GPU:         1,
			GPUMemoryMB: 10000,
		},
	}
	current := &repository.ResourceEstimate{
		Requests: repository.ResourceVector{
			CPU:         resource.MustParse("1"),
			Memory:      resource.MustParse("2Gi"),
			GPU:         1,
			GPUMemoryMB: 10000,
		},
	}
	targetSameSize := &repository.ResourceEstimate{
		Requests: repository.ResourceVector{
			CPU:         resource.MustParse("1"),
			Memory:      resource.MustParse("2Gi"),
			GPU:         1,
			GPUMemoryMB: 10000,
		},
	}

	result, err := CompareBindingQuota(binding, usage, targetSameSize, current)
	if err != nil {
		t.Fatalf("expected update with same resource footprint to fit quota, got %v", err)
	}
	if result.Required.Requests.GPU != 1 || result.Required.Requests.GPUMemoryMB != 10000 {
		t.Fatalf("expected required resources to subtract current release before target, got %#v", result.Required.Requests)
	}

	targetScaledUp := &repository.ResourceEstimate{
		Requests: repository.ResourceVector{
			CPU:         resource.MustParse("2"),
			Memory:      resource.MustParse("4Gi"),
			GPU:         2,
			GPUMemoryMB: 20000,
		},
	}
	result, err = CompareBindingQuota(binding, usage, targetScaledUp, current)
	if !errors.Is(err, ErrQuotaExceeded) {
		t.Fatalf("expected scale-up beyond quota to be rejected, got %v", err)
	}
	if result == nil || result.Allowed {
		t.Fatalf("expected denied quota result, got %#v", result)
	}
}

func TestCompareBindingQuotaTreatsExplicitZeroGPUAsNoGPUAllowed(t *testing.T) {
	t.Parallel()

	binding := &entity.WorkspaceClusterBinding{
		QuotaCPU:    "8",
		QuotaMemory: "32Gi",
		QuotaGPU:    "0",
		QuotaGPUMem: "0",
	}
	vllmLikeEstimate := &repository.ResourceEstimate{
		Requests: repository.ResourceVector{
			CPU:         resource.MustParse("2"),
			Memory:      resource.MustParse("8Gi"),
			GPU:         1,
			GPUMemoryMB: 10000,
		},
	}

	result, err := CompareBindingQuota(binding, &repository.ResourceQuotaUsage{}, vllmLikeEstimate, nil)
	if !errors.Is(err, ErrQuotaExceeded) {
		t.Fatalf("expected GPU request to exceed explicit zero quota, got %v", err)
	}
	exceeded := map[string]bool{}
	for _, item := range result.Exceeded {
		exceeded[item.Name] = true
	}
	for _, name := range []string{"requests.nvidia.com/gpu", "requests.nvidia.com/gpumem"} {
		if !exceeded[name] {
			t.Fatalf("expected %s to be exceeded, got %#v", name, result.Exceeded)
		}
	}
}

func TestBindingQuotaHardKeepsGPUMemoryAsIntegerMB(t *testing.T) {
	t.Parallel()

	hard := bindingQuotaHard(&entity.WorkspaceClusterBinding{QuotaGPU: "1", QuotaGPUMem: "10000"})
	gpuMem := hard[corev1.ResourceName("requests.nvidia.com/gpumem")]
	if gpuMem.Value() != 10000 {
		t.Fatalf("expected gpumem quota to remain integer MB 10000, got %s value=%d", gpuMem.String(), gpuMem.Value())
	}
}

func TestEstimateRenderedManifestResourcesSumsPodTemplates(t *testing.T) {
	t.Parallel()

	manifest := `
apiVersion: apps/v1
kind: Deployment
metadata:
  name: gpu-worker
spec:
  replicas: 3
  template:
    spec:
      initContainers:
      - name: init
        image: busybox
        resources:
          requests:
            cpu: 100m
            memory: 128Mi
      containers:
      - name: app
        image: busybox
        resources:
          requests:
            cpu: 500m
            memory: 1Gi
            nvidia.com/gpu: "1"
            nvidia.com/gpumem: "10000"
          limits:
            cpu: "1"
            memory: 2Gi
            nvidia.com/gpu: "1"
            nvidia.com/gpumem: "12000"
---
apiVersion: v1
kind: Service
metadata:
  name: ignored
`
	estimate, err := EstimateRenderedManifestResources(manifest)
	if err != nil {
		t.Fatalf("EstimateRenderedManifestResources returned error: %v", err)
	}
	if estimate.Requests.CPU.Cmp(resource.MustParse("1800m")) != 0 {
		t.Fatalf("expected requests cpu 1800m, got %s", estimate.Requests.CPU.String())
	}
	if estimate.Requests.Memory.Cmp(resource.MustParse("3456Mi")) != 0 {
		t.Fatalf("expected requests memory 3456Mi, got %s", estimate.Requests.Memory.String())
	}
	if estimate.Requests.GPU != 3 {
		t.Fatalf("expected requests gpu 3, got %d", estimate.Requests.GPU)
	}
	if estimate.Requests.GPUMemoryMB != 30000 {
		t.Fatalf("expected requests gpumem 30000, got %d", estimate.Requests.GPUMemoryMB)
	}
	if estimate.Limits.CPU.Cmp(resource.MustParse("3")) != 0 {
		t.Fatalf("expected limits cpu 3, got %s", estimate.Limits.CPU.String())
	}
	if estimate.Limits.Memory.Cmp(resource.MustParse("6Gi")) != 0 {
		t.Fatalf("expected limits memory 6Gi, got %s", estimate.Limits.Memory.String())
	}
	if estimate.Limits.GPU != 3 {
		t.Fatalf("expected limits gpu 3, got %d", estimate.Limits.GPU)
	}
	if estimate.Limits.GPUMemoryMB != 36000 {
		t.Fatalf("expected limits gpumem 36000, got %d", estimate.Limits.GPUMemoryMB)
	}
}