134 lines
3.3 KiB
Go
134 lines
3.3 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
|
|
"github.com/mitchellh/go-testing-interface"
|
|
"github.com/stretchr/testify/require"
|
|
"google.golang.org/grpc"
|
|
|
|
"github.com/hashicorp/vagrant/internal/protocolversion"
|
|
"github.com/hashicorp/vagrant/internal/server/proto/vagrant_server"
|
|
)
|
|
|
|
// 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 vagrant_server.VagrantServer, opts ...TestOption) vagrant_server.VagrantClient {
|
|
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 vagrant_server.NewVagrantClient(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() *vagrant_server.GetVersionInfoResponse {
|
|
return &vagrant_server.GetVersionInfoResponse{
|
|
Info: &vagrant_server.VersionInfo{
|
|
Api: &vagrant_server.VersionInfo_ProtocolVersion{
|
|
Current: 10,
|
|
Minimum: 1,
|
|
},
|
|
|
|
Entrypoint: &vagrant_server.VersionInfo_ProtocolVersion{
|
|
Current: 10,
|
|
Minimum: 1,
|
|
},
|
|
},
|
|
}
|
|
}
|