Use TestContainers postgress module and snapshots (#103)

go.mod

@@ -112,6 +112,7 @@ require (
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/sosodev/duration v1.3.0 // indirect
 	github.com/spf13/cast v1.5.0 // indirect
+	github.com/testcontainers/testcontainers-go/modules/postgres v0.31.0
 	github.com/tklauser/go-sysconf v0.3.12 // indirect
 	github.com/tklauser/numcpus v0.6.1 // indirect
 	github.com/urfave/cli/v2 v2.27.2 // indirect

go.sum

@@ -651,6 +651,8 @@ github.com/subosito/gotenv v1.3.0/go.mod h1:YzJjq/33h7nrwdY+iHMhEOEEbW0ovIz0tB6t
 github.com/subosito/gotenv v1.4.1/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0=
 github.com/testcontainers/testcontainers-go v0.31.0 h1:W0VwIhcEVhRflwL9as3dhY6jXjVCA27AkmbnZ+UTh3U=
 github.com/testcontainers/testcontainers-go v0.31.0/go.mod h1:D2lAoA0zUFiSY+eAflqK5mcUx/A5hrrORaEQrd0SefI=
+github.com/testcontainers/testcontainers-go/modules/postgres v0.31.0 h1:isAwFS3KNKRbJMbWv+wolWqOFUECmjYZ+sIRZCIBc/E=
+github.com/testcontainers/testcontainers-go/modules/postgres v0.31.0/go.mod h1:ZNYY8vumNCEG9YI59A9d6/YaMY49uwRhmeU563EzFGw=
 github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=
 github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=
 github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=

graph/rewards_test.go

@@ -20,7 +20,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/suite"
-	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
 	"github.com/volatiletech/null/v8"
 	"github.com/volatiletech/sqlboiler/v4/boil"
 	"github.com/volatiletech/sqlboiler/v4/types"
@@ -30,7 +30,7 @@ type RewardsQueryTestSuite struct {
 	suite.Suite
 	ctx       context.Context
 	pdb       db.Store
-	container testcontainers.Container
+	container *postgres.PostgresContainer
 	settings  config.Settings
 	resolver  *Resolver
 	repo      *base.Repository
@@ -51,7 +51,7 @@ func (r *RewardsQueryTestSuite) SetupSuite() {
 
 // TearDownTest after each test truncate tables
 func (r *RewardsQueryTestSuite) TearDownTest() {
-	test.TruncateTables(r.pdb.DBS().Writer.DB, r.T())
+	r.Require().NoError(r.container.Restore(r.ctx))
 }
 
 // TearDownSuite cleanup at end by terminating container

internal/helpers/test.go

@@ -3,171 +3,89 @@ package helpers
 import (
 	"context"
 	"crypto/ecdsa"
-	"database/sql"
 	"fmt"
-	"net/http"
-	"os"
-	"strings"
 	"testing"
 	"time"
 
 	"github.com/DIMO-Network/shared/db"
 	"github.com/docker/go-connections/nat"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
-	"github.com/gofiber/fiber/v2"
-	"github.com/golang-jwt/jwt/v5"
-	"github.com/pkg/errors"
 	"github.com/pressly/goose/v3"
-	"github.com/rs/zerolog"
+	"github.com/stretchr/testify/require"
 	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
 	"github.com/testcontainers/testcontainers-go/wait"
 )
 
 var DBSettings = db.Settings{
-	Name:               "identity_api",
-	Host:               "localhost",
-	Port:               "6669",
-	User:               "postgres",
-	Password:           "postgres",
-	MaxOpenConnections: 2,
-	MaxIdleConnections: 2,
+	Name:     "identity_api",
+	Host:     "localhost",
+	User:     "dimo",
+	Password: "dimo",
+	// MaxOpenConnections: 2,
+	// MaxIdleConnections: 2,
 }
 
 // StartContainerDatabase starts postgres container with default test settings, and migrates the db. Caller must terminate container.
-func StartContainerDatabase(ctx context.Context, t *testing.T, migrationsDirRelPath string) (db.Store, testcontainers.Container) {
-	settings := struct {
-		LogLevel    string
-		DB          db.Settings
-		ServiceName string
-	}{
-		LogLevel:    "info",
-		DB:          DBSettings,
-		ServiceName: DBSettings.Name,
-	}
-
+func StartContainerDatabase(ctx context.Context, t *testing.T, migrationsDirRelPath string) (dbs db.Store, container *postgres.PostgresContainer) {
+	settings := DBSettings // Copy.
 	pgPort := "5432/tcp"
-	dbURL := func(host string, port nat.Port) string {
-		return fmt.Sprintf("postgres://%s:%s@localhost:%s/%s?sslmode=disable", settings.DB.User, settings.DB.Password, port.Port(), settings.DB.Name)
-	}
-	cr := testcontainers.ContainerRequest{
-		Image:        "postgres:12.9-alpine",
-		Env:          map[string]string{"POSTGRES_USER": settings.DB.User, "POSTGRES_PASSWORD": settings.DB.Password, "POSTGRES_DB": settings.DB.Name},
-		ExposedPorts: []string{pgPort},
-		Cmd:          []string{"postgres", "-c", "fsync=off"},
-		WaitingFor:   wait.ForSQL(nat.Port(pgPort), "postgres", dbURL),
-	}
 
-	pgContainer, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
-		ContainerRequest: cr,
-		Started:          true,
-	})
-	if err != nil {
-		return handleContainerStartErr(ctx, err, pgContainer, t)
-	}
-	mappedPort, err := pgContainer.MappedPort(ctx, nat.Port(pgPort))
+	var err error
+
+	container, err = postgres.RunContainer(
+		ctx,
+		testcontainers.WithImage("docker.io/postgres:14.10-alpine"),
+		postgres.WithDatabase("identity_api"),
+		postgres.WithUsername("dimo"),
+		postgres.WithPassword("dimo"),
+		testcontainers.WithWaitStrategy(
+			wait.ForLog("database system is ready to accept connections").
+				WithOccurrence(2).
+				WithStartupTimeout(5*time.Second),
+		),
+	)
+	require.NoError(t, err)
+
+	defer func() {
+		if err != nil {
+			container.Terminate(ctx) //nolint:errcheck
+			t.Fatalf("Failed to set up Postgres container: %v", err)
+		}
+	}()
+
+	mappedPort, err := container.MappedPort(ctx, nat.Port(pgPort))
 	if err != nil {
-		return handleContainerStartErr(ctx, errors.Wrap(err, "failed to get container external port"), pgContainer, t)
+		err = fmt.Errorf("couldn't get Postgres port: %w", err)
+		return
 	}
-	fmt.Printf("postgres container session %s ready and running at port: %s \n", pgContainer.SessionID(), mappedPort)
-	//defer pgContainer.Terminate(ctx) // this should be done by the caller
 
-	settings.DB.Port = mappedPort.Port()
-	pdb := db.NewDbConnectionForTest(ctx, &settings.DB, false)
-	for !pdb.IsReady() {
+	settings.Port = mappedPort.Port()
+	dbs = db.NewDbConnectionForTest(ctx, &settings, false)
+	for !dbs.IsReady() {
 		time.Sleep(500 * time.Millisecond)
 	}
-	// can't connect to db, dsn=user=postgres password=postgres dbname=identity_api host=localhost port=49395 sslmode=disable search_path=identity_api, err=EOF
-	// error happens when calling here
-	_, err = pdb.DBS().Writer.Exec(`
-		grant usage on schema public to public;
-		grant create on schema public to public;
-		CREATE SCHEMA IF NOT EXISTS identity_api;
-		ALTER USER postgres SET search_path = identity_api, public;
-		SET search_path = identity_api, public;
-		`)
-	if err != nil {
-		return handleContainerStartErr(ctx, errors.Wrapf(err, "failed to apply schema. session: %s, port: %s",
-			pgContainer.SessionID(), mappedPort.Port()), pgContainer, t)
-	}
-	// add truncate tables func
-	_, err = pdb.DBS().Writer.Exec(`
-CREATE OR REPLACE FUNCTION truncate_tables() RETURNS void AS $$
-DECLARE
-    statements CURSOR FOR
-        SELECT tablename FROM pg_tables
-        WHERE schemaname = 'identity_api' and tablename != 'migrations';
-BEGIN
-    FOR stmt IN statements LOOP
-        EXECUTE 'TRUNCATE TABLE ' || quote_ident(stmt.tablename) || ' CASCADE;';
-    END LOOP;
-END;
-$$ LANGUAGE plpgsql;
-`)
+
+	_, err = dbs.DBS().Writer.Exec("CREATE SCHEMA IF NOT EXISTS identity_api")
 	if err != nil {
-		return handleContainerStartErr(ctx, errors.Wrap(err, "failed to create truncate func"), pgContainer, t)
+		err = fmt.Errorf("error creating schema: %w", err)
+		return
 	}
 
 	goose.SetTableName("identity_api.migrations")
-	if err := goose.RunContext(context.Background(), "up", pdb.DBS().Writer.DB, migrationsDirRelPath); err != nil {
-		return handleContainerStartErr(ctx, errors.Wrap(err, "failed to apply goose migrations for test"), pgContainer, t)
-	}
-
-	return pdb, pgContainer
-}
-
-func handleContainerStartErr(ctx context.Context, err error, container testcontainers.Container, t *testing.T) (db.Store, testcontainers.Container) {
+	err = goose.RunContext(ctx, "up", dbs.DBS().Writer.DB, migrationsDirRelPath)
 	if err != nil {
-		fmt.Println("start container error: " + err.Error())
-		if container != nil {
-			container.Terminate(ctx) //nolint
-		}
-		t.Fatal(err)
-	}
-	return db.Store{}, container
-}
-
-func BuildRequest(method, url, body string) *http.Request {
-	req, _ := http.NewRequest(
-		method,
-		url,
-		strings.NewReader(body),
-	)
-	req.Header.Set("Content-Type", "application/json")
-
-	return req
-}
-
-// AuthInjectorTestHandler injects fake jwt with sub
-func AuthInjectorTestHandler(userID string) fiber.Handler {
-	return func(c *fiber.Ctx) error {
-		token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
-			"sub": userID,
-			"nbf": time.Now().Unix(),
-		})
-
-		c.Locals("user", token)
-		return c.Next()
+		err = fmt.Errorf("error running migrations: %w", err)
+		return
 	}
-}
 
-// TruncateTables truncates tables for the test db, useful to run as teardown at end of each DB dependent test.
-func TruncateTables(db *sql.DB, t *testing.T) {
-	_, err := db.Exec(`SELECT truncate_tables();`)
+	err = container.Snapshot(ctx)
 	if err != nil {
-		fmt.Println("truncating tables failed.")
-		t.Fatal(err)
+		err = fmt.Errorf("error taking blank snapshot: %w", err)
 	}
-}
-
-/** Test Setup functions. At some point may want to move elsewhere more generic **/
 
-func Logger() *zerolog.Logger {
-	l := zerolog.New(os.Stdout).With().
-		Timestamp().
-		Str("app", DBSettings.Name).
-		Logger()
-	return &l
+	return
 }
 
 func GenerateWallet() (*ecdsa.PrivateKey, *common.Address, error) {

internal/repositories/dcn/dcn_test.go

@@ -17,7 +17,7 @@ import (
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/suite"
-	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
 	"github.com/volatiletech/null/v8"
 	"github.com/volatiletech/sqlboiler/v4/boil"
 )
@@ -26,7 +26,7 @@ type DCNRepoTestSuite struct {
 	suite.Suite
 	ctx       context.Context
 	pdb       db.Store
-	container testcontainers.Container
+	container *postgres.PostgresContainer
 	repo      *Repository
 	settings  config.Settings
 }
@@ -45,7 +45,7 @@ func (o *DCNRepoTestSuite) SetupSuite() {
 
 // TearDownTest after each test truncate tables
 func (s *DCNRepoTestSuite) TearDownTest() {
-	test.TruncateTables(s.pdb.DBS().Writer.DB, s.T())
+	s.Require().NoError(s.container.Restore(s.ctx))
 }
 
 // TearDownSuite cleanup at end by terminating container

internal/repositories/reward/rewards_test.go

@@ -18,7 +18,7 @@ import (
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/suite"
-	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
 	"github.com/volatiletech/null/v8"
 	"github.com/volatiletech/sqlboiler/v4/boil"
 	"github.com/volatiletech/sqlboiler/v4/queries/qm"
@@ -29,7 +29,7 @@ type RewardsRepoTestSuite struct {
 	suite.Suite
 	ctx              context.Context
 	pdb              db.Store
-	container        testcontainers.Container
+	container        *postgres.PostgresContainer
 	repo             *Repository
 	settings         config.Settings
 	paginationHelper helpers.PaginationHelper[RewardsCursor]
@@ -57,7 +57,7 @@ func (r *RewardsRepoTestSuite) SetupSuite() {
 
 // TearDownTest after each test truncate tables
 func (r *RewardsRepoTestSuite) TearDownTest() {
-	helpers.TruncateTables(r.pdb.DBS().Writer.DB, r.T())
+	r.Require().NoError(r.container.Restore(r.ctx))
 }
 
 // TearDownSuite cleanup at end by terminating container

internal/repositories/synthetic/synthetic_devices_test.go

@@ -18,7 +18,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
-	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
 	"github.com/volatiletech/null/v8"
 	"github.com/volatiletech/sqlboiler/v4/boil"
 )
@@ -65,7 +65,7 @@ func Test_SyntheticDeviceToAPI(t *testing.T) {
 type SyntheticTestSuite struct {
 	suite.Suite
 	pdb       db.Store
-	container testcontainers.Container
+	container *postgres.PostgresContainer
 	repo      *Repository
 	settings  config.Settings
 
@@ -170,7 +170,7 @@ func (s *SyntheticTestSuite) SetupSuite() {
 
 // TearDownTest after each test truncate tables.
 func (s *SyntheticTestSuite) TearDownTest() {
-	helpers.TruncateTables(s.pdb.DBS().Writer.DB, s.T())
+	s.Require().NoError(s.container.Restore(context.TODO()))
 }
 
 // TearDownSuite cleanup at end by terminating container.

internal/repositories/vehicle/owned_vehicles_test.go

@@ -18,7 +18,7 @@ import (
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/suite"
-	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
 	"github.com/volatiletech/null/v8"
 	"github.com/volatiletech/sqlboiler/v4/boil"
 )
@@ -27,7 +27,7 @@ type OwnedVehiclesRepoTestSuite struct {
 	suite.Suite
 	ctx       context.Context
 	pdb       db.Store
-	container testcontainers.Container
+	container *postgres.PostgresContainer
 	repo      *Repository
 	settings  config.Settings
 }
@@ -50,7 +50,7 @@ func (s *OwnedVehiclesRepoTestSuite) SetupSuite() {
 
 // TearDownTest after each test truncate tables
 func (s *OwnedVehiclesRepoTestSuite) TearDownTest() {
-	helpers.TruncateTables(s.pdb.DBS().Writer.DB, s.T())
+	s.Require().NoError(s.container.Restore(s.ctx))
 }
 
 // TearDownSuite cleanup at end by terminating container

internal/repositories/vehicle/vehicles_test.go

@@ -20,7 +20,7 @@ import (
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
-	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
 	"github.com/volatiletech/null/v8"
 	"github.com/volatiletech/sqlboiler/v4/boil"
 )
@@ -44,7 +44,7 @@ type AccessibleVehiclesRepoTestSuite struct {
 	suite.Suite
 	ctx       context.Context
 	pdb       db.Store
-	container testcontainers.Container
+	container *postgres.PostgresContainer
 	repo      *Repository
 	settings  config.Settings
 }
@@ -65,7 +65,7 @@ func (o *AccessibleVehiclesRepoTestSuite) SetupSuite() {
 
 // TearDownTest after each test truncate tables
 func (s *AccessibleVehiclesRepoTestSuite) TearDownTest() {
-	test.TruncateTables(s.pdb.DBS().Writer.DB, s.T())
+	s.Require().NoError(s.container.Restore(s.ctx))
 }
 
 // TearDownSuite cleanup at end by terminating container