started moving pg-specific nmap tests to integration_tests/management/network_map_db/pgsql

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
This commit is contained in:
Dmitri Dolguikh
2026-07-31 12:29:35 +02:00
parent 2baeb4bc0d
commit 3b5c8e2298
15 changed files with 168 additions and 14 deletions

View File

@@ -0,0 +1,39 @@
package networkmap_pgsql
import (
"context"
"testing"
networkmap_pgsql "github.com/netbirdio/netbird/management/internals/network_map_db/pgsql"
"github.com/netbirdio/netbird/shared/management/networkmap/nmdata"
"github.com/stretchr/testify/assert"
)
func TestXXX(t *testing.T) {
ctx := context.TODO()
s, err := networkmap_pgsql.NewPostgresqlStore(ctx, dsn)
assert.NoError(t, err)
_, err = s.Pool.Query(ctx,
"insert into accounts (id) VALUES('account-id-1')")
assert.NoError(t, err)
_, err = s.Pool.Query(ctx,
"insert into groups (id, account_id, name, resources, public_id) VALUES('test-group-id-1','account-id-1','test-group-1', '[{\"ID\":\"host-id-1\",\"Type\":\"host\"}]','public-id-1')")
assert.NoError(t, err)
_, err = s.Pool.Query(ctx,
"insert into groups (id, account_id, name, resources, public_id) VALUES('test-group-id-2','account-id-1','test-group-2', '[{\"ID\":\"subnet-id-1\",\"Type\":\"subnet\"}, {\"ID\":\"host-id-2\",\"Type\":\"host\"}]','public-id-2')")
assert.NoError(t, err)
groups, err := s.GetGroups(ctx, "account-id-1")
assert.NoError(t, err)
assert.Contains(t,
groups,
nmdata.Group{Name: "test-group-1", PublicID: "public-id-1", Resources: []nmdata.Resource{{ID: "host-id-1", Type: "host"}}},
)
assert.Contains(t,
groups,
nmdata.Group{Name: "test-group-2", PublicID: "public-id-2", Resources: []nmdata.Resource{{ID: "subnet-id-1", Type: "subnet"}, {ID: "host-id-2", Type: "host"}}},
)
}

View File

@@ -0,0 +1,115 @@
package networkmap_pgsql
import (
"context"
"fmt"
"os"
"regexp"
"strings"
"testing"
"time"
"github.com/google/uuid"
log "github.com/sirupsen/logrus"
"gorm.io/driver/postgres"
"gorm.io/gorm"
gormstore "github.com/netbirdio/netbird/management/server/store"
"github.com/netbirdio/netbird/management/server/testutil"
)
var dsn string
func TestMain(m *testing.M) {
_, tmpdsn, err := testutil.CreatePostgresTestContainer()
if err != nil {
log.Fatalf("error starting postres container %v", err)
}
var db *gorm.DB
for i := range 5 {
db, err = gorm.Open(postgres.Open(tmpdsn), &gorm.Config{})
if err == nil {
break
}
if i < 5 {
waitTime := time.Duration(100*(i+1)) * time.Millisecond
time.Sleep(waitTime)
continue
}
log.Fatalf("error connecting to postres db %v", err)
}
var cleanup func()
dsn, cleanup, err = createRandomDB(tmpdsn, db)
sqlDB, _ := db.DB()
if sqlDB != nil {
sqlDB.Close()
}
if err != nil {
log.Fatalf("error creating postres db %v", err)
}
_, err = gormstore.NewPostgresqlStoreForTests(context.TODO(), dsn, nil, false)
if err != nil {
log.Fatalf("error running migrations %v", err)
}
code := m.Run()
cleanup()
os.Exit(code)
}
func createRandomDB(dsn string, db *gorm.DB) (string, func(), error) {
dbName := fmt.Sprintf("test_db_%s", strings.ReplaceAll(uuid.New().String(), "-", "_"))
if err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbName)).Error; err != nil {
return "", nil, fmt.Errorf("failed to create database: %v", err)
}
originalDSN := dsn
cleanup := func() {
var dropDB *gorm.DB
var err error
dropDB, err = gorm.Open(postgres.Open(originalDSN), &gorm.Config{
SkipDefaultTransaction: true,
PrepareStmt: false,
})
if err != nil {
log.Errorf("failed to connect for dropping database %s: %v", dbName, err)
return
}
defer func() {
if sqlDB, _ := dropDB.DB(); sqlDB != nil {
sqlDB.Close()
}
}()
if sqlDB, _ := dropDB.DB(); sqlDB != nil {
sqlDB.SetMaxOpenConns(1)
sqlDB.SetMaxIdleConns(0)
sqlDB.SetConnMaxLifetime(time.Second)
}
err = dropDB.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s WITH (FORCE)", dbName)).Error
if err != nil {
log.Errorf("failed to drop database %s: %v", dbName, err)
}
}
return replaceDBName(dsn, dbName), cleanup, nil
}
func replaceDBName(dsn, newDBName string) string {
re := regexp.MustCompile(`(?P<pre>[:/@])(?P<dbname>[^/?]+)(?P<post>\?|$)`)
return re.ReplaceAllString(dsn, `${pre}`+newDBName+`${post}`)
}