squirrel/itests/db_test.go at master · n-r-w/squirrel · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
//go:build itest

package itests

import (
	"context"
	"testing"

	"github.com/georgysavva/scany/v2/pgxscan"
	"github.com/jackc/pgx/v5/pgxpool"
	sq "github.com/n-r-w/squirrel"
	"github.com/n-r-w/testdock/v2"
	"github.com/stretchr/testify/require"
)

func newTestPool(t *testing.T) (*pgxpool.Pool, context.Context) {
	t.Helper()

	ctx := context.Background()
	pool, _ := testdock.GetPgxPool(t, testdock.DefaultPostgresDSN)

	return pool, ctx
}

func execSetup(t *testing.T, pool *pgxpool.Pool, ctx context.Context, setupSQL string) {
	t.Helper()

	_, err := pool.Exec(ctx, setupSQL)
	require.NoError(t, err)
}

func queryInt64s(t *testing.T, pool *pgxpool.Pool, ctx context.Context, q sq.Sqlizer) []int64 {
	t.Helper()

	sql, args, err := q.ToSql()
	require.NoError(t, err)

	var ids []int64
	err = pgxscan.Select(ctx, pool, &ids, sql, args...)
	require.NoError(t, err)

	return ids
}

func queryInt64StringPairs(t *testing.T, pool *pgxpool.Pool, ctx context.Context, q sq.Sqlizer) ([]int64, []string) {
	t.Helper()

	sql, args, err := q.ToSql()
	require.NoError(t, err)

	type idName struct {
		ID   int64  `db:"id"`
		Name string `db:"name"`
	}

	var rows []idName
	err = pgxscan.Select(ctx, pool, &rows, sql, args...)
	require.NoError(t, err)

	ids := make([]int64, 0, len(rows))
	names := make([]string, 0, len(rows))
	for _, row := range rows {
		ids = append(ids, row.ID)
		names = append(names, row.Name)
	}

	return ids, names
}