fetch_ml/tests/unit/security/resource_quota_test.go

package security

import (
	"os"
	"strconv"
	"testing"
)

// TestResourceEnvVarParsing verifies that resource environment variables
// are correctly parsed and applied.
func TestResourceEnvVarParsing(t *testing.T) {
	tests := []struct {
		name          string
		envCPU        string
		envMemory     string
		envGPUCount   string
		wantCPUParsed int
		wantGPUParsed int
	}{
		{
			name:          "valid env CPU",
			envCPU:        "4",
			wantCPUParsed: 4,
		},
		{
			name:          "env CPU equals max",
			envCPU:        "8",
			wantCPUParsed: 8,
		},
		{
			name:          "invalid env CPU falls back to 0",
			envCPU:        "invalid",
			wantCPUParsed: 0,
		},
		{
			name:          "negative env CPU parsed as-is",
			envCPU:        "-1",
			wantCPUParsed: -1, // strconv.Atoi parses negative numbers correctly
		},
		{
			name:          "valid env GPU count",
			envGPUCount:   "2",
			wantGPUParsed: 2,
		},
		{
			name:          "invalid env GPU count falls back to 0",
			envGPUCount:   "invalid",
			wantGPUParsed: 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Set env vars
			if tt.envCPU != "" {
				t.Setenv("FETCH_ML_TOTAL_CPU", tt.envCPU)
			}
			if tt.envGPUCount != "" {
				t.Setenv("FETCH_ML_GPU_COUNT", tt.envGPUCount)
			}

			// Parse env values (mimicking how worker/config.go does it)
			var cpuParsed int
			if v := os.Getenv("FETCH_ML_TOTAL_CPU"); v != "" {
				if n, err := strconv.Atoi(v); err == nil {
					cpuParsed = n
				}
			}

			var gpuParsed int
			if v := os.Getenv("FETCH_ML_GPU_COUNT"); v != "" {
				if n, err := strconv.Atoi(v); err == nil {
					gpuParsed = n
				}
			}

			// Verify parsing
			if tt.envCPU != "" && cpuParsed != tt.wantCPUParsed {
				t.Errorf("CPU parsed = %d, want %d", cpuParsed, tt.wantCPUParsed)
			}
			if tt.envGPUCount != "" && gpuParsed != tt.wantGPUParsed {
				t.Errorf("GPU count parsed = %d, want %d", gpuParsed, tt.wantGPUParsed)
			}
		})
	}
}

// TestPodmanCPUParsing verifies podman_cpus parsing
func TestPodmanCPUParsing(t *testing.T) {
	tests := []struct {
		name       string
		podmanCPUs string
		wantParsed int
	}{
		{
			name:       "valid podman cpus",
			podmanCPUs: "2.5",
			wantParsed: 2,
		},
		{
			name:       "podman cpus integer",
			podmanCPUs: "4.0",
			wantParsed: 4,
		},
		{
			name:       "invalid podman cpus falls back",
			podmanCPUs: "invalid",
			wantParsed: 0, // Parse fails, falls back
		},
		{
			name:       "negative podman cpus treated as 0",
			podmanCPUs: "-1.5",
			wantParsed: 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Parse podman_cpus like config does
			var parsedCPUs int
			if f, err := strconv.ParseFloat(tt.podmanCPUs, 64); err == nil {
				if f < 0 {
					parsedCPUs = 0
				} else {
					parsedCPUs = int(f)
				}
			}

			if parsedCPUs != tt.wantParsed {
				t.Errorf("parsed CPUs = %d, want %d", parsedCPUs, tt.wantParsed)
			}
		})
	}
}

// TestResourceEnvVarOverride verifies that env vars override config values
func TestResourceEnvVarOverride(t *testing.T) {
	t.Run("env overrides config", func(t *testing.T) {
		t.Setenv("FETCH_ML_TOTAL_CPU", "8")

		// Simulate config with lower value
		configCPU := 4

		// Parse env override
		var envCPU int
		if v := os.Getenv("FETCH_ML_TOTAL_CPU"); v != "" {
			if n, err := strconv.Atoi(v); err == nil {
				envCPU = n
			}
		}

		// Env should take precedence
		finalCPU := configCPU
		if envCPU > 0 {
			finalCPU = envCPU
		}

		if finalCPU != 8 {
			t.Errorf("final CPU = %d, want 8 (env override)", finalCPU)
		}
	})

	t.Run("empty env uses config", func(t *testing.T) {
		// No env var set
		configCPU := 4

		var envCPU int
		if v := os.Getenv("FETCH_ML_TOTAL_CPU"); v != "" {
			if n, err := strconv.Atoi(v); err == nil {
				envCPU = n
			}
		}

		finalCPU := configCPU
		if envCPU > 0 {
			finalCPU = envCPU
		}

		if finalCPU != 4 {
			t.Errorf("final CPU = %d, want 4 (config value)", finalCPU)
		}
	})
}