73 lines
1.9 KiB
Go
73 lines
1.9 KiB
Go
//go:build integration
|
|
|
|
package ai
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"os"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/logger"
|
|
"senlinai-agent/backend/internal/models"
|
|
)
|
|
|
|
func TestPostgresReserveRateLimitIsAtomicAcrossConcurrentConnections(t *testing.T) {
|
|
dsn := os.Getenv("TEST_DATABASE_URL")
|
|
require.NotEmpty(t, dsn, "TEST_DATABASE_URL is required for integration tests and must point to an isolated database")
|
|
database, err := gorm.Open(postgres.Open(dsn), &gorm.Config{TranslateError: true, Logger: logger.Default.LogMode(logger.Silent)})
|
|
require.NoError(t, err)
|
|
require.NoError(t, models.AutoMigrate(database))
|
|
models.DBService = database
|
|
suffix := fmt.Sprint(time.Now().UnixNano())
|
|
user := createAIRateTestUser(t, database, "postgres-rate-"+suffix+"@example.com")
|
|
action := "postgres_concurrent_" + suffix
|
|
t.Cleanup(func() {
|
|
database.Where("user_id = ?", user.ID).Delete(&models.SenlinAgentAIRateBucket{})
|
|
database.Delete(&user)
|
|
})
|
|
gateway := NewGatewayWithSecret("system-key", "test-encryption-secret")
|
|
const (
|
|
limit = 9
|
|
attempts = 48
|
|
)
|
|
|
|
start := make(chan struct{})
|
|
results := make(chan error, attempts)
|
|
var wait sync.WaitGroup
|
|
for range attempts {
|
|
wait.Add(1)
|
|
go func() {
|
|
defer wait.Done()
|
|
<-start
|
|
results <- gateway.ReserveRateLimit(user.ID, action, limit, time.Hour)
|
|
}()
|
|
}
|
|
close(start)
|
|
wait.Wait()
|
|
close(results)
|
|
|
|
allowed := 0
|
|
limited := 0
|
|
for err := range results {
|
|
switch {
|
|
case err == nil:
|
|
allowed++
|
|
case errors.Is(err, ErrAIRateLimited):
|
|
limited++
|
|
default:
|
|
require.NoError(t, err)
|
|
}
|
|
}
|
|
require.Equal(t, limit, allowed)
|
|
require.Equal(t, attempts-limit, limited)
|
|
var bucket models.SenlinAgentAIRateBucket
|
|
require.NoError(t, database.Where("user_id = ? AND action = ?", user.ID, action).First(&bucket).Error)
|
|
require.Equal(t, limit, bucket.Count)
|
|
}
|