Browse Source

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.
n0ctal 8 hours ago
parent
commit
66df77665f
2 changed files with 107 additions and 1 deletions
  1. 95 0
      internal/testpg/isolate.go
  2. 12 1
      internal/web/service/xray_config_inject_test.go

+ 95 - 0
internal/testpg/isolate.go

@@ -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()
+}

+ 12 - 1
internal/web/service/xray_config_inject_test.go

@@ -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) {