//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.SaAIRateBucket{}) 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.SaAIRateBucket require.NoError(t, database.Where("user_id = ? AND action = ?", user.ID, action).First(&bucket).Error) require.Equal(t, limit, bucket.Count) }