268 lines
8.5 KiB
Go
268 lines
8.5 KiB
Go
package storage_test
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/clause"
|
|
|
|
"gpt-load/internal/platform/config"
|
|
"gpt-load/internal/platform/dbtx"
|
|
"gpt-load/internal/storage"
|
|
"gpt-load/internal/storage/models"
|
|
)
|
|
|
|
// TestExternalDatabaseConcurrentMigrations is opt-in so ordinary unit tests
|
|
// remain hermetic. It starts two migration attempts against the same external
|
|
// database to exercise the driver-level serialization contract.
|
|
func TestExternalDatabaseConcurrentMigrations(t *testing.T) {
|
|
dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_TEST_DSN"))
|
|
if dsn == "" {
|
|
t.Skip("GPT_LOAD_DATABASE_TEST_DSN is not set")
|
|
}
|
|
|
|
databases := make([]*gorm.DB, 2)
|
|
for index := range databases {
|
|
db, err := storage.OpenWithSource(dsn, config.DatabaseSourceExternal)
|
|
if err != nil {
|
|
t.Fatalf("OpenWithSource(%d) error = %v", index, err)
|
|
}
|
|
databases[index] = db
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
t.Fatalf("DB(%d) error = %v", index, err)
|
|
}
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
}
|
|
|
|
errorsCh := make(chan error, len(databases))
|
|
var group sync.WaitGroup
|
|
group.Add(len(databases))
|
|
for _, db := range databases {
|
|
go func(database *gorm.DB) {
|
|
defer group.Done()
|
|
errorsCh <- storage.AutoMigrate(database)
|
|
}(db)
|
|
}
|
|
group.Wait()
|
|
close(errorsCh)
|
|
for err := range errorsCh {
|
|
if err != nil {
|
|
t.Fatalf("concurrent AutoMigrate() error = %v", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// TestExternalDatabaseLifecycle is opt-in so ordinary unit tests remain
|
|
// hermetic. Set GPT_LOAD_DATABASE_TEST_DSN to one MySQL or PostgreSQL URL to
|
|
// exercise connection, migration, reserved-column queries, foreign keys, and
|
|
// the dialect-specific GORM upsert against a real server.
|
|
func TestExternalDatabaseLifecycle(t *testing.T) {
|
|
dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_TEST_DSN"))
|
|
if dsn == "" {
|
|
t.Skip("GPT_LOAD_DATABASE_TEST_DSN is not set")
|
|
}
|
|
|
|
db, err := storage.OpenWithSource(dsn, config.DatabaseSourceExternal)
|
|
if err != nil {
|
|
t.Fatalf("OpenWithSource() error = %v", err)
|
|
}
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
t.Fatalf("DB() error = %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
|
|
if err := storage.AutoMigrate(db); err != nil {
|
|
t.Fatalf("AutoMigrate() error = %v", err)
|
|
}
|
|
if err := storage.AutoMigrate(db); err != nil {
|
|
t.Fatalf("second AutoMigrate() error = %v", err)
|
|
}
|
|
for _, table := range []string{
|
|
"groups", "credentials", "access_keys", "request_logs", "request_log_attempts",
|
|
"usage_aggregation_journal", "usage_stats", "model_prices", "system_settings",
|
|
"jobs", "control_operations", "schema_migrations",
|
|
} {
|
|
if !db.Migrator().HasTable(table) {
|
|
t.Fatalf("table %q is missing", table)
|
|
}
|
|
}
|
|
|
|
suffix := strings.NewReplacer("/", "_", " ", "_").Replace(t.Name())
|
|
group := models.Group{
|
|
Name: fmt.Sprintf("integration-%s-%d", suffix, time.Now().UnixNano()),
|
|
ChannelID: "openai_compatible", Params: models.JSON(`{"base_url":"https://integration.example.com"}`),
|
|
Models: models.JSON(`[]`), Overrides: models.JSON(`{}`), Enabled: true,
|
|
}
|
|
if err := db.Create(&group).Error; err != nil {
|
|
t.Fatalf("create group: %v", err)
|
|
}
|
|
credential := models.Credential{
|
|
GroupID: group.ID, Data: "encrypted-data",
|
|
Fingerprint: fmt.Sprintf("integration-credential-%d", time.Now().UnixNano()),
|
|
}
|
|
if err := db.Create(&credential).Error; err != nil {
|
|
t.Fatalf("create credential: %v", err)
|
|
}
|
|
|
|
setting := models.SystemSetting{
|
|
// Keep this lifecycle-only row outside the public runtime-settings
|
|
// namespace. The same CI database also executes control workflow tests.
|
|
Key: models.InternalSystemSettingPrefix + "integration.key",
|
|
Value: `{}`,
|
|
UpdatedAtMS: time.Now().UnixMilli(),
|
|
}
|
|
if err := dbtx.Run(context.Background(), db, dbtx.Options{
|
|
Mode: dbtx.Write,
|
|
Operation: "external integration write transaction",
|
|
}, func(tx *gorm.DB) error {
|
|
return tx.Create(&setting).Error
|
|
}); err != nil {
|
|
t.Fatalf("create system setting: %v", err)
|
|
}
|
|
var settingCount int64
|
|
if err := dbtx.Run(context.Background(), db, dbtx.Options{
|
|
Mode: dbtx.ReadSnapshot,
|
|
Operation: "external integration read transaction",
|
|
}, func(tx *gorm.DB) error {
|
|
return tx.Model(&models.SystemSetting{}).
|
|
Where(&models.SystemSetting{Key: setting.Key}).Count(&settingCount).Error
|
|
}); err != nil {
|
|
t.Fatalf("read system setting in transaction: %v", err)
|
|
}
|
|
if settingCount != 1 {
|
|
t.Fatalf("system setting count in transaction = %d, want 1", settingCount)
|
|
}
|
|
var loadedSetting models.SystemSetting
|
|
if err := db.Where(&models.SystemSetting{Key: setting.Key}).First(&loadedSetting).Error; err != nil {
|
|
t.Fatalf("query system setting: %v", err)
|
|
}
|
|
|
|
stat := models.UsageStat{
|
|
BucketStartMS: time.Now().UnixMilli(),
|
|
AccessKeyID: 1,
|
|
ChannelID: group.ChannelID,
|
|
GroupID: group.ID,
|
|
CredentialID: credential.ID,
|
|
Model: "integration-model",
|
|
RequestCount: 1,
|
|
SuccessCount: 1,
|
|
}
|
|
upsert := clause.OnConflict{
|
|
Columns: []clause.Column{
|
|
{Name: "bucket_start_ms"}, {Name: "access_key_id"},
|
|
{Name: "channel_id"}, {Name: "group_id"},
|
|
{Name: "credential_id"}, {Name: "model"},
|
|
},
|
|
DoUpdates: clause.AssignmentColumns([]string{"request_count", "success_count"}),
|
|
}
|
|
if err := db.Clauses(upsert).Create(&stat).Error; err != nil {
|
|
t.Fatalf("create usage stat: %v", err)
|
|
}
|
|
stat.RequestCount = 2
|
|
stat.SuccessCount = 2
|
|
if err := db.Clauses(upsert).Create(&stat).Error; err != nil {
|
|
t.Fatalf("upsert usage stat: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestExternalDatabaseModelPriceIdentityUsesExactComparison protects the
|
|
// case-sensitive model ID contract on every supported external database.
|
|
// MySQL requires an explicit binary column collation; PostgreSQL already
|
|
// preserves case under its normal text semantics.
|
|
func TestExternalDatabaseModelPriceIdentityUsesExactComparison(t *testing.T) {
|
|
dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_TEST_DSN"))
|
|
if dsn == "" {
|
|
t.Skip("GPT_LOAD_DATABASE_TEST_DSN is not set")
|
|
}
|
|
|
|
db, err := storage.OpenWithSource(dsn, config.DatabaseSourceExternal)
|
|
if err != nil {
|
|
t.Fatalf("OpenWithSource() error = %v", err)
|
|
}
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
t.Fatalf("DB() error = %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
if err := storage.AutoMigrate(db); err != nil {
|
|
t.Fatalf("AutoMigrate() error = %v", err)
|
|
}
|
|
|
|
base := fmt.Sprintf("external-model-%d", time.Now().UnixNano())
|
|
upper := models.ModelPrice{ChannelID: "openai", ModelID: "Model-" + base}
|
|
lower := models.ModelPrice{ChannelID: "openai", ModelID: "model-" + base}
|
|
if err := db.Create(&upper).Error; err != nil {
|
|
t.Fatalf("create first case-distinct model price: %v", err)
|
|
}
|
|
if err := db.Create(&lower).Error; err != nil {
|
|
t.Fatalf("create second case-distinct model price: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestExternalDatabaseReadSnapshotUsesRepeatableRead is run against a MySQL
|
|
// connection configured with READ COMMITTED. The transaction helper must
|
|
// explicitly promote the one report transaction before requesting its stable
|
|
// snapshot; WITH CONSISTENT SNAPSHOT alone follows the session isolation.
|
|
func TestExternalDatabaseReadSnapshotUsesRepeatableRead(t *testing.T) {
|
|
dsn := strings.TrimSpace(os.Getenv("GPT_LOAD_DATABASE_TEST_DSN"))
|
|
if dsn == "" {
|
|
t.Skip("GPT_LOAD_DATABASE_TEST_DSN is not set")
|
|
}
|
|
|
|
db, err := storage.OpenWithSource(dsn, config.DatabaseSourceExternal)
|
|
if err != nil {
|
|
t.Fatalf("OpenWithSource() error = %v", err)
|
|
}
|
|
sqlDB, err := db.DB()
|
|
if err != nil {
|
|
t.Fatalf("DB() error = %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = sqlDB.Close() })
|
|
if err := storage.AutoMigrate(db); err != nil {
|
|
t.Fatalf("AutoMigrate() error = %v", err)
|
|
}
|
|
if db.Dialector.Name() != "mysql" {
|
|
t.Skip("repeatable-read setup is specific to MySQL")
|
|
}
|
|
|
|
modelID := fmt.Sprintf("external-snapshot-%d", time.Now().UnixNano())
|
|
err = dbtx.Run(t.Context(), db, dbtx.Options{
|
|
Mode: dbtx.ReadSnapshot,
|
|
Operation: "external read snapshot contract",
|
|
}, func(tx *gorm.DB) error {
|
|
var before int64
|
|
if err := tx.Model(&models.ModelPrice{}).
|
|
Where("model_id = ?", modelID).
|
|
Count(&before).Error; err != nil {
|
|
return err
|
|
}
|
|
if before != 0 {
|
|
return fmt.Errorf("initial model price count = %d, want 0", before)
|
|
}
|
|
if err := db.Create(&models.ModelPrice{ChannelID: "openai", ModelID: modelID}).Error; err != nil {
|
|
return err
|
|
}
|
|
var after int64
|
|
if err := tx.Model(&models.ModelPrice{}).
|
|
Where("model_id = ?", modelID).
|
|
Count(&after).Error; err != nil {
|
|
return err
|
|
}
|
|
if after != 0 {
|
|
return fmt.Errorf("snapshot model price count after concurrent write = %d, want 0", after)
|
|
}
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ReadSnapshot transaction error = %v", err)
|
|
}
|
|
}
|