Files
NianAIGC/backend/internal/postgres/runtime_settings_test.go
2026-08-18 12:41:39 +08:00

122 lines
3.5 KiB
Go

package postgres
import (
"context"
"errors"
"reflect"
"testing"
)
func TestLoadRuntimeSettingsUsesRequestedKeys(t *testing.T) {
pool := &runtimeSettingsPool{rows: [][]any{{"EVOLINK_API_KEY", "secret"}, {"IMAGE_GENERATE_ENGINE", "evolink"}}}
database := NewDatabase(Config{Backend: BackendPostgres}, pool)
keys := []string{"EVOLINK_API_KEY", "IMAGE_GENERATE_ENGINE"}
values, err := database.LoadRuntimeSettings(context.Background(), keys)
if err != nil {
t.Fatal(err)
}
if pool.query != LoadRuntimeSettingsSQL || !reflect.DeepEqual(pool.args, []any{keys}) {
t.Fatalf("query=%q args=%#v", pool.query, pool.args)
}
want := map[string]string{"EVOLINK_API_KEY": "secret", "IMAGE_GENERATE_ENGINE": "evolink"}
if !reflect.DeepEqual(values, want) {
t.Fatalf("values=%#v want=%#v", values, want)
}
}
func TestSaveRuntimeSettingsUpsertsSortedKeysInOneTransaction(t *testing.T) {
transaction := &runtimeSettingsTransaction{}
pool := &runtimeSettingsPool{transaction: transaction}
database := NewDatabase(Config{Backend: BackendPostgres}, pool)
err := database.SaveRuntimeSettings(context.Background(), map[string]string{
"VIDEO_GENERATE_ENGINE": "seedance",
"BAILIAN_API_KEY": "secret",
})
if err != nil {
t.Fatal(err)
}
want := [][]any{{"BAILIAN_API_KEY", "secret"}, {"VIDEO_GENERATE_ENGINE", "seedance"}}
if !reflect.DeepEqual(transaction.execArgs, want) {
t.Fatalf("exec args=%#v want=%#v", transaction.execArgs, want)
}
if transaction.commits != 1 {
t.Fatalf("commits=%d", transaction.commits)
}
if transaction.rollbacks != 0 {
t.Fatalf("rollbacks=%d", transaction.rollbacks)
}
}
type runtimeSettingsPool struct {
rows [][]any
query string
args []any
transaction *runtimeSettingsTransaction
}
func (pool *runtimeSettingsPool) Query(_ context.Context, query string, args ...any) (Rows, error) {
pool.query, pool.args = query, args
return &runtimeSettingsRows{rows: pool.rows}, nil
}
func (pool *runtimeSettingsPool) Begin(context.Context) (Transaction, error) {
if pool.transaction == nil {
return nil, errors.New("transaction unavailable")
}
return pool.transaction, nil
}
type runtimeSettingsTransaction struct {
execArgs [][]any
commits int
rollbacks int
}
func (*runtimeSettingsTransaction) Query(context.Context, string, ...any) (Rows, error) {
return &runtimeSettingsRows{}, nil
}
func (transaction *runtimeSettingsTransaction) Exec(_ context.Context, query string, args ...any) error {
if query != UpsertRuntimeSettingSQL {
return errors.New("unexpected runtime settings query")
}
transaction.execArgs = append(transaction.execArgs, args)
return nil
}
func (transaction *runtimeSettingsTransaction) Commit(context.Context) error {
transaction.commits++
return nil
}
func (transaction *runtimeSettingsTransaction) Rollback(context.Context) error {
transaction.rollbacks++
return nil
}
type runtimeSettingsRows struct {
rows [][]any
index int
}
func (*runtimeSettingsRows) Close() {}
func (*runtimeSettingsRows) Err() error { return nil }
func (rows *runtimeSettingsRows) Next() bool { return rows.index < len(rows.rows) }
func (rows *runtimeSettingsRows) Scan(dest ...any) error {
if rows.index >= len(rows.rows) || len(dest) != len(rows.rows[rows.index]) {
return errors.New("invalid runtime settings row")
}
for index, value := range rows.rows[rows.index] {
pointer, ok := dest[index].(*string)
if !ok {
return errors.New("unexpected runtime settings destination")
}
*pointer = value.(string)
}
rows.index++
return nil
}