synapsecns/sanguine

package db_test

import (
    "database/sql"
    "fmt"
    "os"
    "sync"
    "sync/atomic"
    "testing"
    "time"

    "github.com/synapsecns/sanguine/core"
    "github.com/synapsecns/sanguine/core/dbcommon"
    "github.com/synapsecns/sanguine/core/metrics"
    "github.com/synapsecns/sanguine/core/metrics/localmetrics"
    "github.com/synapsecns/sanguine/core/testsuite"
    "github.com/synapsecns/sanguine/services/scribe/db"
    "github.com/synapsecns/sanguine/services/scribe/metadata"

    "github.com/Flaque/filet"
    . "github.com/stretchr/testify/assert"
    "github.com/stretchr/testify/suite"
    "github.com/synapsecns/sanguine/services/scribe/db/datastore/sql/mysql"
    "github.com/synapsecns/sanguine/services/scribe/db/datastore/sql/sqlite"
    "gorm.io/gorm/schema"
)

type DBSuite struct {
    *testsuite.TestSuite
    dbs           []db.EventDB
    logIndex      atomic.Int64
    scribeMetrics metrics.Handler
}

// NewEventDBSuite creates a new EventDBSuite.
func NewEventDBSuite(tb testing.TB) *DBSuite {
    tb.Helper()
    return &DBSuite{
        TestSuite: testsuite.NewTestSuite(tb),
        dbs:       []db.EventDB{},
    }
}

func (t *DBSuite) SetupTest() {
    t.TestSuite.SetupTest()
    t.logIndex.Store(0)

    sqliteStore, err := sqlite.NewSqliteStore(t.GetTestContext(), filet.TmpDir(t.T(), ""), t.scribeMetrics, false)
    Nil(t.T(), err)

    t.dbs = []db.EventDB{sqliteStore}
    t.setupMysqlDB()
}

func (t *DBSuite) SetupSuite() {
    t.TestSuite.SetupSuite()

    // don't use metrics on ci for integration tests
    isCI := core.GetEnvBool("CI", false)
    useMetrics := !isCI
    metricsHandler := metrics.Null

    if useMetrics {
        localmetrics.SetupTestJaeger(t.GetSuiteContext(), t.T())
        metricsHandler = metrics.Jaeger
    }

    var err error
    t.scribeMetrics, err = metrics.NewByType(t.GetSuiteContext(), metadata.BuildInfo(), metricsHandler)
    t.Require().Nil(err)
}

func (t *DBSuite) setupMysqlDB() {
    // skip if mysql test disabled, this really only needs to be run in ci
    // skip if mysql test disabled
    if os.Getenv(dbcommon.EnableMysqlTestVar) == "" {
        return
    }
    // sets up the conn string to the default database
    connString := dbcommon.GetTestConnString()
    // sets up the myqsl db
    testDB, err := sql.Open("mysql", connString)
    Nil(t.T(), err)
    // close the db once the ocnnection is odne
    defer func() {
        Nil(t.T(), testDB.Close())
    }()

    // override the naming strategy to prevent tests from messing with each other.
    // todo this should be solved via a proper teardown process or transactions.
    mysql.NamingStrategy = schema.NamingStrategy{
        TablePrefix: fmt.Sprintf("test%d_%d_", t.GetTestID(), time.Now().Unix()),
    }

    mysql.MaxIdleConns = 10
    mysql.MaxOpenConns = 10

    // create the sql store
    mysqlStore, err := mysql.NewMysqlStore(t.GetTestContext(), connString, t.scribeMetrics, false)
    Nil(t.T(), err)
    // add the db
    t.dbs = append(t.dbs, mysqlStore)
}

func (t *DBSuite) RunOnAllDBs(testFunc func(testDB db.EventDB)) {
    t.T().Helper()

    wg := sync.WaitGroup{}
    for _, testDB := range t.dbs {
        wg.Add(1)
        // capture the value
        go func(testDB db.EventDB) {
            defer wg.Done()
            testFunc(testDB)
        }(testDB)
    }
    wg.Wait()
}

// TestDBSuite tests the db suite.
func TestEventDBSuite(t *testing.T) {
    suite.Run(t, NewEventDBSuite(t))
}