Golang测试数据库操作 使用测试容器方案

testcontainers-go

package main

import (
    "context"
    "database/sql"
    "fmt"
    "log"
    "testing"
    "time"

    _ "github.com/lib/pq" // PostgreSQL driver
    "github.com/testcontainers/testcontainers-go"
    "github.com/testcontainers/testcontainers-go/wait"
)

// DBContainer represents our database test container
type DBContainer struct {
    testcontainers.Container
    ConnectionString string
}

// setupDBContainer starts a PostgreSQL container and returns its connection string
func setupDBContainer(ctx context.Context) (*DBContainer, error) {
    req := testcontainers.ContainerRequest{
        Image:        "postgres:13",
        ExposedPorts: []string{"5432/tcp"},
        Env: map[string]string{
            "POSTGRES_DB":       "testdb",
            "POSTGRES_USER":     "user",
            "POSTGRES_PASSWORD": "password",
        },
        WaitingFor: wait.ForLog("database system is ready to accept connections").WithStartupTimeout(120 * time.Second),
    }

    container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
        ContainerRequest: req,
        Started:          true,
    })
    if err != nil {
        return nil, fmt.Errorf("failed to start container: %w", err)
    }

    hostIP, err := container.Host(ctx)
    if err != nil {
        return nil, fmt.Errorf("failed to get host IP: %w", err)
    }
    port, err := container.MappedPort(ctx, "5432/tcp")
    if err != nil {
        return nil, fmt.Errorf("failed to get mapped port: %w", err)
    }

    connStr := fmt.Sprintf("host=%s port=%s user=user password=password dbname=testdb sslmode=disable", hostIP, port.Port())

    // A small hack: sometimes the driver needs a tiny bit more time after the log message
    time.Sleep(1 * time.Second) 

    return &DBContainer{Container: container, ConnectionString: connStr}, nil
}

// TestMain is crucial for managing the container lifecycle
func TestMain(m *testing.M) {
    ctx := context.Background()
    dbContainer, err := setupDBContainer(ctx)
    if err != nil {
        log.Fatalf("Could not setup database container: %v", err)
    }

    // Set the connection string globally or pass it around
    // For simplicity, let's assume a global variable or a test helper passes it
    // In a real project, you'd likely inject this into your test suite
    testDBConnString = dbContainer.ConnectionString 

    code := m.Run()

    // Clean up the container after all tests are done
    if err := dbContainer.Terminate(ctx); err != nil {
        log.Fatalf("Could not terminate database container: %v", err)
    }
    // os.Exit(code) // This is handled by m.Run()
}

var testDBConnString string // Global for example, better to pass via context/struct

// Example test function
func TestDatabaseOperation(t *testing.T) {
    if testDBConnString == "" {
        t.Fatal("Database connection string not set up by TestMain")
    }

    db, err := sql.Open("postgres", testDBConnString)
    if err != nil {
        t.Fatalf("Failed to connect to DB: %v", err)
    }
    defer db.Close()

    err = db.Ping()
    if err != nil {
        t.Fatalf("Failed to ping DB: %v", err)
    }
    t.Log("Successfully connected and pinged the database.")

    // Now, perform