This repository has been archived by the owner on Jan 8, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 330
/
testing.go
133 lines (110 loc) · 3.21 KB
/
testing.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
121
122
123
124
125
126
127
128
129
130
131
132
133
package server
import (
"context"
"net"
"github.com/mitchellh/go-testing-interface"
"github.com/stretchr/testify/require"
"google.golang.org/grpc"
"github.com/hashicorp/waypoint/internal/protocolversion"
pb "github.com/hashicorp/waypoint/internal/server/gen"
)
// TestServer starts a server and returns a gRPC client to that server.
// We use t.Cleanup to ensure resources are automatically cleaned up.
func TestServer(t testing.T, impl pb.WaypointServer, opts ...TestOption) pb.WaypointClient {
require := require.New(t)
c := testConfig{
ctx: context.Background(),
}
for _, opt := range opts {
opt(&c)
}
// Listen on a random port
ln, err := net.Listen("tcp", "127.0.0.1:")
require.NoError(err)
t.Cleanup(func() { ln.Close() })
// We make run a function since we'll call it to restart too
run := func(ctx context.Context) context.CancelFunc {
ctx, cancel := context.WithCancel(ctx)
go Run(
WithContext(ctx),
WithGRPC(ln),
WithImpl(impl),
)
t.Cleanup(func() { cancel() })
return cancel
}
// Create the server
cancel := run(c.ctx)
// If we have a restart channel, then listen to that for restarts.
if c.restartCh != nil {
doneCh := make(chan struct{})
t.Cleanup(func() { close(doneCh) })
go func() {
for {
select {
case <-c.restartCh:
// Cancel the old context
cancel()
// This can fail, but it probably won't. Can't think of
// a cleaner way since gRPC force closes its listener.
ln, err = net.Listen("tcp", ln.Addr().String())
if err != nil {
return
}
// Create a new one
cancel = run(context.Background())
case <-doneCh:
return
}
}
}()
}
// Get our version info we'll set on the client
vsnInfo := testVersionInfoResponse().Info
// Connect, this should retry in the case Run is not going yet
conn, err := grpc.DialContext(context.Background(), ln.Addr().String(),
grpc.WithBlock(),
grpc.WithInsecure(),
grpc.WithUnaryInterceptor(protocolversion.UnaryClientInterceptor(vsnInfo)),
grpc.WithStreamInterceptor(protocolversion.StreamClientInterceptor(vsnInfo)),
)
require.NoError(err)
t.Cleanup(func() { conn.Close() })
return pb.NewWaypointClient(conn)
}
// TestOption is used with TestServer to configure test behavior.
type TestOption func(*testConfig)
type testConfig struct {
ctx context.Context
restartCh <-chan struct{}
}
// TestWithContext specifies a context to use with the test server. When
// this is done then the server will exit.
func TestWithContext(ctx context.Context) TestOption {
return func(c *testConfig) {
c.ctx = ctx
}
}
// TestWithRestart specifies a channel that will be sent to to trigger
// a restart. The restart happens asynchronously. If you want to ensure the
// server is shutdown first, use TestWithContext, shut it down, wait for
// errors on the API, then restart.
func TestWithRestart(ch <-chan struct{}) TestOption {
return func(c *testConfig) {
c.restartCh = ch
}
}
func testVersionInfoResponse() *pb.GetVersionInfoResponse {
return &pb.GetVersionInfoResponse{
Info: &pb.VersionInfo{
Api: &pb.VersionInfo_ProtocolVersion{
Current: 10,
Minimum: 1,
},
Entrypoint: &pb.VersionInfo_ProtocolVersion{
Current: 10,
Minimum: 1,
},
},
}
}