mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-18 02:31:10 -04:00
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:
1 parent
b8316f2a2a
commit
115aeb47c0
2 files changed
+166
-34
No files matched your search
@@ -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"))
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -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())
|
||||
|
||||
Reference in new issue
Block a user