87 lines
3.5 KiB
Go
87 lines
3.5 KiB
Go
package environment
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"fmt"
|
|
"net/url"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestOpenAISettingsSchemaFreshUpgradeAndRestart(t *testing.T) {
|
|
databaseURL := os.Getenv("CREATORHUB_POSTGRES_TEST_URL")
|
|
if databaseURL == "" {
|
|
t.Skip("set CREATORHUB_POSTGRES_TEST_URL to run migration tests")
|
|
}
|
|
ctx := context.Background()
|
|
admin, err := sql.Open("pgx", databaseURL)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer admin.Close()
|
|
schema := fmt.Sprintf("openai_upgrade_%d", time.Now().UnixNano())
|
|
if _, err := admin.ExecContext(ctx, "CREATE SCHEMA "+schema); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer func() {
|
|
if _, err := admin.ExecContext(ctx, "DROP SCHEMA "+schema+" CASCADE"); err != nil {
|
|
t.Errorf("cleanup: %v", err)
|
|
}
|
|
}()
|
|
parsed, err := url.Parse(databaseURL)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
query := parsed.Query()
|
|
query.Set("search_path", schema)
|
|
parsed.RawQuery = query.Encode()
|
|
store, err := Open(ctx, parsed.String())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer store.Close()
|
|
assertColumns := func() {
|
|
t.Helper()
|
|
var columns int
|
|
if err := store.db.QueryRowContext(ctx, `SELECT count(*) FROM information_schema.columns WHERE table_schema=current_schema() AND table_name='creator_settings' AND column_name IN ('ai_base_url','ai_api_key','ai_model')`).Scan(&columns); err != nil || columns != 3 {
|
|
t.Fatalf("AI columns=%d err=%v", columns, err)
|
|
}
|
|
if err := store.db.QueryRowContext(ctx, `SELECT count(*) FROM information_schema.columns WHERE table_schema=current_schema() AND table_name='creator_settings' AND column_name IN ('ai_provider','ai_configured')`).Scan(&columns); err != nil || columns != 0 {
|
|
t.Fatalf("legacy columns=%d err=%v", columns, err)
|
|
}
|
|
}
|
|
assertColumns()
|
|
// Recreate the previous schema only in this isolated database, preserving its settings row.
|
|
if _, err := store.db.ExecContext(ctx, `ALTER TABLE creator_settings DROP COLUMN ai_base_url,DROP COLUMN ai_api_key,ADD COLUMN ai_provider text NOT NULL DEFAULT '',ADD COLUMN ai_configured boolean NOT NULL DEFAULT false;
|
|
UPDATE creator_settings SET lookback_days=12,ai_model='legacy-model',ai_provider='bailian',ai_configured=true;
|
|
DELETE FROM schema_migration WHERE version=1057`); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := store.migrate(ctx); err != nil {
|
|
t.Fatalf("upgrade: %v", err)
|
|
}
|
|
assertColumns()
|
|
var days int
|
|
var model, base, key string
|
|
if err := store.db.QueryRowContext(ctx, `SELECT lookback_days,ai_model,ai_base_url,ai_api_key FROM creator_settings`).Scan(&days, &model, &base, &key); err != nil || days != 12 || model != "legacy-model" || base != "" || key != "" {
|
|
t.Fatalf("preserved settings=%d/%s/%s/%s err=%v", days, model, base, key, err)
|
|
}
|
|
if _, err := store.db.ExecContext(ctx, `UPDATE creator_settings SET ai_base_url='https://example.com/v1',ai_api_key='saved-key',ai_model='chosen-model'`); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for i := 0; i < 2; i++ {
|
|
if err := store.migrate(ctx); err != nil {
|
|
t.Fatalf("repeat startup: %v", err)
|
|
}
|
|
}
|
|
if err := store.db.QueryRowContext(ctx, `SELECT lookback_days,ai_model,ai_base_url,ai_api_key FROM creator_settings`).Scan(&days, &model, &base, &key); err != nil || days != 12 || model != "chosen-model" || base != "https://example.com/v1" || key != "saved-key" {
|
|
t.Fatalf("restart changed settings: %d/%s/%s/%s err=%v", days, model, base, key, err)
|
|
}
|
|
var applied int
|
|
if err := store.db.QueryRowContext(ctx, `SELECT count(*) FROM schema_migration WHERE version=1057`).Scan(&applied); err != nil || applied != 1 {
|
|
t.Fatalf("applied=%d err=%v", applied, err)
|
|
}
|
|
}
|