promptkit/internal/capacity/manager_test.go

package capacity

import (
	"context"
	"errors"
	"testing"

	"gitea.maximumdirect.net/eric/promptkit/internal/domain"
)

func TestNewManagerRejectsInvalidPolicies(t *testing.T) {
	maxInt := int(^uint(0) >> 1)
	tests := []struct {
		name   string
		id     string
		policy domain.BackendCapacityPolicy
	}{
		{
			name:   "blank ID",
			id:     " \t ",
			policy: domain.BackendCapacityPolicy{ConcurrencyLimit: 1},
		},
		{
			name:   "zero concurrency",
			id:     "backend",
			policy: domain.BackendCapacityPolicy{},
		},
		{
			name:   "negative concurrency",
			id:     "backend",
			policy: domain.BackendCapacityPolicy{ConcurrencyLimit: -1},
		},
		{
			name: "negative queue",
			id:   "backend",
			policy: domain.BackendCapacityPolicy{
				ConcurrencyLimit: 1,
				QueueCapacity:    -1,
			},
		},
		{
			name: "total overflow",
			id:   "backend",
			policy: domain.BackendCapacityPolicy{
				ConcurrencyLimit: maxInt,
				QueueCapacity:    1,
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			_, err := NewManager(map[string]domain.BackendCapacityPolicy{
				tc.id: tc.policy,
			})
			if err == nil {
				t.Fatal("expected invalid policy error")
			}
		})
	}
}

func TestManagerAdmissionIsBoundedAndReleaseIsIdempotent(t *testing.T) {
	policies := map[string]domain.BackendCapacityPolicy{
		"limited": {
			ConcurrencyLimit: 2,
			QueueCapacity:    1,
		},
		"independent": {
			ConcurrencyLimit: 1,
		},
	}
	manager, err := NewManager(policies)
	if err != nil {
		t.Fatalf("construct manager: %v", err)
	}
	policies["limited"] = domain.BackendCapacityPolicy{
		ConcurrencyLimit: 100,
		QueueCapacity:    100,
	}

	releases := make([]func(), 0, 3)
	for range 3 {
		release, err := manager.Admit(context.Background(), "limited")
		if err != nil {
			t.Fatalf("admit within configured capacity: %v", err)
		}
		releases = append(releases, release)
	}
	if release, err := manager.Admit(context.Background(), "limited"); release != nil ||
		!errors.Is(err, ErrCapacityExceeded) {
		t.Fatalf("admission beyond capacity=(release=%t, err=%v), want ErrCapacityExceeded",
			release != nil, err)
	}
	independentRelease, err := manager.Admit(context.Background(), "independent")
	if err != nil {
		t.Fatalf("admit independent backend while first is full: %v", err)
	}
	independentRelease()

	releases[0]()
	releases[0]()
	replacement, err := manager.Admit(context.Background(), "limited")
	if err != nil {
		t.Fatalf("admit after release: %v", err)
	}
	replacement()
	releases[1]()
	releases[2]()

	pool := manager.pools["limited"]
	pool.mu.Lock()
	admitted := pool.admitted
	pool.mu.Unlock()
	if admitted != 0 {
		t.Fatalf("admitted runs after releases=%d, want 0", admitted)
	}
}

func TestManagerAdmissionHonorsContextAndUnlimitedBackends(t *testing.T) {
	manager, err := NewManager(map[string]domain.BackendCapacityPolicy{
		"limited": {ConcurrencyLimit: 1},
	})
	if err != nil {
		t.Fatalf("construct manager: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	if release, err := manager.Admit(ctx, "limited"); release != nil ||
		!errors.Is(err, context.Canceled) {
		t.Fatalf("canceled limited admission=(release=%t, err=%v), want context cancellation",
			release != nil, err)
	}

	var nilManager *Manager
	for _, tc := range []struct {
		name      string
		manager   *Manager
		backendID string
	}{
		{name: "nil manager", manager: nilManager, backendID: "limited"},
		{name: "blank ID", manager: manager},
		{name: "unknown ID", manager: manager, backendID: "unknown"},
	} {
		t.Run(tc.name, func(t *testing.T) {
			release, err := tc.manager.Admit(ctx, tc.backendID)
			if err != nil {
				t.Fatalf("unlimited admission: %v", err)
			}
			if release == nil {
				t.Fatal("unlimited admission returned nil release")
			}
			release()
			release()
		})
	}
}