mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-09-27 12:16:41 +08:00
test(database): give each package its own schema when tests run on PostgreSQL (#6594)
With XUI_DB_TYPE=postgres every test package shared one database and worked in public. Go runs package test binaries concurrently, so migrations raced and rows a previous run left behind leaked into the next. testpg.IsolatePackage creates a schema for the calling package, puts it first on search_path and drops it when the package finishes. It returns at once unless XUI_DB_TYPE is postgres. internal/web/service's TestMain adopts it.
This commit is contained in:
@@ -0,0 +1,95 @@
|
||||
package testpg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const (
|
||||
dbTypeEnv = "XUI_DB_TYPE"
|
||||
dbDSNEnv = "XUI_DB_DSN"
|
||||
)
|
||||
|
||||
// IsolatePackage gives one test package its own PostgreSQL schema: package test
|
||||
// binaries run concurrently, and sharing public lets their migrations race.
|
||||
func IsolatePackage(packageName string) (func(), error) {
|
||||
if os.Getenv(dbTypeEnv) != "postgres" {
|
||||
return func() {}, nil
|
||||
}
|
||||
baseDSN := strings.TrimSpace(os.Getenv(dbDSNEnv))
|
||||
if baseDSN == "" {
|
||||
return func() {}, nil
|
||||
}
|
||||
|
||||
suffix := make([]byte, 8)
|
||||
if _, err := rand.Read(suffix); err != nil {
|
||||
return nil, fmt.Errorf("generate PostgreSQL test schema suffix: %w", err)
|
||||
}
|
||||
schema := fmt.Sprintf("xui_%s_%d_%s", sanitize(packageName), os.Getpid(), hex.EncodeToString(suffix))
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
admin, err := pgxpool.New(ctx, baseDSN)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open PostgreSQL test database: %w", err)
|
||||
}
|
||||
if _, err := admin.Exec(ctx, "CREATE SCHEMA "+pgx.Identifier{schema}.Sanitize()); err != nil {
|
||||
admin.Close()
|
||||
return nil, fmt.Errorf("create PostgreSQL test schema: %w", err)
|
||||
}
|
||||
isolatedDSN, err := withSearchPath(baseDSN, schema)
|
||||
if err != nil {
|
||||
admin.Close()
|
||||
return nil, err
|
||||
}
|
||||
if err := os.Setenv(dbDSNEnv, isolatedDSN); err != nil {
|
||||
admin.Close()
|
||||
return nil, fmt.Errorf("set isolated PostgreSQL test DSN: %w", err)
|
||||
}
|
||||
|
||||
return func() {
|
||||
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cleanupCancel()
|
||||
_, _ = admin.Exec(cleanupCtx, "DROP SCHEMA "+pgx.Identifier{schema}.Sanitize()+" CASCADE")
|
||||
admin.Close()
|
||||
_ = os.Setenv(dbDSNEnv, baseDSN)
|
||||
}, nil
|
||||
}
|
||||
|
||||
func withSearchPath(dsn, schema string) (string, error) {
|
||||
u, err := url.Parse(dsn)
|
||||
if err == nil && (u.Scheme == "postgres" || u.Scheme == "postgresql") {
|
||||
query := u.Query()
|
||||
query.Set("search_path", schema)
|
||||
u.RawQuery = query.Encode()
|
||||
return u.String(), nil
|
||||
}
|
||||
if strings.ContainsAny(schema, " '[]=\\") {
|
||||
return "", fmt.Errorf("unsafe PostgreSQL test schema name")
|
||||
}
|
||||
return strings.TrimSpace(dsn) + " search_path=" + schema, nil
|
||||
}
|
||||
|
||||
func sanitize(value string) string {
|
||||
var result strings.Builder
|
||||
for _, r := range strings.ToLower(value) {
|
||||
if r >= 'a' && r <= 'z' || r >= '0' && r <= '9' || r == '_' {
|
||||
result.WriteRune(r)
|
||||
} else {
|
||||
result.WriteByte('_')
|
||||
}
|
||||
}
|
||||
if result.Len() == 0 {
|
||||
return "pkg"
|
||||
}
|
||||
return result.String()
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
@@ -10,6 +11,7 @@ import (
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/amneziawgnet"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/database/model"
|
||||
xuilogger "github.com/mhsanaei/3x-ui/v3/internal/logger"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/testpg"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/util/json_util"
|
||||
"github.com/mhsanaei/3x-ui/v3/internal/xray"
|
||||
|
||||
@@ -25,7 +27,16 @@ func TestMain(m *testing.M) {
|
||||
// injectPanelEgress logs when it skips injection; the package logger must
|
||||
// exist before any test exercises a skipped path.
|
||||
xuilogger.InitLogger(logging.ERROR)
|
||||
os.Exit(m.Run())
|
||||
// Against PostgreSQL every package shares one database; give this one its
|
||||
// own schema so a parallel package and a previous run cannot reach it.
|
||||
cleanup, err := testpg.IsolatePackage("internal_web_service")
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, err)
|
||||
os.Exit(1)
|
||||
}
|
||||
code := m.Run()
|
||||
cleanup()
|
||||
os.Exit(code)
|
||||
}
|
||||
|
||||
func TestEnsureAPIServices(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user