Skip to content

Commit

Permalink
Remove use of testify/suite
Browse files Browse the repository at this point in the history
  • Loading branch information
@serprex
serprex committed Dec 22, 2023
1 parent 08861aa commit 47a9538
Show file tree
Hide file tree
Showing 16 changed files with 412 additions and 412 deletions.
60 changes: 34 additions & 26 deletions flow/connectors/postgres/postgres_repl_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,104 +7,112 @@ import (
"testing"
"time"

"github.com/PeerDB-io/peer-flow/e2e/e2eshared"
"github.com/PeerDB-io/peer-flow/generated/protos"
"github.com/PeerDB-io/peer-flow/shared"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgtype"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
"github.com/ysmood/got"
)

type PostgresReplicationSnapshotTestSuite struct {
suite.Suite
got.G
t *testing.T

connector *PostgresConnector
}

func (suite *PostgresReplicationSnapshotTestSuite) SetupSuite() {
var err error
suite.connector, err = NewPostgresConnector(context.Background(), &protos.PostgresConfig{
func setupSuite(t *testing.T, g got.G) PostgresReplicationSnapshotTestSuite {
connector, err := NewPostgresConnector(context.Background(), &protos.PostgresConfig{
Host: "localhost",
Port: 7132,
User: "postgres",
Password: "postgres",
Database: "postgres",
})
require.NoError(suite.T(), err)
require.NoError(t, err)

setupTx, err := suite.connector.pool.Begin(context.Background())
require.NoError(suite.T(), err)
setupTx, err := connector.pool.Begin(context.Background())
require.NoError(t, err)
defer func() {
err := setupTx.Rollback(context.Background())
if err != pgx.ErrTxClosed {
require.NoError(suite.T(), err)
require.NoError(t, err)
}
}()

_, err = setupTx.Exec(context.Background(), "DROP SCHEMA IF EXISTS pgpeer_repl_test CASCADE")
require.NoError(suite.T(), err)
require.NoError(t, err)

_, err = setupTx.Exec(context.Background(), "CREATE SCHEMA pgpeer_repl_test")
require.NoError(suite.T(), err)
require.NoError(t, err)

// setup 3 tables in pgpeer_repl_test schema
// test_1, test_2, test_3, all have 5 columns all text, c1, c2, c3, c4, c5
tables := []string{"test_1", "test_2", "test_3"}
for _, table := range tables {
_, err = setupTx.Exec(context.Background(),
fmt.Sprintf("CREATE TABLE pgpeer_repl_test.%s (c1 text, c2 text, c3 text, c4 text, c5 text)", table))
require.NoError(suite.T(), err)
require.NoError(t, err)
}

err = setupTx.Commit(context.Background())
require.NoError(suite.T(), err)
require.NoError(t, err)

return PostgresReplicationSnapshotTestSuite{
G: g,
t: t,
connector: connector,
}
}

func (suite *PostgresReplicationSnapshotTestSuite) TearDownSuite() {
func (suite PostgresReplicationSnapshotTestSuite) TearDownSuite() {
teardownTx, err := suite.connector.pool.Begin(context.Background())
require.NoError(suite.T(), err)
require.NoError(suite.t, err)
defer func() {
err := teardownTx.Rollback(context.Background())
if err != pgx.ErrTxClosed {
require.NoError(suite.T(), err)
require.NoError(suite.t, err)
}
}()

_, err = teardownTx.Exec(context.Background(), "DROP SCHEMA IF EXISTS pgpeer_test CASCADE")
require.NoError(suite.T(), err)
require.NoError(suite.t, err)

// Fetch all the publications
rows, err := teardownTx.Query(context.Background(),
"SELECT pubname FROM pg_publication WHERE pubname LIKE $1", fmt.Sprintf("%%%s", "test_simple_slot_creation"))
require.NoError(suite.T(), err)
require.NoError(suite.t, err)

// Iterate over the publications and drop them
for rows.Next() {
var pubname pgtype.Text
err := rows.Scan(&pubname)
require.NoError(suite.T(), err)
require.NoError(suite.t, err)

// Drop the publication in a new transaction
_, err = suite.connector.pool.Exec(context.Background(), fmt.Sprintf("DROP PUBLICATION %s", pubname.String))
require.NoError(suite.T(), err)
require.NoError(suite.t, err)
}

_, err = teardownTx.Exec(context.Background(),
"SELECT pg_drop_replication_slot(slot_name) FROM pg_replication_slots WHERE slot_name LIKE $1",
fmt.Sprintf("%%%s", "test_simple_slot_creation"))
require.NoError(suite.T(), err)
require.NoError(suite.t, err)

err = teardownTx.Commit(context.Background())
require.NoError(suite.T(), err)
require.NoError(suite.t, err)

suite.True(suite.connector.ConnectionActive() == nil)

err = suite.connector.Close()
require.NoError(suite.T(), err)
require.NoError(suite.t, err)

suite.False(suite.connector.ConnectionActive() == nil)
}

func (suite *PostgresReplicationSnapshotTestSuite) TestSimpleSlotCreation() {
func (suite PostgresReplicationSnapshotTestSuite) TestSimpleSlotCreation() {
tables := map[string]string{
"pgpeer_repl_test.test_1": "test_1_dst",
}
Expand All @@ -121,7 +129,7 @@ func (suite *PostgresReplicationSnapshotTestSuite) TestSimpleSlotCreation() {
// Moved to a go routine
go func() {
err := suite.connector.SetupReplication(signal, setupReplicationInput)
require.NoError(suite.T(), err)
require.NoError(suite.t, err)
}()

slog.Info("waiting for slot creation to complete", flowLog)
Expand All @@ -136,5 +144,5 @@ func (suite *PostgresReplicationSnapshotTestSuite) TestSimpleSlotCreation() {
}

func TestPostgresReplTestSuite(t *testing.T) {
suite.Run(t, new(PostgresReplicationSnapshotTestSuite))
got.Each(t, e2eshared.GotSuite(setupSuite))
}
84 changes: 43 additions & 41 deletions flow/connectors/postgres/postgres_schema_delta_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,79 +5,81 @@ import (
"fmt"
"testing"

"github.com/PeerDB-io/peer-flow/e2e/e2eshared"
"github.com/PeerDB-io/peer-flow/generated/protos"
"github.com/PeerDB-io/peer-flow/model/qvalue"
"github.com/jackc/pgx/v5"
"github.com/stretchr/testify/suite"
"github.com/stretchr/testify/require"
"github.com/ysmood/got"
)

type PostgresSchemaDeltaTestSuite struct {
suite.Suite
got.G
t *testing.T

connector *PostgresConnector
}

const schemaDeltaTestSchemaName = "pgschema_delta_test"

func (suite *PostgresSchemaDeltaTestSuite) failTestError(err error) {
if err != nil {
suite.FailNow(err.Error())
}
}

func (suite *PostgresSchemaDeltaTestSuite) SetupSuite() {
var err error
suite.connector, err = NewPostgresConnector(context.Background(), &protos.PostgresConfig{
func setupSchemaDeltaSuite(t *testing.T, g got.G) PostgresSchemaDeltaTestSuite {
connector, err := NewPostgresConnector(context.Background(), &protos.PostgresConfig{
Host: "localhost",
Port: 7132,
User: "postgres",
Password: "postgres",
Database: "postgres",
})
suite.failTestError(err)
require.NoError(t, err)

setupTx, err := suite.connector.pool.Begin(context.Background())
suite.failTestError(err)
setupTx, err := connector.pool.Begin(context.Background())
require.NoError(t, err)
defer func() {
err := setupTx.Rollback(context.Background())
if err != pgx.ErrTxClosed {
suite.failTestError(err)
require.NoError(t, err)
}
}()
_, err = setupTx.Exec(context.Background(), fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE",
schemaDeltaTestSchemaName))
suite.failTestError(err)
require.NoError(t, err)
_, err = setupTx.Exec(context.Background(), fmt.Sprintf("CREATE SCHEMA %s", schemaDeltaTestSchemaName))
suite.failTestError(err)
require.NoError(t, err)
err = setupTx.Commit(context.Background())
suite.failTestError(err)
require.NoError(t, err)
return PostgresSchemaDeltaTestSuite{
G: g,
t: t,
connector: connector,
}
}

func (suite *PostgresSchemaDeltaTestSuite) TearDownSuite() {
func (suite PostgresSchemaDeltaTestSuite) TearDownSuite() {
teardownTx, err := suite.connector.pool.Begin(context.Background())
suite.failTestError(err)
require.NoError(suite.t, err)
defer func() {
err := teardownTx.Rollback(context.Background())
if err != pgx.ErrTxClosed {
suite.failTestError(err)
require.NoError(suite.t, err)
}
}()
_, err = teardownTx.Exec(context.Background(), fmt.Sprintf("DROP SCHEMA IF EXISTS %s CASCADE",
schemaDeltaTestSchemaName))
suite.failTestError(err)
require.NoError(suite.t, err)
err = teardownTx.Commit(context.Background())
suite.failTestError(err)
require.NoError(suite.t, err)

suite.True(suite.connector.ConnectionActive() == nil)
err = suite.connector.Close()
suite.failTestError(err)
require.NoError(suite.t, err)
suite.False(suite.connector.ConnectionActive() == nil)
}

func (suite *PostgresSchemaDeltaTestSuite) TestSimpleAddColumn() {
func (suite PostgresSchemaDeltaTestSuite) TestSimpleAddColumn() {
tableName := fmt.Sprintf("%s.simple_add_column", schemaDeltaTestSchemaName)
_, err := suite.connector.pool.Exec(context.Background(),
fmt.Sprintf("CREATE TABLE %s(id INT PRIMARY KEY)", tableName))
suite.failTestError(err)
require.NoError(suite.t, err)

err = suite.connector.ReplayTableSchemaDeltas("schema_delta_flow", []*protos.TableSchemaDelta{{
SrcTableName: tableName,
Expand All @@ -87,12 +89,12 @@ func (suite *PostgresSchemaDeltaTestSuite) TestSimpleAddColumn() {
ColumnType: string(qvalue.QValueKindInt64),
}},
}})
suite.failTestError(err)
require.NoError(suite.t, err)

output, err := suite.connector.GetTableSchema(&protos.GetTableSchemaBatchInput{
TableIdentifiers: []string{tableName},
})
suite.failTestError(err)
require.NoError(suite.t, err)
suite.Equal(&protos.TableSchema{
TableIdentifier: tableName,
Columns: map[string]string{
Expand All @@ -103,11 +105,11 @@ func (suite *PostgresSchemaDeltaTestSuite) TestSimpleAddColumn() {
}, output.TableNameSchemaMapping[tableName])
}

func (suite *PostgresSchemaDeltaTestSuite) TestAddAllColumnTypes() {
func (suite PostgresSchemaDeltaTestSuite) TestAddAllColumnTypes() {
tableName := fmt.Sprintf("%s.add_drop_all_column_types", schemaDeltaTestSchemaName)
_, err := suite.connector.pool.Exec(context.Background(),
fmt.Sprintf("CREATE TABLE %s(id INT PRIMARY KEY)", tableName))
suite.failTestError(err)
require.NoError(suite.t, err)

expectedTableSchema := &protos.TableSchema{
TableIdentifier: tableName,
Expand Down Expand Up @@ -148,20 +150,20 @@ func (suite *PostgresSchemaDeltaTestSuite) TestAddAllColumnTypes() {
DstTableName: tableName,
AddedColumns: addedColumns,
}})
suite.failTestError(err)
require.NoError(suite.t, err)

output, err := suite.connector.GetTableSchema(&protos.GetTableSchemaBatchInput{
TableIdentifiers: []string{tableName},
})
suite.failTestError(err)
require.NoError(suite.t, err)
suite.Equal(expectedTableSchema, output.TableNameSchemaMapping[tableName])
}

func (suite *PostgresSchemaDeltaTestSuite) TestAddTrickyColumnNames() {
func (suite PostgresSchemaDeltaTestSuite) TestAddTrickyColumnNames() {
tableName := fmt.Sprintf("%s.add_drop_tricky_column_names", schemaDeltaTestSchemaName)
_, err := suite.connector.pool.Exec(context.Background(),
fmt.Sprintf("CREATE TABLE %s(id INT PRIMARY KEY)", tableName))
suite.failTestError(err)
require.NoError(suite.t, err)

expectedTableSchema := &protos.TableSchema{
TableIdentifier: tableName,
Expand Down Expand Up @@ -194,20 +196,20 @@ func (suite *PostgresSchemaDeltaTestSuite) TestAddTrickyColumnNames() {
DstTableName: tableName,
AddedColumns: addedColumns,
}})
suite.failTestError(err)
require.NoError(suite.t, err)

output, err := suite.connector.GetTableSchema(&protos.GetTableSchemaBatchInput{
TableIdentifiers: []string{tableName},
})
suite.failTestError(err)
require.NoError(suite.t, err)
suite.Equal(expectedTableSchema, output.TableNameSchemaMapping[tableName])
}

func (suite *PostgresSchemaDeltaTestSuite) TestAddDropWhitespaceColumnNames() {
func (suite PostgresSchemaDeltaTestSuite) TestAddDropWhitespaceColumnNames() {
tableName := fmt.Sprintf("%s.add_drop_whitespace_column_names", schemaDeltaTestSchemaName)
_, err := suite.connector.pool.Exec(context.Background(),
fmt.Sprintf("CREATE TABLE %s(\" \" INT PRIMARY KEY)", tableName))
suite.failTestError(err)
require.NoError(suite.t, err)

expectedTableSchema := &protos.TableSchema{
TableIdentifier: tableName,
Expand All @@ -234,15 +236,15 @@ func (suite *PostgresSchemaDeltaTestSuite) TestAddDropWhitespaceColumnNames() {
DstTableName: tableName,
AddedColumns: addedColumns,
}})
suite.failTestError(err)
require.NoError(suite.t, err)

output, err := suite.connector.GetTableSchema(&protos.GetTableSchemaBatchInput{
TableIdentifiers: []string{tableName},
})
suite.failTestError(err)
require.NoError(suite.t, err)
suite.Equal(expectedTableSchema, output.TableNameSchemaMapping[tableName])
}

func TestPostgresSchemaDeltaTestSuite(t *testing.T) {
suite.Run(t, new(PostgresSchemaDeltaTestSuite))
got.Each(t, e2eshared.GotSuite(setupSchemaDeltaSuite))
}
Loading

0 comments on commit 47a9538

Please sign in to comment.