fetch_ml/internal/security/resource_quota_test.go
Jeremie Fraeys 61660dc925
refactor: co-locate security, storage, telemetry, tracking, worker tests
Move unit tests from tests/unit/ to internal/ following Go conventions:

Security tests:
- tests/unit/security/* -> internal/security/* (audit, config_integrity, filetype, gpu_audit, hipaa_validation, manifest_filename, path_traversal, resource_quota, secrets)

Storage tests:
- tests/unit/storage/* -> internal/storage/* (db, experiment_metadata)

Telemetry tests:
- tests/unit/telemetry/* -> internal/telemetry/* (telemetry)

Tracking tests:
- tests/unit/reproducibility/* -> internal/tracking/* (config_hash, environment_capture)

Worker tests:
- tests/unit/worker/* -> internal/worker/* (artifacts, config, hash_bench, plugins/jupyter_task, plugins/vllm, prewarm_v1, run_manifest_execution, snapshot_stage, snapshot_store, worker)

Update import paths in test files to reflect new locations.
2026-03-12 16:37:03 -04:00

183 lines
3.9 KiB
Go

package security_test
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)
}
})
}