bootstrap_test.go

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

     1  package db_test
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  
     8  	"github.com/jackc/pgx/v5/pgxpool"
     9  	"github.com/pashagolub/pgxmock/v4"
    10  	"github.com/stretchr/testify/assert"
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/cybertec-postgresql/pgwatch/v5/internal/db"
    14  	"github.com/cybertec-postgresql/pgwatch/v5/internal/testutil"
    15  )
    16  
    17  func TestPing(t *testing.T) {
    18  	connStr := "foo_boo"
    19  	assert.Error(t, db.Ping(testutil.TestContext, connStr))
    20  
    21  	pg, pgTeardown, err := testutil.SetupPostgresContainer()
    22  	require.NoError(t, err)
    23  	defer pgTeardown()
    24  
    25  	connStr, err = pg.ConnectionString(testutil.TestContext)
    26  	assert.NoError(t, err)
    27  	assert.NoError(t, db.Ping(testutil.TestContext, connStr))
    28  	assert.NoError(t, pg.Terminate(testutil.TestContext))
    29  }
    30  
    31  func TestDoesSchemaExist(t *testing.T) {
    32  	conn, err := pgxmock.NewPool()
    33  	assert.NoError(t, err)
    34  	conn.ExpectQuery("SELECT EXISTS").
    35  		WithArgs("public").
    36  		WillReturnRows(pgxmock.NewRows([]string{"exists"}).AddRow(true))
    37  	exists, err := db.DoesSchemaExist(testutil.TestContext, conn, "public")
    38  	assert.NoError(t, err)
    39  	assert.True(t, exists)
    40  }
    41  func TestInit(t *testing.T) {
    42  
    43  	conn, err := pgxmock.NewPool()
    44  	assert.NoError(t, err)
    45  	initCalled := false
    46  	initFunc := func(context.Context, db.PgxIface) error {
    47  		initCalled = true
    48  		return nil
    49  	}
    50  
    51  	// Test successful initialization
    52  	conn.ExpectPing()
    53  	err = db.Init(testutil.TestContext, conn, initFunc)
    54  	assert.NoError(t, err)
    55  	assert.True(t, initCalled)
    56  
    57  	// Test failed initialization with 3 retries
    58  	conn.ExpectPing().Times(1 + 3).WillReturnError(errors.New("connection failed"))
    59  	initCalled = false
    60  	err = db.Init(testutil.TestContext, conn, initFunc)
    61  	assert.Error(t, err)
    62  	assert.False(t, initCalled)
    63  
    64  	assert.NoError(t, conn.ExpectationsWereMet())
    65  }
    66  
    67  func TestNew(t *testing.T) {
    68  	pg, pgTeardown, err := testutil.SetupPostgresContainer()
    69  	require.NoError(t, err)
    70  	defer pgTeardown()
    71  
    72  	connStr, err := pg.ConnectionString(testutil.TestContext)
    73  	t.Log(connStr)
    74  	assert.NoError(t, err)
    75  
    76  	initCalled := false
    77  	initFunc := func(*pgxpool.Config) error {
    78  		initCalled = true
    79  		return nil
    80  	}
    81  	// Test successful initialization
    82  	pool, err := db.New(context.Background(), connStr, initFunc)
    83  	assert.NoError(t, err)
    84  	assert.NotNil(t, pool)
    85  	assert.True(t, initCalled)
    86  	_, err = pool.Exec(testutil.TestContext, `DO $$
    87  BEGIN
    88     RAISE NOTICE 'This is a notice';
    89  END $$;`)
    90  	assert.NoError(t, err)
    91  	pool.Close()
    92  
    93  	// Test failed initialization
    94  	initCalled = false
    95  	pool, err = db.New(context.Background(), "foo", initFunc)
    96  	assert.Error(t, err)
    97  	assert.Nil(t, pool)
    98  	assert.False(t, initCalled)
    99  
   100  	// Test failed initialization with callback
   101  	initFunc = func(*pgxpool.Config) error {
   102  		return errors.New("callback failed")
   103  	}
   104  	initCalled = false
   105  	pool, err = db.New(context.Background(), connStr, initFunc)
   106  	assert.Error(t, err)
   107  	assert.Nil(t, pool)
   108  	assert.False(t, initCalled)
   109  }
   110