easyai-ai-gateway/apps/api/internal/store/rate_limit_policy_test.go

package store

import "testing"

func TestEffectiveRateLimitPolicyPrecedence(t *testing.T) {
	policy := func(limit float64) map[string]any {
		return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
	}
	tests := []struct {
		name  string
		input EffectiveRateLimitPolicyInput
		want  float64
		ok    bool
	}{
		{name: "base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2)}, want: 2, ok: true},
		{name: "platform", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: policy(4)}, want: 4, ok: true},
		{name: "empty platform inherits base", input: EffectiveRateLimitPolicyInput{BasePolicy: policy(2), PlatformPolicy: map[string]any{"rules": []any{}}}, want: 2, ok: true},
		{name: "runtime", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicy: policy(8), RuntimePolicyExplicit: true}, want: 8, ok: true},
		{name: "runtime override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), RuntimePolicyOverride: map[string]any{"rateLimitPolicy": policy(16)}}, want: 16, ok: true},
		{name: "model override", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "override"}, want: 32, ok: true},
		{name: "model explicit unlimited", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: map[string]any{}, ModelPolicyMode: "override"}, ok: false},
		{name: "model inherit", input: EffectiveRateLimitPolicyInput{PlatformPolicy: policy(4), ModelPolicy: policy(32), ModelPolicyMode: "inherit"}, want: 4, ok: true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, ok := ConcurrentPolicyCapacity(EffectiveRateLimitPolicy(tt.input))
			if ok != tt.ok || (ok && got != int(tt.want)) {
				t.Fatalf("capacity = (%d, %v), want (%d, %v)", got, ok, int(tt.want), tt.ok)
			}
		})
	}
}

func TestNormalizeRateLimitPolicyLegacyShapes(t *testing.T) {
	tests := []struct {
		name   string
		policy map[string]any
		metric string
		want   float64
	}{
		{
			name:   "platform concurrent",
			policy: map[string]any{"platformLimits": map[string]any{"max_concurrent_requests": 5.0}},
			metric: "concurrent",
			want:   5,
		},
		{
			name:   "model concurrent camel case",
			policy: map[string]any{"modelLimits": map[string]any{"maxConcurrentRequests": 10.0}},
			metric: "concurrent",
			want:   10,
		},
		{
			name: "stricter duplicate wins",
			policy: map[string]any{
				"platformLimits": map[string]any{"max_concurrent_requests": 8.0},
				"modelLimits":    map[string]any{"max_concurrent_requests": 3.0},
			},
			metric: "concurrent",
			want:   3,
		},
		{
			name:   "requests per minute",
			policy: map[string]any{"model_limits": map[string]any{"max_request_per_minute": 60.0}},
			metric: "rpm",
			want:   60,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got, ok := RateLimitPolicyMetric(tt.policy, tt.metric)
			if !ok || got != tt.want {
				t.Fatalf("metric %s = (%v, %v), want (%v, true)", tt.metric, got, ok, tt.want)
			}
		})
	}
}

func TestConcurrentPolicyCapacityRoundsDown(t *testing.T) {
	policy := map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": 96.9}}}
	got, ok := ConcurrentPolicyCapacity(policy)
	if !ok || got != 96 {
		t.Fatalf("capacity = (%d, %v), want (96, true)", got, ok)
	}
}

func TestAsyncWorkerCapacityAggregation(t *testing.T) {
	policy := func(limit float64) map[string]any {
		return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
	}
	tests := []struct {
		name         string
		policies     []map[string]any
		hardLimit    int
		wantCapacity int
		wantDesired  int
		wantCapped   bool
	}{
		{name: "no enabled models", hardLimit: 2048, wantCapacity: 1, wantDesired: 1},
		{name: "finite sum", policies: []map[string]any{policy(64), policy(32)}, hardLimit: 2048, wantCapacity: 96, wantDesired: 96},
		{name: "unlimited model", policies: []map[string]any{policy(64), {}}, hardLimit: 2048, wantCapacity: 2048, wantDesired: 2048},
		{name: "hard limit cap", policies: []map[string]any{policy(80), policy(80)}, hardLimit: 96, wantCapacity: 96, wantDesired: 160, wantCapped: true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := asyncWorkerCapacityFromPolicies(tt.policies, tt.hardLimit)
			if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped {
				t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped)
			}
		})
	}
}

func TestAsyncWorkerCapacityRespectsUserGroupCeiling(t *testing.T) {
	policy := func(limit float64) map[string]any {
		return map[string]any{"rules": []any{map[string]any{"metric": "concurrent", "limit": limit}}}
	}
	tests := []struct {
		name         string
		models       []map[string]any
		groups       []map[string]any
		hardLimit    int
		wantCapacity int
		wantDesired  int
		wantCapped   bool
	}{
		{
			name:         "group ceiling prevents worker oversubscription",
			models:       []map[string]any{{}},
			groups:       []map[string]any{policy(3), policy(10)},
			hardLimit:    2048,
			wantCapacity: 13,
			wantDesired:  13,
		},
		{
			name:         "model ceiling is stricter",
			models:       []map[string]any{policy(5), policy(7)},
			groups:       []map[string]any{policy(300)},
			hardLimit:    2048,
			wantCapacity: 12,
			wantDesired:  12,
		},
		{
			name:         "both policy sets unlimited use hard limit",
			models:       []map[string]any{{}},
			groups:       []map[string]any{{}},
			hardLimit:    256,
			wantCapacity: 256,
			wantDesired:  256,
		},
		{
			name:         "finite group desired still reports hard cap",
			models:       []map[string]any{{}},
			groups:       []map[string]any{policy(500)},
			hardLimit:    256,
			wantCapacity: 256,
			wantDesired:  500,
			wantCapped:   true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := asyncWorkerCapacityFromPolicySets(tt.models, tt.groups, tt.hardLimit)
			if got.Capacity != tt.wantCapacity || got.Desired != tt.wantDesired || got.Capped != tt.wantCapped {
				t.Fatalf("snapshot=%+v, want capacity=%d desired=%d capped=%v", got, tt.wantCapacity, tt.wantDesired, tt.wantCapped)
			}
		})
	}
}