/
testsuite.go
99 lines (86 loc) · 2.73 KB
/
testsuite.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
package testsuite
import (
"context"
"fmt"
"net"
"sync"
"time"
"github.com/prometheus/client_golang/api"
apiv1 "github.com/prometheus/client_golang/api/prometheus/v1"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
)
// Suite holds a dbtest.Suite and a lazily started prometheus server.
// Each Suite holds its own Prometheus server. Suites can be run in parallel.
type Suite struct {
suite.Suite
// promMutex guards all prom* variables below.
promMutex sync.Mutex
promAddr string
promCtx context.Context
promWaitShutdown func() error
promCancelCtx context.CancelFunc
}
// PrometheusURL starts a prometheus server and returns the api url to it.
// The started prometheus is shared in a testsuite.
func (ts *Suite) PrometheusURL() string {
ts.promMutex.Lock()
defer ts.promMutex.Unlock()
if ts.promAddr != "" {
return ts.promAddr
}
port, err := getFreePort()
require.NoError(ts.T(), err)
url := fmt.Sprintf("http://127.0.0.1:%d", port)
ts.promCtx, ts.promCancelCtx = context.WithCancel(context.Background())
ts.Require().FileExists(PromBin, "Prometheus binary is required to run tests. Download with `make ensure-prometheus`")
wait, err := StartPrometheus(ts.promCtx, port)
require.NoError(ts.T(), err)
ts.promWaitShutdown = wait
require.NoError(ts.T(), waitForPrometheusReady(ts.promCtx, url, 5*time.Second))
ts.promAddr = url
return url
}
// PrometheusAPIClient starts a prometheus server and returns a client to it.
// The started prometheus is shared in a testsuite.
func (ts *Suite) PrometheusAPIClient() apiv1.API {
client, err := api.NewClient(api.Config{Address: ts.PrometheusURL()})
require.NoError(ts.T(), err)
return apiv1.NewAPI(client)
}
// TearDownSuite stops prometheus and drops the temporary database.
func (ts *Suite) TearDownSuite() {
ts.promMutex.Lock()
defer ts.promMutex.Unlock()
if ts.promCancelCtx != nil {
defer ts.promWaitShutdown()
ts.promCancelCtx()
}
}
func waitForPrometheusReady(ctx context.Context, url string, timeout time.Duration) error {
client, err := api.NewClient(api.Config{Address: url})
if err != nil {
return fmt.Errorf("error creating prometheus client: %w", err)
}
clientv1 := apiv1.NewAPI(client)
timeoutCtx, timeoutCancel := context.WithTimeout(ctx, timeout)
defer timeoutCancel()
for {
_, err := clientv1.Runtimeinfo(timeoutCtx)
if err == nil {
break
} else if timeoutCtx.Err() != nil {
return fmt.Errorf("failed waiting for prometheus to become ready, last error: %w", err)
}
time.Sleep(50 * time.Millisecond)
}
return nil
}
func getFreePort() (port int, err error) {
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
return 0, err
}
defer l.Close()
return l.Addr().(*net.TCPAddr).Port, nil
}