package store

import "testing"

func TestAllocateWorkerCapacitiesHonorsInstanceLimits(t *testing.T) {
	workers := []activeWorkerCapacity{
		{InstanceID: "worker-a", CapacityLimit: 24},
		{InstanceID: "worker-b", CapacityLimit: 24},
	}
	allocations, global := allocateWorkerCapacities(workers, 310)
	if global != 48 || allocations["worker-a"] != 24 || allocations["worker-b"] != 24 {
		t.Fatalf("allocations=%v global=%d, want 24/24 and 48", allocations, global)
	}
}

func TestAllocateWorkerCapacitiesBalancesBelowLimits(t *testing.T) {
	workers := []activeWorkerCapacity{
		{InstanceID: "worker-a", CapacityLimit: 24},
		{InstanceID: "worker-b", CapacityLimit: 24},
	}
	allocations, global := allocateWorkerCapacities(workers, 5)
	if global != 5 || allocations["worker-a"] != 3 || allocations["worker-b"] != 2 {
		t.Fatalf("allocations=%v global=%d, want 3/2 and 5", allocations, global)
	}
}

func TestAllocateWorkerCapacitiesKeepsSurvivorBounded(t *testing.T) {
	workers := []activeWorkerCapacity{{InstanceID: "worker-a", CapacityLimit: 24}}
	allocations, global := allocateWorkerCapacities(workers, 310)
	if global != 24 || allocations["worker-a"] != 24 {
		t.Fatalf("allocations=%v global=%d, want survivor capped at 24", allocations, global)
	}
}

func TestAllocateWorkerCapacitiesSupportsUnequalLimits(t *testing.T) {
	workers := []activeWorkerCapacity{
		{InstanceID: "worker-a", CapacityLimit: 1},
		{InstanceID: "worker-b", CapacityLimit: 4},
	}
	allocations, global := allocateWorkerCapacities(workers, 5)
	if global != 5 || allocations["worker-a"] != 1 || allocations["worker-b"] != 4 {
		t.Fatalf("allocations=%v global=%d, want 1/4 and 5", allocations, global)
	}
}

func TestAllocateWorkerCapacitiesRedistributesFromPressuredWorker(t *testing.T) {
	workers := []activeWorkerCapacity{
		{InstanceID: "worker-a", CapacityLimit: 0},
		{InstanceID: "worker-b", CapacityLimit: 6},
		{InstanceID: "worker-c", CapacityLimit: 6},
	}
	allocations, global := allocateWorkerCapacities(workers, 8)
	if global != 8 || allocations["worker-a"] != 0 || allocations["worker-b"] != 4 || allocations["worker-c"] != 4 {
		t.Fatalf("allocations=%v global=%d, want 0/4/4 and 8", allocations, global)
	}
}
