Skip to content

Tests reliability improvements: Connection parameters extraction logic unification #4192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pfcoperez merged 16 commits into from
Apr 22, 2026
Merged

Tests reliability improvements: Connection parameters extraction logic unification #4192

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
5a5788a
Connection parameters extraction logic unification
pfcoperez Apr 17, 2026
dd47750
Colocate IAM helpers in the same package and file.
pfcoperez Apr 21, 2026
591968f
Further use of common test helper library in MySQL/Maria tests
pfcoperez Apr 21, 2026
0d2b17a
Use common test library for PSQL RDS parameters
pfcoperez Apr 21, 2026
c496c3c
Use common test helper library to get MongoDB connection parameters. …
pfcoperez Apr 21, 2026
9eccded
Expanded helper functions
pfcoperez Apr 21, 2026
5d93ddb
Refactor MYSQL/Maria parameter extraction using the default-as-parame…
pfcoperez Apr 21, 2026
feeec76
TO-REVERT: Temporarily re-enable `TestAwsRDSIAMAuthConnectForPostgres…
pfcoperez Apr 22, 2026
606029e
PR Comments: Move MongoDB env vars validations to the common test too…
pfcoperez Apr 22, 2026
e63cd4c
Add checks to RDS env var parameters
pfcoperez Apr 22, 2026
7c49893
Merge remote-tracking branch 'upstream/main' into DBI-674/unified-con…
pfcoperez Apr 22, 2026
6058d5c
Moved `GetAncillaryPostgresConfigFromEnv` to test helpers module.
pfcoperez Apr 22, 2026
552396b
PR Comments: Consistent MySQL format in `MySQLTestVersionIsMySQLPos` …
pfcoperez Apr 22, 2026
639031e
Revert "TO-REVERT: Temporarily re-enable `TestAwsRDSIAMAuthConnectFor…
pfcoperez Apr 22, 2026
10c9fdb
Factor out AWS credentials extraction
pfcoperez Apr 22, 2026
e96f0a6
Factor out `TOXIPROXY_POSTGRES_HOST` access.
pfcoperez Apr 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 4 additions & 23 deletions flow/connectors/clickhouse/s3_iam_role_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,14 @@ package connclickhouse
import (
"fmt"
"os"
"strconv"
"testing"
"time"

"github.com/stretchr/testify/require"

"github.com/PeerDB-io/peerdb/flow/connectors/utils"
"github.com/PeerDB-io/peerdb/flow/generated/protos"
"github.com/PeerDB-io/peerdb/flow/internal"
)

func TestIAMRoleCanIssueSelectFromS3(t *testing.T) {
Expand All @@ -33,33 +33,14 @@ func TestIAMRoleCanIssueSelectFromS3(t *testing.T) {
} {
t.Setenv(envVar, "")
}
for _, envVar := range []string{
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_SESSION_TOKEN",
} {
t.Setenv(envVar, os.Getenv("FLOW_TESTS_"+envVar))
}
internal.SetupFlowAWSCredentialsFromEnv(t)
t.Setenv("PEERDB_CLICKHOUSE_AWS_S3_BUCKET_NAME", os.Getenv(bucketNameEnvVar))
ctx := t.Context()

chHost := "localhost"
if host := os.Getenv("CI_CLICKHOUSE_HOST"); host != "" {
chHost = host
}
var chPort uint32 = 9000
if port := os.Getenv("CI_CLICKHOUSE_NATIVE_PORT"); port != "" {
p, err := strconv.ParseUint(port, 10, 32)
if err != nil {
t.Fatalf("invalid CI_CLICKHOUSE_NATIVE_PORT %q: %v", port, err)
}
chPort = uint32(p)
}

conn, err := NewClickHouseConnector(ctx, nil,
&protos.ClickhouseConfig{
Host: chHost,
Port: chPort,
Host: internal.ClickHouseTestHost(),
Port: internal.ClickHouseTestPort(),
Database: "default",
DisableTls: true,
})
Comment thread
pfcoperez marked this conversation as resolved.
Expand Down
48 changes: 12 additions & 36 deletions flow/connectors/mysql/rds_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,68 +1,44 @@
package connmysql

import (
"os"
"testing"

"github.com/stretchr/testify/require"

"github.com/PeerDB-io/peerdb/flow/generated/protos"
"github.com/PeerDB-io/peerdb/flow/internal"
)

func TestAwsRDSIAMAuthConnectForMYSQL(t *testing.T) {
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
host := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_MYSQL")
username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_MYSQL")
internal.SetupRDSIAMAuthAWSCredentials(t)
conn := internal.RDSIAMAuthMySQLTestConnectionInfo(t)
mysqlConnector, err := NewMySqlConnector(t.Context(),
&protos.MySqlConfig{
Host: host,
Host: conn.Host,
Database: "postgres",
User: username,
User: conn.Username,
Port: 5432,
AuthType: protos.MySqlAuthType_MYSQL_IAM_AUTH,
DisableTls: false, // Assumed that AWS Root CA is installed
AwsAuth: &protos.AwsAuthenticationConfig{
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
AuthConfig: &protos.AwsAuthenticationConfig_Role{
Role: &protos.AWSAuthAssumeRoleConfig{
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
ChainedRoleArn: new(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
},
},
},
AwsAuth: internal.RDSIAMAuthAssumeRoleConfig(t),
})
Comment thread
pfcoperez marked this conversation as resolved.
require.NoError(t, err)
require.NoError(t, mysqlConnector.Ping(t.Context()))
}

func TestAwsRDSIAMAuthConnectForMYSQLViaProxy(t *testing.T) {
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
rdsHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_MYSQL")
proxyHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_MYSQL_PROXY")

username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_MYSQL")
internal.SetupRDSIAMAuthAWSCredentials(t)
conn := internal.RDSIAMAuthMySQLTestConnectionInfo(t)
mysqlConnector, err := NewMySqlConnector(t.Context(),
&protos.MySqlConfig{
Host: proxyHost,
Host: conn.ProxyHost,
Database: "postgres",
User: username,
User: conn.Username,
Port: 5432,
TlsHost: rdsHost,
TlsHost: conn.Host,
AuthType: protos.MySqlAuthType_MYSQL_IAM_AUTH,
DisableTls: false, // Assumed that AWS Root CA is installed
AwsAuth: &protos.AwsAuthenticationConfig{
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
AuthConfig: &protos.AwsAuthenticationConfig_Role{
Role: &protos.AWSAuthAssumeRoleConfig{
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
ChainedRoleArn: new(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
},
},
},
AwsAuth: internal.RDSIAMAuthAssumeRoleConfig(t),
})
require.NoError(t, err)
require.NoError(t, mysqlConnector.Ping(t.Context()))
Expand Down
37 changes: 10 additions & 27 deletions flow/connectors/mysql/ssh_keepalive_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,7 @@ package connmysql

import (
"context"
"os"
"strconv"
"strings"
"sync/atomic"
"testing"
"time"
Expand All @@ -14,6 +12,7 @@ import (

"github.com/PeerDB-io/peerdb/flow/connectors/utils"
"github.com/PeerDB-io/peerdb/flow/generated/protos"
"github.com/PeerDB-io/peerdb/flow/internal"
"github.com/PeerDB-io/peerdb/flow/model"
"github.com/PeerDB-io/peerdb/flow/otel_metrics"
"github.com/PeerDB-io/peerdb/flow/shared"
Expand All @@ -31,25 +30,9 @@ const (

func resolveMySQL(t *testing.T) (string, uint32, string) {
t.Helper()

mysqlHost := "mysql"
if envHost := os.Getenv("CI_MYSQL_HOST"); envHost != "" {
mysqlHost = envHost
}

var mysqlPort uint32 = 3306
if envPortStr := os.Getenv("CI_MYSQL_PORT"); envPortStr != "" {
envPort, err := strconv.ParseUint(strings.TrimSpace(strings.Split(envPortStr, "#")[0]), 10, 32)
require.NoError(t, err, "Failed to parse CI_MYSQL_PORT")
mysqlPort = uint32(envPort)
}

mysqlRootPass := "cipass"
if envPass := os.Getenv("CI_MYSQL_ROOT_PASSWORD"); envPass != "" {
mysqlRootPass = envPass
}

return mysqlHost, mysqlPort, mysqlRootPass
return internal.MySQLTestHostWithFallback("mysql"),
internal.MySQLTestPortWithFallback(3306),
internal.MySQLTestRootPasswordWithFallback("cipass")
}

// Connector -> Toxi -> SSH -> MySQL
Expand Down Expand Up @@ -141,7 +124,7 @@ func setupMySQLSSHKeepaliveHarness(ctx context.Context, t *testing.T, proxyName

func TestMySQLSSHKeepaliveTunnelDown(t *testing.T) {
t.Parallel()
if os.Getenv("CI_MYSQL_VERSION") == "maria" {
if internal.MySQLTestVersionIsMaria() {
t.Skip("Skipping SSH keepalive test for MariaDB")
}
connector, cfg := setupMySQLSSHKeepaliveHarness(t.Context(), t, "my-ssh-keepalive-test", toxiproxyDownProxyPort)
Expand All @@ -151,7 +134,7 @@ func TestMySQLSSHKeepaliveTunnelDown(t *testing.T) {

func TestMySQLSSHKeepaliveLatency(t *testing.T) {
t.Parallel()
if os.Getenv("CI_MYSQL_VERSION") == "maria" {
if internal.MySQLTestVersionIsMaria() {
t.Skip("Skipping SSH keepalive test for MariaDB")
}
connector, cfg := setupMySQLSSHKeepaliveHarness(t.Context(), t, "my-ssh-latency-test", toxiproxyLatencyProxyPort)
Expand All @@ -161,7 +144,7 @@ func TestMySQLSSHKeepaliveLatency(t *testing.T) {

func TestMySQLSSHResetPeer(t *testing.T) {
t.Parallel()
if os.Getenv("CI_MYSQL_VERSION") == "maria" {
if internal.MySQLTestVersionIsMaria() {
t.Skip("Skipping SSH keepalive test for MariaDB")
}
connector, cfg := setupMySQLSSHKeepaliveHarness(t.Context(), t, "my-ssh-reset-peer-test", toxiproxyResetProxyPort)
Expand Down Expand Up @@ -207,7 +190,7 @@ func setupCDCPullRecords(

func TestMySQLSSHKeepaliveCDCHang(t *testing.T) {
t.Parallel()
if os.Getenv("CI_MYSQL_VERSION") == "maria" {
if internal.MySQLTestVersionIsMaria() {
t.Skip("Skipping SSH keepalive test for MariaDB")
}
ctx := t.Context()
Expand Down Expand Up @@ -260,7 +243,7 @@ func TestMySQLSSHKeepaliveCDCHang(t *testing.T) {

func TestMySQLSSHKeepaliveCDCCloseHang(t *testing.T) {
t.Parallel()
if os.Getenv("CI_MYSQL_VERSION") == "maria" {
if internal.MySQLTestVersionIsMaria() {
t.Skip("Skipping SSH keepalive test for MariaDB")
}

Expand Down Expand Up @@ -320,7 +303,7 @@ func TestMySQLSSHKeepaliveCDCCloseHang(t *testing.T) {

func TestMySQLCloseSyncerWithTimeout(t *testing.T) {
t.Parallel()
if os.Getenv("CI_MYSQL_VERSION") == "maria" {
if internal.MySQLTestVersionIsMaria() {
t.Skip("Skipping for MariaDB")
}

Expand Down
48 changes: 12 additions & 36 deletions flow/connectors/postgres/rds_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,39 +1,28 @@
package connpostgres

import (
"os"
"testing"

"github.com/stretchr/testify/require"

"github.com/PeerDB-io/peerdb/flow/generated/protos"
"github.com/PeerDB-io/peerdb/flow/internal"
)

func TestAwsRDSIAMAuthConnectForPostgres(t *testing.T) {
t.Skip("flaky")
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
host := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES")
username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_POSTGRES")
internal.SetupRDSIAMAuthAWSCredentials(t)
conn := internal.RDSIAMAuthPostgresTestConnectionInfo(t)
postgresConnector, err := NewPostgresConnector(t.Context(),
nil,
&protos.PostgresConfig{
Host: host,
Host: conn.Host,
Database: "postgres",
User: username,
User: conn.Username,
Port: 5432,
AuthType: protos.PostgresAuthType_POSTGRES_IAM_AUTH,
RequireTls: true, // Assumed that AWS Root CA is installed
AwsAuth: &protos.AwsAuthenticationConfig{
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
AuthConfig: &protos.AwsAuthenticationConfig_Role{
Role: &protos.AWSAuthAssumeRoleConfig{
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
ChainedRoleArn: new(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
},
},
},
AwsAuth: internal.RDSIAMAuthAssumeRoleConfig(t),
})
require.NoError(t, err)
defer postgresConnector.Close()
Expand All @@ -51,32 +40,19 @@ func TestAwsRDSIAMAuthConnectForPostgres(t *testing.T) {

func TestAwsRDSIAMAuthConnectForPostgresViaProxy(t *testing.T) {
t.Skip("flaky")
t.Setenv("AWS_ACCESS_KEY_ID", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_ACCESS_KEY_ID"))
t.Setenv("AWS_SECRET_ACCESS_KEY", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SECRET_ACCESS_KEY"))
t.Setenv("AWS_SESSION_TOKEN", os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_AWS_SESSION_TOKEN"))
rdsHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES")
proxyHost := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_HOST_POSTGRES_PROXY")

username := os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_USERNAME_POSTGRES")
internal.SetupRDSIAMAuthAWSCredentials(t)
conn := internal.RDSIAMAuthPostgresTestConnectionInfo(t)
postgresConnector, err := NewPostgresConnector(t.Context(),
nil,
&protos.PostgresConfig{
Host: proxyHost,
Host: conn.ProxyHost,
Port: 5432,
User: username,
User: conn.Username,
Database: "postgres",
TlsHost: rdsHost,
TlsHost: conn.Host,
RequireTls: true, // Assumed that AWS Root CA is installed
AuthType: protos.PostgresAuthType_POSTGRES_IAM_AUTH,
AwsAuth: &protos.AwsAuthenticationConfig{
AuthType: protos.AwsIAMAuthConfigType_IAM_AUTH_ASSUME_ROLE,
AuthConfig: &protos.AwsAuthenticationConfig_Role{
Role: &protos.AWSAuthAssumeRoleConfig{
AssumeRoleArn: os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_ASSUME_ROLE"),
ChainedRoleArn: new(os.Getenv("FLOW_TESTS_RDS_IAM_AUTH_CHAINED_ROLE")),
},
},
},
AwsAuth: internal.RDSIAMAuthAssumeRoleConfig(t),
})
require.NoError(t, err)
defer postgresConnector.Close()
Expand Down
44 changes: 11 additions & 33 deletions flow/connectors/postgres/ssh_keepalive_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,13 @@ package connpostgres

import (
"context"
"os"
"strconv"
"strings"
"testing"

"github.com/stretchr/testify/require"

"github.com/PeerDB-io/peerdb/flow/connectors/utils"
"github.com/PeerDB-io/peerdb/flow/generated/protos"
"github.com/PeerDB-io/peerdb/flow/internal"
)

const (
Expand All @@ -26,40 +24,20 @@ func setupPostgresConnectorWithSSH(ctx context.Context, t *testing.T, proxyName
toxiproxyClient := utils.NewToxiproxyClient(t)
sshProxy := utils.CreateSSHProxy(t, toxiproxyClient, proxyName, proxyPort)

pgHost := "localhost" // Default for local environments

// In local set-up environments like Tilt, when a non catalog
// instance of Postgres is present, we use it instead of the catalog.
if envHost := os.Getenv("PG_HOST"); envHost != "" {
pgHost = envHost
}

// In CI, TOXIPROXY_POSTGRES_HOST is set, usually pointing to 'catalog'
// network.
if toxiproxyPgHost := os.Getenv("TOXIPROXY_POSTGRES_HOST"); toxiproxyPgHost != "" {
pgHost = toxiproxyPgHost
pgConfig := internal.GetAncillaryPostgresConfigFromEnv()
// In CI, TOXIPROXY_POSTGRES_HOST is set, usually pointing to 'catalog' network,
// in this case, we override the host in pgConfig with the value from TOXIPROXY_POSTGRES_HOST.
pgConfig.Host = internal.PostgresToxiproxyUpstreamHostWithFallback(pgConfig.Host)
pgConfig.SshConfig = &protos.SSHConfig{
Host: "localhost",
Port: uint32(proxyPort),
User: "testuser",
Password: "testpass",
}

var pgPort uint32 = 5432
if envPortStr := os.Getenv("PG_PORT"); envPortStr != "" {
envPort, err := strconv.ParseUint(strings.TrimSpace(strings.Split(envPortStr, "#")[0]), 10, 32)
require.NoError(t, err, "Failed to parse PG_PORT")
pgPort = uint32(envPort)
}

connector, err := NewPostgresConnector(ctx, nil, &protos.PostgresConfig{
Host: pgHost,
Port: pgPort,
User: "postgres",
Password: "postgres",
Database: "postgres",
SshConfig: &protos.SSHConfig{
Host: "localhost",
Port: uint32(proxyPort),
User: "testuser",
Password: "testpass",
},
})
connector, err := NewPostgresConnector(ctx, nil, pgConfig)
require.NoError(t, err)

// Test initial connection works
Expand Down
Loading
Loading