-
Notifications
You must be signed in to change notification settings - Fork 19
/
setup.go
120 lines (97 loc) · 2.76 KB
/
setup.go
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
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package databasetest
import (
"bytes"
"context"
"fmt"
"io/fs"
"math/rand"
"net"
"os"
"path/filepath"
"testing"
"time"
"github.com/cenkalti/backoff/v4"
"github.com/jackc/pgx/v4"
"code.vegaprotocol.io/vega/datanode/sqlstore"
"code.vegaprotocol.io/vega/logging"
)
var (
sqlTestsEnabled = true
minPort = 30000
maxPort = 40000
postgresServerTimeout = time.Second * 10
)
func TestMain(m *testing.M, onSetupComplete func(sqlstore.Config, *sqlstore.ConnectionSource, *bytes.Buffer),
postgresRuntimePath string, sqlFs fs.FS,
) int {
testDBSocketDir := filepath.Join(postgresRuntimePath)
testDBPort := 5432
testDBHost := ""
sqlConfig := NewTestConfig(testDBPort, testDBHost, testDBSocketDir)
if sqlTestsEnabled {
log := logging.NewTestLogger()
err := os.Mkdir(postgresRuntimePath, fs.ModePerm)
if err != nil {
panic(err)
}
defer os.RemoveAll(postgresRuntimePath)
postgresLog := &bytes.Buffer{}
embeddedPostgres, err := sqlstore.StartEmbeddedPostgres(log, sqlConfig, postgresRuntimePath, postgresLog)
if err != nil {
log.Errorf("failed to start postgres: %s", postgresLog.String())
panic(err)
}
log.Infof("Test DB Port: %d", testDBPort)
// Make sure the database has started before we run the tests.
ctx, cancel := context.WithTimeout(context.Background(), postgresServerTimeout)
op := func() error {
connStr := sqlConfig.ConnectionConfig.GetConnectionString()
conn, err := pgx.Connect(ctx, connStr)
if err != nil {
return err
}
return conn.Ping(ctx)
}
if err := backoff.Retry(op, backoff.NewExponentialBackOff()); err != nil {
cancel()
panic(err)
}
cancel()
connectionSource, err := sqlstore.NewTransactionalConnectionSource(log, sqlConfig.ConnectionConfig)
if err != nil {
panic(err)
}
defer embeddedPostgres.Stop()
if err = sqlstore.WipeDatabaseAndMigrateSchemaToLatestVersion(log, sqlConfig.ConnectionConfig, sqlFs, false); err != nil {
panic(err)
}
if err = sqlstore.ApplyDataRetentionPolicies(sqlConfig, log); err != nil {
panic(err)
}
onSetupComplete(sqlConfig, connectionSource, postgresLog)
return m.Run()
}
return 0
}
func NewTestConfig(port int, host, socketDir string) sqlstore.Config {
sqlConfig := sqlstore.NewDefaultConfig()
sqlConfig.UseEmbedded = true
sqlConfig.ConnectionConfig.Port = port
sqlConfig.ConnectionConfig.Host = host
sqlConfig.ConnectionConfig.SocketDir = socketDir
return sqlConfig
}
func GetNextFreePort() int {
rand.Seed(time.Now().UnixNano())
for {
port := rand.Intn(maxPort-minPort+1) + minPort
timeout := time.Millisecond * 100
conn, err := net.DialTimeout("tcp", net.JoinHostPort("localhost", fmt.Sprintf("%d", port)), timeout)
if err != nil {
return port
}
if conn != nil {
conn.Close()
}
}
}