bootstrap_test.go

Documentation: github.com/cybertec-postgresql/pgwatch/v3/internal/db

     1  package db_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/jackc/pgx/v5/pgxpool"
    10  	"github.com/pashagolub/pgxmock/v4"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/cybertec-postgresql/pgwatch/v3/internal/db"
    15  	testcontainers "github.com/testcontainers/testcontainers-go"
    16  	"github.com/testcontainers/testcontainers-go/modules/postgres"
    17  	"github.com/testcontainers/testcontainers-go/wait"
    18  )
    19  
    20  const ImageName = "docker.io/postgres:17-alpine"
    21  
    22  var ctx = context.Background()
    23  
    24  func TestPing(t *testing.T) {
    25  	connStr := "foo_boo"
    26  	assert.Error(t, db.Ping(ctx, connStr))
    27  
    28  	pg, err := initTestContainer()
    29  	require.NoError(t, err)
    30  	connStr, err = pg.ConnectionString(ctx)
    31  	assert.NoError(t, err)
    32  	assert.NoError(t, db.Ping(ctx, connStr))
    33  	assert.NoError(t, pg.Terminate(ctx))
    34  }
    35  
    36  func TestDoesSchemaExist(t *testing.T) {
    37  	conn, err := pgxmock.NewPool()
    38  	assert.NoError(t, err)
    39  	conn.ExpectQuery("SELECT EXISTS").
    40  		WithArgs("public").
    41  		WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
    42  	exists, err := db.DoesSchemaExist(ctx, conn, "public")
    43  	assert.NoError(t, err)
    44  	assert.True(t, exists)
    45  }
    46  func TestInit(t *testing.T) {
    47  
    48  	conn, err := pgxmock.NewPool()
    49  	assert.NoError(t, err)
    50  	initCalled := false
    51  	initFunc := func(context.Context, db.PgxIface) error {
    52  		initCalled = true
    53  		return nil
    54  	}
    55  
    56  	// Test successful initialization
    57  	conn.ExpectPing()
    58  	err = db.Init(ctx, conn, initFunc)
    59  	assert.NoError(t, err)
    60  	assert.True(t, initCalled)
    61  
    62  	// Test failed initialization with 3 retries
    63  	conn.ExpectPing().Times(1 + 3).WillReturnError(errors.New("connection failed"))
    64  	initCalled = false
    65  	err = db.Init(ctx, conn, initFunc)
    66  	assert.Error(t, err)
    67  	assert.False(t, initCalled)
    68  
    69  	assert.NoError(t, conn.ExpectationsWereMet())
    70  }
    71  
    72  func initTestContainer() (*postgres.PostgresContainer, error) {
    73  	dbName := "pgwatch"
    74  	dbUser := "pgwatch"
    75  	dbPassword := "pgwatchadmin"
    76  
    77  	return postgres.Run(ctx,
    78  		ImageName,
    79  		postgres.WithDatabase(dbName),
    80  		postgres.WithUsername(dbUser),
    81  		postgres.WithPassword(dbPassword),
    82  		testcontainers.WithWaitStrategy(
    83  			wait.ForLog("database system is ready to accept connections").
    84  				WithOccurrence(2).
    85  				WithStartupTimeout(5*time.Second)),
    86  	)
    87  }
    88  
    89  func TestNew(t *testing.T) {
    90  	pg, err := initTestContainer()
    91  	require.NoError(t, err)
    92  	defer func() { assert.NoError(t, pg.Terminate(ctx)) }()
    93  	connStr, err := pg.ConnectionString(ctx)
    94  	t.Log(connStr)
    95  	assert.NoError(t, err)
    96  
    97  	initCalled := false
    98  	initFunc := func(*pgxpool.Config) error {
    99  		initCalled = true
   100  		return nil
   101  	}
   102  	// Test successful initialization
   103  	pool, err := db.New(context.Background(), connStr, initFunc)
   104  	assert.NoError(t, err)
   105  	assert.NotNil(t, pool)
   106  	assert.True(t, initCalled)
   107  	_, err = pool.Exec(ctx, `DO $$
   108  BEGIN
   109     RAISE NOTICE 'This is a notice';
   110  END $$;`)
   111  	assert.NoError(t, err)
   112  	pool.Close()
   113  
   114  	// Test failed initialization
   115  	initCalled = false
   116  	pool, err = db.New(context.Background(), "foo", initFunc)
   117  	assert.Error(t, err)
   118  	assert.Nil(t, pool)
   119  	assert.False(t, initCalled)
   120  
   121  	// Test failed initialization with callback
   122  	initFunc = func(*pgxpool.Config) error {
   123  		return errors.New("callback failed")
   124  	}
   125  	initCalled = false
   126  	pool, err = db.New(context.Background(), connStr, initFunc)
   127  	assert.Error(t, err)
   128  	assert.Nil(t, pool)
   129  	assert.False(t, initCalled)
   130  }
   131