From 115aeb47c0f9ea4ca383c149059e68f10e6cf4a9 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 09:59:26 +0000 Subject: [PATCH] test(distributed): share suite containers, isolate specs by database Starting a Postgres and a NATS container per spec cost roughly 48 minutes of startup across the 213 specs behind SetupInfra, which is why this suite was never wired into CI. Containers move to BeforeSuite and isolation comes from CREATE DATABASE, which the dbName argument already described. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- tests/e2e/distributed/dbname_test.go | 33 +++++ tests/e2e/distributed/testhelpers_test.go | 167 +++++++++++++++++----- 2 files changed, 166 insertions(+), 34 deletions(-) create mode 100644 tests/e2e/distributed/dbname_test.go 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())