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 <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-08-31 09:59:26 +00:00
1 parent b8316f2a2a
commit 115aeb47c0
2 files changed
+166 -34

No files matched your search

+33
View File
@@ -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"))
})
})
})
+133 -34
View File
@@ -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 "_<counter>".
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())