1  //go:build !js && !(windows && (arm || 386)) && !(linux && (ppc64 || mips || mipsle || mips64)) && !(netbsd || openbsd)
  2  
  3  package sqldb
  4  
  5  import (
  6  	"context"
  7  	"crypto/rand"
  8  	"database/sql"
  9  	"encoding/hex"
 10  	"fmt"
 11  	"strconv"
 12  	"strings"
 13  	"testing"
 14  	"time"
 15  
 16  	_ "github.com/jackc/pgx/v5"
 17  	"github.com/ory/dockertest/v3"
 18  	"github.com/ory/dockertest/v3/docker"
 19  	"github.com/stretchr/testify/require"
 20  )
 21  
 22  const (
 23  	testPgUser   = "test"
 24  	testPgPass   = "test"
 25  	testPgDBName = "test"
 26  	PostgresTag  = "11"
 27  )
 28  
 29  // TestPgFixture is a test fixture that starts a Postgres 11 instance in a
 30  // docker container.
 31  type TestPgFixture struct {
 32  	db       *sql.DB
 33  	pool     *dockertest.Pool
 34  	resource *dockertest.Resource
 35  	host     string
 36  	port     int
 37  }
 38  
 39  // NewTestPgFixture constructs a new TestPgFixture starting up a docker
 40  // container running Postgres 11. The started container will expire in after
 41  // the passed duration.
 42  func NewTestPgFixture(t testing.TB, expiry time.Duration) *TestPgFixture {
 43  	// Use a sensible default on Windows (tcp/http) and linux/osx (socket)
 44  	// by specifying an empty endpoint.
 45  	pool, err := dockertest.NewPool("")
 46  	require.NoError(t, err, "Could not connect to docker")
 47  
 48  	// Pulls an image, creates a container based on it and runs it.
 49  	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
 50  		Repository: "postgres",
 51  		Tag:        PostgresTag,
 52  		Env: []string{
 53  			fmt.Sprintf("POSTGRES_USER=%v", testPgUser),
 54  			fmt.Sprintf("POSTGRES_PASSWORD=%v", testPgPass),
 55  			fmt.Sprintf("POSTGRES_DB=%v", testPgDBName),
 56  			"listen_addresses='*'",
 57  		},
 58  		Cmd: []string{
 59  			"postgres",
 60  			"-c", "log_statement=all",
 61  			"-c", "log_destination=stderr",
 62  			"-c", "max_connections=5000",
 63  		},
 64  	}, func(config *docker.HostConfig) {
 65  		// Set AutoRemove to true so that stopped container goes away
 66  		// by itself.
 67  		config.AutoRemove = true
 68  		config.RestartPolicy = docker.RestartPolicy{Name: "no"}
 69  	})
 70  	require.NoError(t, err, "Could not start resource")
 71  
 72  	hostAndPort := resource.GetHostPort("5432/tcp")
 73  	parts := strings.Split(hostAndPort, ":")
 74  	host := parts[0]
 75  	port, err := strconv.ParseInt(parts[1], 10, 64)
 76  	require.NoError(t, err)
 77  
 78  	fixture := &TestPgFixture{
 79  		host: host,
 80  		port: int(port),
 81  	}
 82  	databaseURL := fixture.GetConfig(testPgDBName).Dsn
 83  	log.Infof("Connecting to Postgres fixture: %v\n", databaseURL)
 84  
 85  	// Tell docker to hard kill the container in "expiry" seconds.
 86  	require.NoError(t, resource.Expire(uint(expiry.Seconds())))
 87  
 88  	// Exponential backoff-retry, because the application in the container
 89  	// might not be ready to accept connections yet.
 90  	pool.MaxWait = 120 * time.Second
 91  
 92  	var testDB *sql.DB
 93  	err = pool.Retry(func() error {
 94  		testDB, err = sql.Open("pgx", databaseURL)
 95  		if err != nil {
 96  			return err
 97  		}
 98  
 99  		return testDB.Ping()
100  	})
101  	require.NoError(t, err, "Could not connect to docker")
102  
103  	// Now fill in the rest of the fixture.
104  	fixture.db = testDB
105  	fixture.pool = pool
106  	fixture.resource = resource
107  
108  	return fixture
109  }
110  
111  // GetConfig returns the full config of the Postgres node.
112  func (f *TestPgFixture) GetConfig(dbName string) *PostgresConfig {
113  	return &PostgresConfig{
114  		Dsn: fmt.Sprintf(
115  			"postgres://%v:%v@%v:%v/%v?sslmode=disable",
116  			testPgUser, testPgPass, f.host, f.port, dbName,
117  		),
118  	}
119  }
120  
121  // TearDown stops the underlying docker container.
122  func (f *TestPgFixture) TearDown(t testing.TB) {
123  	err := f.pool.Purge(f.resource)
124  	require.NoError(t, err, "Could not purge resource")
125  }
126  
127  // randomDBName generates a random database name.
128  func randomDBName(t testing.TB) string {
129  	randBytes := make([]byte, 8)
130  	_, err := rand.Read(randBytes)
131  	require.NoError(t, err)
132  
133  	return "test_" + hex.EncodeToString(randBytes)
134  }
135  
136  // NewTestPostgresDB is a helper function that creates a Postgres database for
137  // testing using the given fixture.
138  func NewTestPostgresDB(t testing.TB, fixture *TestPgFixture) *PostgresStore {
139  	t.Helper()
140  
141  	dbName := randomDBName(t)
142  
143  	t.Logf("Creating new Postgres DB '%s' for testing", dbName)
144  
145  	_, err := fixture.db.ExecContext(
146  		context.Background(), "CREATE DATABASE "+dbName,
147  	)
148  	require.NoError(t, err)
149  
150  	cfg := fixture.GetConfig(dbName)
151  	store, err := NewPostgresStore(cfg)
152  	require.NoError(t, err)
153  
154  	require.NoError(t, store.ApplyAllMigrations(
155  		context.Background(), GetMigrations()),
156  	)
157  
158  	t.Cleanup(func() {
159  		require.NoError(t, store.DB.Close())
160  	})
161  
162  	return store
163  }
164  
165  // NewTestPostgresDBWithVersion is a helper function that creates a Postgres
166  // database for testing and migrates it to the given version.
167  func NewTestPostgresDBWithVersion(t *testing.T, fixture *TestPgFixture,
168  	version uint) *PostgresStore {
169  
170  	t.Helper()
171  
172  	t.Logf("Creating new Postgres DB for testing, migrating to version %d",
173  		version)
174  
175  	dbName := randomDBName(t)
176  	_, err := fixture.db.ExecContext(
177  		context.Background(), "CREATE DATABASE "+dbName,
178  	)
179  	require.NoError(t, err)
180  
181  	storeCfg := fixture.GetConfig(dbName)
182  	storeCfg.SkipMigrations = true
183  	store, err := NewPostgresStore(storeCfg)
184  	require.NoError(t, err)
185  
186  	err = store.ExecuteMigrations(TargetVersion(version))
187  	require.NoError(t, err)
188  
189  	t.Cleanup(func() {
190  		require.NoError(t, store.DB.Close())
191  	})
192  
193  	return store
194  }