1// Package pgtest starts a postgres with pico's schema for tests.
  2package pgtest
  3
  4import (
  5	"context"
  6	"errors"
  7	"fmt"
  8	"net/url"
  9	"os"
 10	"os/exec"
 11	"path/filepath"
 12	"runtime"
 13	"sort"
 14	"sync"
 15	"testing"
 16	"time"
 17
 18	"github.com/jmoiron/sqlx"
 19	_ "github.com/lib/pq"
 20	"github.com/picosh/pico/pkg/db"
 21	"github.com/testcontainers/testcontainers-go"
 22	"github.com/testcontainers/testcontainers-go/modules/postgres"
 23	"github.com/testcontainers/testcontainers-go/wait"
 24)
 25
 26var (
 27	once     sync.Once
 28	dbURL    string
 29	stop     = func() {}
 30	startErr error
 31)
 32
 33// URL is the connection string of a postgres with pico's schema, started
 34// the first time a test in the package asks for one. It is the database
 35// TEST_DATABASE_URL names, or a container. A test skips when there is
 36// neither, and fails when REQUIRE_TESTCONTAINERS=1.
 37func URL(t testing.TB) string {
 38	t.Helper()
 39	once.Do(func() {
 40		var started func()
 41		dbURL, started, startErr = start(context.Background())
 42		if startErr == nil {
 43			stop = started
 44		}
 45	})
 46	if startErr != nil {
 47		if os.Getenv("REQUIRE_TESTCONTAINERS") == "1" {
 48			t.Fatalf("postgres: %v", startErr)
 49		}
 50		t.Skipf("postgres unavailable, set TEST_DATABASE_URL or start docker: %v", startErr)
 51	}
 52	return dbURL
 53}
 54
 55// Main runs a package's tests and stops the postgres they started.
 56func Main(m *testing.M) {
 57	code := m.Run()
 58	stop()
 59	os.Exit(code)
 60}
 61
 62// Clean deletes every row the tests wrote.
 63func Clean(t testing.TB, d *sqlx.DB) {
 64	t.Helper()
 65	tables := []string{
 66		"access_logs", "tuns_event_logs", "analytics_visits",
 67		"feed_items", "post_aliases", "post_tags", "posts",
 68		"projects", "feature_flags", "payment_history", "tokens",
 69		"public_keys", "pipe_monitors", "app_users",
 70	}
 71	for _, table := range tables {
 72		if _, err := d.Exec(fmt.Sprintf("DELETE FROM %s", table)); err != nil {
 73			t.Fatalf("failed to clean up %s: %v", table, err)
 74		}
 75	}
 76}
 77
 78// InsertPost inserts post, publishing and updating it now unless it says
 79// otherwise.
 80func InsertPost(t testing.TB, d interface {
 81	InsertPost(post *db.Post) (*db.Post, error)
 82}, post *db.Post) *db.Post {
 83	t.Helper()
 84	now := time.Now()
 85	if post.UpdatedAt == nil {
 86		post.UpdatedAt = &now
 87	}
 88	if post.PublishAt == nil {
 89		post.PublishAt = &now
 90	}
 91	created, err := d.InsertPost(post)
 92	if err != nil {
 93		t.Fatalf("InsertPost failed: %v", err)
 94	}
 95	return created
 96}
 97
 98func start(ctx context.Context) (string, func(), error) {
 99	// Packages test in parallel, so each gets its own database.
100	if base := os.Getenv("TEST_DATABASE_URL"); base != "" {
101		return createDatabase(base)
102	}
103
104	if err := exec.Command("docker", "info").Run(); err != nil {
105		return "", nil, fmt.Errorf("docker: %w", err)
106	}
107	container, err := postgres.Run(ctx,
108		"postgres:14",
109		postgres.WithDatabase("pico_test"),
110		postgres.WithUsername("postgres"),
111		postgres.WithPassword("postgres"),
112		testcontainers.WithWaitStrategy(
113			wait.ForLog("database system is ready to accept connections").
114				WithOccurrence(2).
115				WithStartupTimeout(30*time.Second)),
116	)
117	if err != nil {
118		return "", nil, err
119	}
120	terminate := func() { _ = container.Terminate(ctx) }
121	connStr, err := container.ConnectionString(ctx, "sslmode=disable")
122	if err == nil {
123		err = migrate(connStr)
124	}
125	if err != nil {
126		terminate()
127		return "", nil, err
128	}
129	return connStr, terminate, nil
130}
131
132func createDatabase(base string) (string, func(), error) {
133	server, err := sqlx.Connect("postgres", base)
134	if err != nil {
135		return "", nil, err
136	}
137	name := fmt.Sprintf("pico_test_%d_%d", os.Getpid(), time.Now().UnixNano())
138	if _, err := server.Exec("CREATE DATABASE " + name); err != nil {
139		_ = server.Close()
140		return "", nil, err
141	}
142	drop := func() {
143		_, _ = server.Exec("DROP DATABASE " + name + " WITH (FORCE)")
144		_ = server.Close()
145	}
146	u, err := url.Parse(base)
147	if err == nil {
148		u.Path = "/" + name
149		err = migrate(u.String())
150	}
151	if err != nil {
152		drop()
153		return "", nil, err
154	}
155	return u.String(), drop, nil
156}
157
158func migrate(connStr string) error {
159	d, err := sqlx.Connect("postgres", connStr)
160	if err != nil {
161		return err
162	}
163	defer func() { _ = d.Close() }()
164
165	_, filename, _, ok := runtime.Caller(0)
166	if !ok {
167		return errors.New("failed to get current file path")
168	}
169	migrationsDir := filepath.Join(filepath.Dir(filename), "..", "..", "..", "..", "sql", "migrations")
170	entries, err := os.ReadDir(migrationsDir)
171	if err != nil {
172		return fmt.Errorf("failed to read migrations directory: %w", err)
173	}
174	// They're date-prefixed, so alphabetical order is correct.
175	var migrationFiles []string
176	for _, entry := range entries {
177		if !entry.IsDir() && filepath.Ext(entry.Name()) == ".sql" {
178			migrationFiles = append(migrationFiles, entry.Name())
179		}
180	}
181	sort.Strings(migrationFiles)
182	for _, name := range migrationFiles {
183		content, err := os.ReadFile(filepath.Join(migrationsDir, name))
184		if err != nil {
185			return fmt.Errorf("failed to read migration %s: %w", name, err)
186		}
187		if _, err := d.Exec(string(content)); err != nil {
188			return fmt.Errorf("failed to execute migration %s: %w", name, err)
189		}
190	}
191	return nil
192}