xinmi/internal/workbench/quota_test.go

package workbench

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

func seedQuota(t *testing.T, s *Service, quota UsageQuota) UsageQuota {
	t.Helper()
	quota.ID = uuid.NewString()
	if quota.Period == "" {
		quota.Period = "daily"
	}
	if quota.OveragePolicy == "" {
		quota.OveragePolicy = "hard_limit"
	}
	if err := s.db.Create(&quota).Error; err != nil {
		t.Fatal(err)
	}
	return quota
}

func TestQuotaReservationPoliciesAndRollback(t *testing.T) {
	for _, mode := range []string{"token", "points", "concurrency", "soft", "warn", "disabled", "unlimited"} {
		t.Run(mode, func(t *testing.T) {
			s := sqliteWorkbench(t)
			ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
			seedBalance(t, s, 1, "alice", 100000)
			quota := UsageQuota{TenantID: 1, UserID: "alice", TokenLimit: 100, Enabled: true, WarningRatio: .8}
			switch mode {
			case "points":
				quota.TokenLimit, quota.PointsLimit = 0, 1000
			case "concurrency":
				quota.TokenLimit, quota.ConcurrentLimit = 0, 1
			case "soft":
				quota.OveragePolicy = "soft_limit"
			case "warn":
				quota.OveragePolicy = "warn"
			case "disabled":
				quota.Enabled = false
			case "unlimited":
				quota.TokenLimit, quota.WarningRatio = 0, 0
			}
			seedQuota(t, s, quota)
			first, err := s.reserveBalance(ctx, 900)
			if err != nil {
				t.Fatal(err)
			}
			second, err := s.reserveBalance(ctx, 200)
			blocked := mode == "token" || mode == "points" || mode == "concurrency" || mode == "soft"
			if blocked != errors.Is(err, ErrInvalid) {
				t.Fatalf("blocked=%v, error=%v", blocked, err)
			}
			wantReserved := int64(900)
			if !blocked {
				wantReserved += 200
			}
			if got := accountBalance(t, s, 1, "alice"); got.Reserved != wantReserved || got.Balance != 100000 {
				t.Fatalf("failed reservation leaked: %+v", got)
			}
			settled := Usage{ID: uuid.NewString(), CallID: first.ID, TenantID: 1, OwnerID: "alice", InputTokens: 30, Points: 300, CreatedAt: first.CreatedAt, Metadata: JSON{}}
			if err := s.settleBalance(ctx, first, settled); err != nil {
				t.Fatal(err)
			}
			if blocked {
				if _, err := s.reserveBalance(ctx, 200); err != nil {
					t.Fatalf("released capacity unavailable: %v", err)
				}
			} else if second == nil {
				t.Fatal("missing accepted hold")
			}
			var warnings int64
			s.db.Model(&Event{}).Count(&warnings)
			if mode == "unlimited" || mode == "disabled" || mode == "concurrency" {
				if warnings != 0 {
					t.Fatal("unlimited or disabled quota warned")
				}
			} else if warnings != 1 {
				t.Fatalf("expected one warning per period, got %d", warnings)
			}
		})
	}
}

func TestQuotaPeriodsIsolationAndExpiredCalls(t *testing.T) {
	s := sqliteWorkbench(t)
	now := time.Date(2026, 9, 30, 16, 0, 0, 0, time.UTC) // October 1 midnight in Shanghai.
	for _, period := range []string{"daily", "monthly"} {
		start, end := quotaPeriod(period, now)
		if !start.Equal(now) || end.In(quotaLocation).Hour() != 0 {
			t.Fatalf("wrong period: %v %v", start, end)
		}
		if period == "monthly" && end.In(quotaLocation).Month() != time.November {
			t.Fatal("month boundary wrong")
		}
	}
	month, monthEnd := quotaPeriod("monthly", time.Date(2028, 2, 29, 6, 0, 0, 0, time.UTC))
	if month.Day() != 1 || monthEnd.Month() != time.March || monthEnd.Sub(month) != 29*24*time.Hour {
		t.Fatal("leap year boundary wrong")
	}
	quota := UsageQuota{TenantID: 1, UserID: "alice", Period: "daily", TokenLimit: 100, PointsLimit: 1000, ConcurrentLimit: 2, WarningRatio: .8, Enabled: true, OveragePolicy: "hard_limit"}
	_, end := quotaPeriod("daily", now)
	for _, row := range []Usage{
		{TenantID: 1, OwnerID: "alice", InputTokens: 1000, CreatedAt: now.Add(-time.Nanosecond)},
		{TenantID: 1, OwnerID: "alice", InputTokens: 20, Points: 200, CreatedAt: now},
		{TenantID: 1, OwnerID: "alice", InputTokens: 1000, CreatedAt: end},
		{TenantID: 2, OwnerID: "alice", InputTokens: 1000, CreatedAt: now},
		{TenantID: 1, OwnerID: "bob", InputTokens: 1000, CreatedAt: now},
		{TenantID: 1, OwnerID: "alice", InputTokens: 1000, CreatedAt: now, Status: "void"},
	} {
		row.ID, row.CallID, row.Metadata = uuid.NewString(), uuid.NewString(), JSON{}
		if err := s.db.Create(&row).Error; err != nil {
			t.Fatal(err)
		}
	}
	for _, row := range []BalanceHold{
		{TenantID: 1, UserID: "alice", Amount: 500, Status: "reserved", CreatedAt: now, ExpiresAt: now.Add(time.Minute)},
		{TenantID: 1, UserID: "alice", Amount: 900, Status: "reserved", CreatedAt: now.Add(-time.Minute), ExpiresAt: now.Add(time.Minute)},
		{TenantID: 1, UserID: "alice", Amount: 900, Status: "reserved", CreatedAt: now, ExpiresAt: now},
		{TenantID: 2, UserID: "alice", Amount: 900, Status: "reserved", CreatedAt: now, ExpiresAt: now.Add(time.Minute)},
	} {
		row.ID = uuid.NewString()
		if err := s.db.Create(&row).Error; err != nil {
			t.Fatal(err)
		}
	}
	state, err := loadQuotaState(s.db, quota, now)
	if err != nil {
		t.Fatal(err)
	}
	if state.Used.Tokens != 20 || state.Used.Points != 200 || state.Used.ReservedTokens != 50 || state.Used.ReservedPoints != 500 || state.Used.ActiveCalls != 2 || state.Allowed || state.Warning {
		t.Fatalf("bad state: %+v", state)
	}
	quota.WarningRatio = 0
	state, err = loadQuotaState(s.db, quota, now)
	if err != nil || !state.Warning {
		t.Fatalf("zero warning threshold ignored: %+v %v", state, err)
	}
}

func TestQuotaSettlementNormalizesCallerTimezone(t *testing.T) {
	s := sqliteWorkbench(t)
	ctx := workbenchContext(1, "alice", types.TenantRoleContributor)
	now := time.Date(2026, 10, 1, 0, 0, 0, 0, quotaLocation)
	for i, at := range []time.Time{now.Add(-time.Nanosecond), now, now.AddDate(0, 0, 1)} {
		usage := Usage{ID: uuid.NewString(), CallID: uuid.NewString(), TenantID: 1, OwnerID: "alice", InputTokens: int64(i + 1), Points: int64((i + 1) * 10), CreatedAt: at, Metadata: JSON{}}
		if err := s.settleBalance(ctx, nil, usage); err != nil {
			t.Fatal(err)
		}
	}
	quota := UsageQuota{TenantID: 1, UserID: "alice", Period: "daily"}
	for _, clock := range []time.Time{now, now.UTC(), now.In(time.FixedZone("west", -7*60*60))} {
		state, err := loadQuotaState(s.db, quota, clock)
		if err != nil || state.Used.Tokens != 2 || state.Used.Points != 20 {
			t.Fatalf("timezone changed period usage: %+v, %v", state, err)
		}
	}
}

func quotaRequest(t *testing.T, s *Service, ctx context.Context, method, path string, body any) *httptest.ResponseRecorder {
	t.Helper()
	gin.SetMode(gin.TestMode)
	router := gin.New()
	h := NewHandler(s)
	router.GET("/quotas", h.Quotas)
	router.PUT("/quotas/:userID", h.SaveQuota)
	router.DELETE("/quotas/:id", h.DeleteQuota)
	encoded, _ := json.Marshal(body)
	request := httptest.NewRequest(method, path, bytes.NewReader(encoded)).WithContext(ctx)
	request.Header.Set("Content-Type", "application/json")
	response := httptest.NewRecorder()
	router.ServeHTTP(response, request)
	return response
}

func TestQuotaAdminValidationUpsertAndTenantIsolation(t *testing.T) {
	s := sqliteWorkbench(t)
	if err := s.db.AutoMigrate(&types.TenantMember{}, &types.User{}); err != nil {
		t.Fatal(err)
	}
	if err := s.db.Create(&types.User{ID: "alice", Username: "Alice", Email: "alice@example.test"}).Error; err != nil {
		t.Fatal(err)
	}
	if err := s.db.Create(&types.TenantMember{TenantID: 1, UserID: "alice", Role: types.TenantRoleContributor}).Error; err != nil {
		t.Fatal(err)
	}
	admin := workbenchContext(1, "admin", types.TenantRoleAdmin)
	body := UsageQuota{Period: "daily", Enabled: true, TokenLimit: 100, WarningRatio: .8, OveragePolicy: "hard_limit"}
	for _, role := range []types.TenantRole{types.TenantRoleViewer, types.TenantRoleContributor} {
		ctx := workbenchContext(1, "alice", role)
		for _, method := range []string{http.MethodGet, http.MethodPut, http.MethodDelete} {
			path := "/quotas"
			if method != http.MethodGet {
				path += "/alice"
			}
			if got := quotaRequest(t, s, ctx, method, path, body).Code; got != 403 {
				t.Fatalf("%s role %s status %d", method, role, got)
			}
		}
	}
	if response := quotaRequest(t, s, admin, "PUT", "/quotas/alice", body); response.Code != 200 {
		t.Fatal(response.Body.String())
	}
	body.TokenLimit, body.Enabled, body.WarningRatio = 0, false, 0
	if response := quotaRequest(t, s, admin, "PUT", "/quotas/alice", body); response.Code != 200 {
		t.Fatal(response.Body.String())
	}
	var rows []UsageQuota
	s.db.Find(&rows)
	if len(rows) != 1 || rows[0].TokenLimit != 0 || rows[0].Enabled || rows[0].WarningRatio != 0 {
		t.Fatalf("upsert duplicated or zero fields lost: %+v", rows)
	}
	if response := quotaRequest(t, s, admin, "GET", "/quotas", nil); response.Code != 200 || !strings.Contains(response.Body.String(), "Alice") {
		t.Fatal(response.Body.String())
	}
	other := workbenchContext(2, "admin", types.TenantRoleAdmin)
	if response := quotaRequest(t, s, other, "GET", "/quotas", nil); response.Code != 200 || strings.Contains(response.Body.String(), "Alice") {
		t.Fatal("cross tenant list")
	}
	if got := quotaRequest(t, s, other, "DELETE", "/quotas/"+rows[0].ID, nil).Code; got != 404 {
		t.Fatalf("cross tenant delete: %d", got)
	}
	if got := quotaRequest(t, s, other, "PUT", "/quotas/alice", body).Code; got != 404 {
		t.Fatalf("cross tenant upsert: %d", got)
	}
	for _, invalid := range []UsageQuota{
		{Period: "weekly", OveragePolicy: "warn"}, {Period: "daily", OveragePolicy: "free"}, {Period: "daily", OveragePolicy: "warn", PointsLimit: -1}, {Period: "daily", OveragePolicy: "warn", WarningRatio: 1.1}, {Period: "daily", OveragePolicy: "warn", ConcurrentLimit: 10001},
	} {
		if got := quotaRequest(t, s, admin, "PUT", "/quotas/alice", invalid).Code; got != 400 {
			t.Fatalf("invalid quota accepted: %+v", invalid)
		}
	}
	if got := quotaRequest(t, s, admin, "DELETE", "/quotas/"+rows[0].ID, nil).Code; got != 200 {
		t.Fatalf("delete status %d", got)
	}
}

func TestPostgresQuotaReservationRaces(t *testing.T) {
	dsn := os.Getenv("WORKBENCH_TEST_POSTGRES_DSN")
	if dsn == "" {
		t.Skip("WORKBENCH_TEST_POSTGRES_DSN is not set")
	}
	db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal("PostgreSQL unavailable")
	}
	connection, _ := db.DB()
	defer connection.Close()
	schema := "wb_quota_test_" + strings.ReplaceAll(uuid.NewString(), "-", "")
	if err := db.Exec("CREATE SCHEMA " + schema).Error; err != nil {
		t.Fatal(err)
	}
	defer db.Exec("DROP SCHEMA " + schema + " CASCADE")
	testDSN := dsn + " search_path=" + schema
	if strings.Contains(dsn, "://") {
		separator := "?"
		if strings.Contains(dsn, "?") {
			separator = "&"
		}
		testDSN = dsn + separator + "search_path=" + schema
	}
	testDB, err := gorm.Open(postgres.Open(testDSN), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	if err != nil {
		t.Fatal("isolated PostgreSQL connection unavailable")
	}
	pool, _ := testDB.DB()
	defer pool.Close()
	if err := testDB.AutoMigrate(&Balance{}, &BalanceHold{}, &Usage{}, &UsageQuota{}, &Transaction{}, &Event{}); err != nil {
		t.Fatal(err)
	}
	s := NewService(testDB)
	for _, mode := range []string{"token", "points", "concurrency", "unlimited"} {
		t.Run(mode, func(t *testing.T) {
			ctx := workbenchContext(1, mode, types.TenantRoleContributor)
			seedBalance(t, s, 1, mode, 1000000)
			q := UsageQuota{TenantID: 1, UserID: mode, Enabled: true, WarningRatio: .8}
			want := 2
			switch mode {
			case "token":
				q.TokenLimit = 200
			case "points":
				q.PointsLimit = 2000
			case "concurrency":
				q.ConcurrentLimit = 2
			case "unlimited":
				want = 12
			}
			seedQuota(t, s, q)
			var workers sync.WaitGroup
			results := make(chan *BalanceHold, 12)
			failures := make(chan error, 12)
			start := make(chan struct{})
			for i := 0; i < 12; i++ {
				workers.Add(1)
				go func() {
					defer workers.Done()
					<-start
					hold, err := s.reserveBalance(ctx, 1000)
					if err != nil {
						failures <- err
					} else {
						results <- hold
					}
				}()
			}
			close(start)
			workers.Wait()
			close(results)
			close(failures)
			if len(results) != want {
				t.Fatalf("accepted %d calls; want %d", len(results), want)
			}
			for err := range failures {
				if !errors.Is(err, ErrInvalid) {
					t.Fatalf("unexpected reservation error: %v", err)
				}
			}
			if balance := accountBalance(t, s, 1, mode); balance.Reserved != int64(want*1000) {
				t.Fatalf("hold/account mismatch: %+v", balance)
			}
			var holds []*BalanceHold
			for hold := range results {
				holds = append(holds, hold)
			}
			// Concurrent settlements release all slots and remain idempotent.
			settleErrors := make(chan error, len(holds)*2)
			for _, hold := range holds {
				for repeat := 0; repeat < 2; repeat++ {
					workers.Add(1)
					go func(hold *BalanceHold) {
						defer workers.Done()
						settleErrors <- s.settleBalance(ctx, hold, Usage{ID: hold.ID, CallID: hold.ID, TenantID: 1, OwnerID: mode, InputTokens: 10, Points: 100, CreatedAt: hold.CreatedAt, Metadata: JSON{}})
					}(hold)
				}
			}
			workers.Wait()
			close(settleErrors)
			for err := range settleErrors {
				if err != nil {
					t.Fatal(err)
				}
			}
			if balance := accountBalance(t, s, 1, mode); balance.Reserved != 0 || balance.Balance != 1000000-int64(want*100) {
				t.Fatalf("settlement mismatch: %+v", balance)
			}
			if _, err := s.reserveBalance(ctx, 1000); err != nil {
				t.Fatalf("settled capacity not reusable: %v", err)
			}
		})
	}
}