package postgres
import (
    "context"
    "time"
    "github.com/testcontainers/testcontainers-go"
    "github.com/testcontainers/testcontainers-go/wait"
)
// postgresContainer represents the postgres container type used in the module
type postgresContainer struct {
    testcontainers.Container
}
type postgresContainerOption func(req *testcontainers.ContainerRequest)
func WithWaitStrategy(strategies ...wait.Strategy) func(req *testcontainers.ContainerRequest) {
    return func(req *testcontainers.ContainerRequest) {
        req.WaitingFor = wait.ForAll(strategies...).WithDeadline(1 * time.Minute)
    }
}
func WithPort(port string) func(req *testcontainers.ContainerRequest) {
    return func(req *testcontainers.ContainerRequest) {
        req.ExposedPorts = append(req.ExposedPorts, port)
    }
}
func WithInitialDatabase(user string, password string, dbName string) func(req *testcontainers.ContainerRequest) {
    return func(req *testcontainers.ContainerRequest) {
        req.Env["POSTGRES_USER"] = user
        req.Env["POSTGRES_PASSWORD"] = password
        req.Env["POSTGRES_DB"] = dbName
    }
}
// setupPostgres creates an instance of the postgres container type
func setupPostgres(ctx context.Context, opts ...postgresContainerOption) (*postgresContainer, error) {
    req := testcontainers.ContainerRequest{
        Image:        "postgres:11-alpine",
        Env:          map[string]string{},
        ExposedPorts: []string{},
        Cmd:          []string{"postgres", "-c", "fsync=off"},
    }
    for _, opt := range opts {
        opt(&req)
    }
    container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
        ContainerRequest: req,
        Started:          true,
    })
    if err != nil {
        return nil, err
    }
    return &postgresContainer{Container: container}, nil
}
 
package postgres
import (
    "context"
    "database/sql"
    "fmt"
    "testing"
    "time"
    "github.com/docker/go-connections/nat"
    _ "github.com/lib/pq"
    "github.com/stretchr/testify/assert"
    "github.com/stretchr/testify/require"
    "github.com/testcontainers/testcontainers-go/wait"
)
func TestPostgres(t *testing.T) {
    ctx := context.Background()
    const dbname = "test-db"
    const user = "postgres"
    const password = "password"
    port, err := nat.NewPort("tcp", "5432")
    require.NoError(t, err)
    container, err := setupPostgres(ctx,
        WithPort(port.Port()),
        WithInitialDatabase(user, password, dbname),
        WithWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5*time.Second)),
    )
    if err != nil {
        t.Fatal(err)
    }
    // Clean up the container after the test is complete
    t.Cleanup(func() {
        if err := container.Terminate(ctx); err != nil {
            t.Fatalf("failed to terminate container: %s", err)
        }
    })
    containerPort, err := container.MappedPort(ctx, port)
    assert.NoError(t, err)
    host, err := container.Host(ctx)
    assert.NoError(t, err)
    connStr := fmt.Sprintf("host=%s port=%s user=%s password=%s dbname=%s sslmode=disable", host, containerPort.Port(), user, password, dbname)
    // perform assertions
    db, err := sql.Open("postgres", connStr)
    assert.NoError(t, err)
    assert.NotNil(t, db)
    defer db.Close()
    result, err := db.Exec("CREATE TABLE IF NOT EXISTS test (id int, name varchar(255));")
    assert.NoError(t, err)
    assert.NotNil(t, result)
    result, err = db.Exec("INSERT INTO test (id, name) VALUES (1, 'test');")
    assert.NoError(t, err)
    assert.NotNil(t, result)
}
func TestContainerWithWaitForSQL(t *testing.T) {
    const dbname = "test-db"
    const user = "postgres"
    const password = "password"
    ctx := context.Background()
    var port = "5432/tcp"
    dbURL := func(host string, port nat.Port) string {
        return fmt.Sprintf("postgres://postgres:password@%s:%s/%s?sslmode=disable", host, port.Port(), dbname)
    }
    t.Run("default query", func(t *testing.T) {
        container, err := setupPostgres(ctx, WithPort(port), WithInitialDatabase("postgres", "password", dbname), WithWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL)))
        require.NoError(t, err)
        require.NotNil(t, container)
    })
    t.Run("custom query", func(t *testing.T) {
        container, err := setupPostgres(
            ctx,
            WithPort(port),
            WithInitialDatabase(user, password, dbname),
            WithWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL).WithStartupTimeout(time.Second*5).WithQuery("SELECT 10")),
        )
        require.NoError(t, err)
        require.NotNil(t, container)
    })
    t.Run("custom bad query", func(t *testing.T) {
        container, err := setupPostgres(
            ctx,
            WithPort(port),
            WithInitialDatabase(user, password, dbname),
            WithWaitStrategy(wait.ForSQL(nat.Port(port), "postgres", dbURL).WithStartupTimeout(time.Second*5).WithQuery("SELECT 'a' from b")),
        )
        require.Error(t, err)
        require.Nil(t, container)
    })
}