diff --git a/tests/e2e/distributed/dbname_test.go b/tests/e2e/distributed/dbname_test.go new file mode 100644 index 000000000..76f103074 --- /dev/null +++ b/tests/e2e/distributed/dbname_test.go @@ -0,0 +1,33 @@ +package distributed_test + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Test database naming", Label("Distributed"), func() { + Describe("sanitizeDBName", func() { + It("lowercases and replaces characters Postgres will not accept unquoted", func() { + Expect(sanitizeDBName("LocalAI-Test.Suite")).To(Equal("localai_test_suite")) + }) + + It("truncates to fit the 63-byte identifier limit with room for a suffix", func() { + long := "" + for i := 0; i < 100; i++ { + long += "a" + } + Expect(len(sanitizeDBName(long))).To(BeNumerically("<=", 50)) + }) + + It("never produces an empty name", func() { + Expect(sanitizeDBName("---")).ToNot(BeEmpty()) + }) + }) + + Describe("replaceDBName", func() { + It("swaps the database in a testcontainers DSN and keeps the query string", func() { + dsn := "postgres://test:test@127.0.0.1:32768/localai_suite?sslmode=disable" + Expect(replaceDBName(dsn, "spec_7")).To(Equal("postgres://test:test@127.0.0.1:32768/spec_7?sslmode=disable")) + }) + }) +}) diff --git a/tests/e2e/distributed/testhelpers_test.go b/tests/e2e/distributed/testhelpers_test.go index 68cf537e3..564e8c406 100644 --- a/tests/e2e/distributed/testhelpers_test.go +++ b/tests/e2e/distributed/testhelpers_test.go @@ -2,6 +2,10 @@ package distributed_test import ( "context" + "fmt" + "net/url" + "strings" + "sync/atomic" "time" "github.com/mudler/LocalAI/core/services/messaging" @@ -13,6 +17,9 @@ import ( tcnats "github.com/testcontainers/testcontainers-go/modules/nats" tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres" "github.com/testcontainers/testcontainers-go/wait" + "gorm.io/driver/postgres" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" ) // TestInfra holds shared test containers and connection strings. @@ -25,71 +32,167 @@ type TestInfra struct { NC *messaging.Client } -// SetupInfra starts PostgreSQL and NATS containers and connects a messaging client. -// Call in BeforeEach. Use DeferCleanup or call Teardown in AfterEach. -func SetupInfra(dbName string) *TestInfra { - GinkgoHelper() +// Containers are suite-scoped, not spec-scoped. Starting a Postgres (~10s) and a +// NATS (~3.5s) per spec cost roughly 48 minutes of pure startup across the 213 +// specs behind SetupInfra, which is why this suite was never wired into CI. +// Isolation now comes from a database per spec (~67ms), which is what the dbName +// argument was always describing. +// +// Plain BeforeSuite rather than SynchronizedBeforeSuite is deliberate: under +// `ginkgo -p` each process gets its own container pair, which keeps NATS subjects +// isolated per process. A single shared NATS across parallel processes would let +// specs on different processes see each other's messages on the same subject. +var ( + suitePG *tcpostgres.PostgresContainer + suiteNATS *tcnats.NATSContainer + suitePGDSN string + suiteNatsURL string + dbCounter atomic.Int64 +) - infra := &TestInfra{Ctx: context.Background()} +var _ = BeforeSuite(func() { + ctx := context.Background() var err error - // Start PostgreSQL container - infra.PGContainer, err = tcpostgres.Run(infra.Ctx, "postgres:16-alpine", - tcpostgres.WithDatabase(dbName), + suitePG, err = tcpostgres.Run(ctx, "postgres:16-alpine", + tcpostgres.WithDatabase("localai_suite"), tcpostgres.WithUsername("test"), tcpostgres.WithPassword("test"), testcontainers.WithWaitStrategy( wait.ForLog("database system is ready to accept connections"). WithOccurrence(2). - WithStartupTimeout(30*time.Second), + WithStartupTimeout(90*time.Second), ), ) Expect(err).ToNot(HaveOccurred()) - infra.PGURL, err = infra.PGContainer.ConnectionString(infra.Ctx, "sslmode=disable") + suitePGDSN, err = suitePG.ConnectionString(ctx, "sslmode=disable") Expect(err).ToNot(HaveOccurred()) - // Start NATS container - infra.NATSContainer, err = tcnats.Run(infra.Ctx, "nats:2-alpine") + suiteNATS, err = tcnats.Run(ctx, "nats:2-alpine") Expect(err).ToNot(HaveOccurred()) - infra.NatsURL, err = infra.NATSContainer.ConnectionString(infra.Ctx) + suiteNatsURL, err = suiteNATS.ConnectionString(ctx) Expect(err).ToNot(HaveOccurred()) +}) - // Connect messaging client +var _ = AfterSuite(func() { + ctx := context.Background() + if suitePG != nil { + _ = suitePG.Terminate(ctx) + } + if suiteNATS != nil { + _ = suiteNATS.Terminate(ctx) + } +}) + +// sanitizeDBName maps a spec-supplied label onto a legal unquoted Postgres +// identifier, leaving headroom for the uniqueness suffix appended by SetupInfra. +func sanitizeDBName(name string) string { + var b strings.Builder + for _, r := range strings.ToLower(name) { + switch { + case r >= 'a' && r <= 'z', r >= '0' && r <= '9', r == '_': + b.WriteRune(r) + default: + b.WriteRune('_') + } + } + out := strings.Trim(b.String(), "_") + if out == "" { + out = "spec" + } + // Postgres identifiers cap at 63 bytes; reserve the rest for "_". + if len(out) > 50 { + out = out[:50] + } + return out +} + +// replaceDBName swaps the database component of a DSN, preserving credentials, +// host, port and query parameters. +func replaceDBName(dsn, name string) string { + GinkgoHelper() + u, err := url.Parse(dsn) + Expect(err).ToNot(HaveOccurred()) + u.Path = "/" + name + return u.String() +} + +// adminDB opens a short-lived connection to the suite's maintenance database. +// CREATE/DROP DATABASE cannot run inside a transaction or against the target +// database itself, so every call gets its own connection and closes it. +func adminDB() *gorm.DB { + GinkgoHelper() + db, err := gorm.Open(postgres.Open(suitePGDSN), &gorm.Config{Logger: gormlogger.Discard}) + Expect(err).ToNot(HaveOccurred()) + return db +} + +func closeDB(db *gorm.DB) { + if db == nil { + return + } + if sqlDB, err := db.DB(); err == nil { + _ = sqlDB.Close() + } +} + +// SetupInfra provisions a dedicated database on the suite-scoped Postgres and +// returns a client connected to the suite-scoped NATS. Call in BeforeEach; +// cleanup is registered with DeferCleanup. +func SetupInfra(dbName string) *TestInfra { + GinkgoHelper() + Expect(suitePG).ToNot(BeNil(), "SetupInfra called before BeforeSuite started the shared containers") + + infra := &TestInfra{ + Ctx: context.Background(), + PGContainer: suitePG, + NATSContainer: suiteNATS, + NatsURL: suiteNatsURL, + } + + db := fmt.Sprintf("%s_%d", sanitizeDBName(dbName), dbCounter.Add(1)) + + admin := adminDB() + Expect(admin.Exec(fmt.Sprintf("CREATE DATABASE %q", db)).Error).To(Succeed()) + closeDB(admin) + + infra.PGURL = replaceDBName(suitePGDSN, db) + + var err error infra.NC, err = messaging.New(infra.NatsURL) Expect(err).ToNot(HaveOccurred()) - // Register cleanup in LIFO order DeferCleanup(func() { if infra.NC != nil { infra.NC.Close() } - if infra.PGContainer != nil { - infra.PGContainer.Terminate(context.Background()) - } - if infra.NATSContainer != nil { - infra.NATSContainer.Terminate(context.Background()) + // FORCE terminates any connection the spec left open (Postgres 13+). + // Failure to drop must not fail the spec: the container dies at AfterSuite. + drop := adminDB() + if err := drop.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %q WITH (FORCE)", db)).Error; err != nil { + AddReportEntry("drop database failed", fmt.Sprintf("%s: %v", db, err)) } + closeDB(drop) }) return infra } -// SetupNATSOnly starts only a NATS container and connects a messaging client. -// Useful for tests that don't need PostgreSQL. +// SetupNATSOnly returns a client on the suite-scoped NATS for specs that need no +// database. func SetupNATSOnly() *TestInfra { GinkgoHelper() + Expect(suiteNATS).ToNot(BeNil(), "SetupNATSOnly called before BeforeSuite started the shared containers") + + infra := &TestInfra{ + Ctx: context.Background(), + NATSContainer: suiteNATS, + NatsURL: suiteNatsURL, + } - infra := &TestInfra{Ctx: context.Background()} var err error - - infra.NATSContainer, err = tcnats.Run(infra.Ctx, "nats:2-alpine") - Expect(err).ToNot(HaveOccurred()) - - infra.NatsURL, err = infra.NATSContainer.ConnectionString(infra.Ctx) - Expect(err).ToNot(HaveOccurred()) - infra.NC, err = messaging.New(infra.NatsURL) Expect(err).ToNot(HaveOccurred()) @@ -97,16 +200,12 @@ func SetupNATSOnly() *TestInfra { if infra.NC != nil { infra.NC.Close() } - if infra.NATSContainer != nil { - infra.NATSContainer.Terminate(context.Background()) - } }) return infra } // FlushNATS ensures all subscriptions are registered server-side before publishing. -// Replaces time.Sleep(100ms) after Subscribe calls. func FlushNATS(nc *messaging.Client) { GinkgoHelper() Expect(nc.Conn().Flush()).To(Succeed())