alby-hub/tests/db/test_db.go

package db

import (
	"fmt"
	"net/url"
	"os"
	"strings"
	"testing"

	"github.com/peterldowns/pgtestdb"
	"gorm.io/gorm"

	"github.com/getAlby/hub/db"
	"github.com/getAlby/hub/logger"
)

const defaultTestDB = "test.db"

func GetTestDatabaseURI() string {
	ret := os.Getenv("TEST_DATABASE_URI")
	if ret == "" {
		// TODO: use in-memory DB, or a temporary file
		return defaultTestDB
	}

	return ret
}

func NewDB(t *testing.T) (*gorm.DB, error) {
	dbUri := GetTestDatabaseURI()

	if dbUri == defaultTestDB {
		//in case the file was not removed in the last run, remove it before starting the test
		logger.Logger.WithField("uri", defaultTestDB).Info("removing test db")
		os.Remove(defaultTestDB)
	}

	logger.Logger.WithField("uri", dbUri).Info("Creating new test DB with URI")
	return NewDBWithURI(t, dbUri)
}

func NewDBWithURI(t *testing.T, uri string) (*gorm.DB, error) {
	if db.IsPostgresURI(uri) {
		parsedURI, err := url.Parse(uri)
		if err != nil {
			return nil, fmt.Errorf("failed to parse postgres DB URI: %w", err)
		}

		var user, password string
		if userInfo := parsedURI.User; userInfo != nil {
			user = userInfo.Username()
			password, _ = userInfo.Password()
		}

		dbName := strings.TrimPrefix(parsedURI.Path, "/")

		config := pgtestdb.Custom(t, pgtestdb.Config{
			DriverName: "pgx",
			Host:       parsedURI.Hostname(),
			Port:       parsedURI.Port(),
			User:       user,
			Password:   password,
			Database:   dbName,
		}, pgtestdb.NoopMigrator{})

		uri = config.URL()
	}

	return db.NewDBWithConfig(&db.Config{
		URI:        uri,
		LogQueries: true,
	})
}

func CloseDB(d *gorm.DB) {
	if err := db.Stop(d); err != nil {
		panic("failed to close database: " + err.Error())
	}

	if GetTestDatabaseURI() == defaultTestDB {
		logger.Logger.WithField("uri", defaultTestDB).Info("removing test db")
		os.Remove(defaultTestDB)
	}
}