easyai-ai-gateway/apps/api/internal/capacitycontroller/planner_test.go

package capacitycontroller

import (
	"encoding/json"
	"strings"
	"testing"
)

func TestCalculatePlanRespectsResourceDatabaseAndStepLimits(t *testing.T) {
	input := PlanInput{
		Queued: 1000, Running: 48, InstanceSlots: 24,
		MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
		DatabaseConnections: 80, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
		SynchronousDatabasePeers: 1, ScaleUpEligible: true,
		Sites: []SiteResources{
			{
				Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
				AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 6 << 30,
				WorkerRequestMemoryBytes: 1 << 30,
				AllocatableMilliCPU:      4000, UsedMilliCPU: 2000, WorkerRequestMilliCPU: 500,
			},
			{
				Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
				AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
				WorkerRequestMemoryBytes: 1 << 30,
				AllocatableMilliCPU:      4000, UsedMilliCPU: 500, WorkerRequestMilliCPU: 500,
			},
		},
	}
	plan := CalculatePlan(input)
	if plan.RawDesired != 22 {
		t.Fatalf("raw desired=%d, want 22", plan.RawDesired)
	}
	if plan.DesiredTotal != 4 {
		t.Fatalf("desired total=%d, want step-limited 4: %+v", plan.DesiredTotal, plan)
	}
	if plan.Sites[0].Site != "hongkong" || plan.Sites[0].DesiredReplicas != 3 {
		t.Fatalf("site allocation=%+v, want hongkong=3 ningbo=1", plan.Sites)
	}
}

func TestCalculatePlanFreezesScaleUpOnHardMemoryPressure(t *testing.T) {
	plan := CalculatePlan(PlanInput{
		Queued: 200, InstanceSlots: 24,
		DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
		SynchronousDatabasePeers: 1, ScaleUpEligible: true,
		Sites: []SiteResources{{
			Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
			AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 7 << 30,
			WorkerRequestMemoryBytes: 1 << 30,
		}},
	})
	if plan.DesiredTotal != 1 || plan.FrozenReason != "node_memory_pressure" {
		t.Fatalf("plan=%+v, want frozen at one replica", plan)
	}
}

func TestCalculatePlanReportsConfiguredMinimumOutsideResourceBudget(t *testing.T) {
	plan := CalculatePlan(PlanInput{
		Queued: 200, InstanceSlots: 24,
		DatabaseConnections: 10, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 24,
		SynchronousDatabasePeers: 1, ScaleUpEligible: true,
		Sites: []SiteResources{{
			Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4,
			AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 13 << 29,
			WorkerRequestMemoryBytes: 2 << 30,
			AllocatableMilliCPU:      4000, UsedMilliCPU: 2800, WorkerRequestMilliCPU: 500,
		}},
	})
	if plan.DesiredTotal != 1 || plan.FrozenReason != "site_resource_budget" {
		t.Fatalf("plan=%+v, want the existing minimum preserved and expansion frozen", plan)
	}
}

func TestCalculatePlanWaitsForSafeScaleDownWindow(t *testing.T) {
	input := PlanInput{
		InstanceSlots: 24, DatabaseConnections: 20, DatabaseConnectionBudget: 150,
		WorkerDatabasePoolMax: 24, SynchronousDatabasePeers: 1,
		ScaleDownEligible: false,
		Sites: []SiteResources{
			{Site: "ningbo", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
			{Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 4},
		},
	}
	plan := CalculatePlan(input)
	if plan.DesiredTotal != 4 || plan.FrozenReason != "scale_down_stabilization" {
		t.Fatalf("plan=%+v, want current replicas during stabilization", plan)
	}
	input.ScaleDownEligible = true
	plan = CalculatePlan(input)
	if plan.DesiredTotal != 2 {
		t.Fatalf("desired total=%d, want min total 2", plan.DesiredTotal)
	}
}

func TestCalculatePlanAddsCapacityAcrossLabeledSiteNodes(t *testing.T) {
	plan := CalculatePlan(PlanInput{
		Queued: 500, InstanceSlots: 24,
		MemoryTargetPercent: 75, MemoryHardPercent: 85, CPUTargetPercent: 70,
		DatabaseConnections: 20, DatabaseConnectionBudget: 150, WorkerDatabasePoolMax: 20,
		SynchronousDatabasePeers: 1, ScaleUpEligible: true,
		Sites: []SiteResources{{
			Site: "hongkong", CurrentReplicas: 2, MinReplicas: 1, MaxReplicas: 8,
			WorkerRequestMemoryBytes: 1 << 30, WorkerRequestMilliCPU: 500,
			Nodes: []NodeResources{
				{
					NodeName: "hk-a", CurrentReplicas: 1,
					AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 3 << 30,
					AllocatableMilliCPU: 4000, UsedMilliCPU: 1000,
				},
				{
					NodeName: "hk-b", CurrentReplicas: 1,
					AllocatableMemoryBytes: 8 << 30, UsedMemoryBytes: 2 << 30,
					AllocatableMilliCPU: 4000, UsedMilliCPU: 500,
				},
			},
		}},
	})
	if plan.Sites[0].ResourceMax != 7 {
		t.Fatalf("resource max=%d, want database-budgeted multi-node maximum: %+v", plan.Sites[0].ResourceMax, plan)
	}
	if plan.DesiredTotal != 4 {
		t.Fatalf("desired=%d, want two-wave step from 2 to 4", plan.DesiredTotal)
	}
}

func TestCalculatePlanCapsReplicaMaximumByDeclaredDatabasePools(t *testing.T) {
	plan := CalculatePlan(PlanInput{
		Queued: 500, InstanceSlots: 24,
		DatabaseConnections: 20, DatabaseConnectionBudget: 150,
		NonWorkerConnectionBudget: 72, WorkerDatabasePoolMax: 32,
		SynchronousDatabasePeers: 1, ScaleUpEligible: true,
		Sites: []SiteResources{
			{Site: "ningbo", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
			{Site: "hongkong", CurrentReplicas: 1, MinReplicas: 1, MaxReplicas: 4},
		},
	})
	if plan.DesiredTotal != 2 {
		t.Fatalf("desired=%d, want the current 1+1 database-safe topology: %+v", plan.DesiredTotal, plan)
	}
	resourceMax := 0
	for _, site := range plan.Sites {
		resourceMax += site.ResourceMax
	}
	if resourceMax != 2 {
		t.Fatalf("resource maximum=%d, want floor((150-72)/32)=2: %+v", resourceMax, plan)
	}
}

func TestPlanJSONUsesAcceptanceContractFieldNames(t *testing.T) {
	payload, err := json.Marshal(Plan{
		RawDesired: 3,
		Sites:      []SitePlan{{Site: "hongkong", ResourceMax: 2}},
	})
	if err != nil {
		t.Fatal(err)
	}
	text := string(payload)
	for _, field := range []string{`"rawDesired":3`, `"sites"`, `"site":"hongkong"`, `"resourceMax":2`} {
		if !strings.Contains(text, field) {
			t.Fatalf("plan JSON %s does not contain %s", text, field)
		}
	}
}