126 lines
2.7 KiB
Go
126 lines
2.7 KiB
Go
package grpc
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"net"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/hashicorp/consul/agent/grpc/internal/testservice"
|
|
"github.com/hashicorp/consul/agent/metadata"
|
|
"github.com/hashicorp/consul/agent/pool"
|
|
"github.com/stretchr/testify/require"
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
type testServer struct {
|
|
addr net.Addr
|
|
name string
|
|
dc string
|
|
shutdown func()
|
|
}
|
|
|
|
func (s testServer) Metadata() *metadata.Server {
|
|
return &metadata.Server{ID: s.name, Datacenter: s.dc, Addr: s.addr}
|
|
}
|
|
|
|
func newTestServer(t *testing.T, name string, dc string) testServer {
|
|
addr := &net.IPAddr{IP: net.ParseIP("127.0.0.1")}
|
|
handler := NewHandler(addr)
|
|
|
|
testservice.RegisterSimpleServer(handler.srv, &simple{name: name, dc: dc})
|
|
|
|
lis, err := net.Listen("tcp", "127.0.0.1:0")
|
|
require.NoError(t, err)
|
|
|
|
rpc := &fakeRPCListener{t: t, handler: handler}
|
|
|
|
g := errgroup.Group{}
|
|
g.Go(func() error {
|
|
return rpc.listen(lis)
|
|
})
|
|
g.Go(func() error {
|
|
return handler.Run()
|
|
})
|
|
return testServer{
|
|
addr: lis.Addr(),
|
|
name: name,
|
|
dc: dc,
|
|
shutdown: func() {
|
|
if err := lis.Close(); err != nil {
|
|
t.Logf("listener closed with error: %v", err)
|
|
}
|
|
if err := handler.Shutdown(); err != nil {
|
|
t.Logf("grpc server shutdown: %v", err)
|
|
}
|
|
if err := g.Wait(); err != nil {
|
|
t.Logf("grpc server error: %v", err)
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
type simple struct {
|
|
name string
|
|
dc string
|
|
}
|
|
|
|
func (s *simple) Flow(_ *testservice.Req, flow testservice.Simple_FlowServer) error {
|
|
for flow.Context().Err() == nil {
|
|
resp := &testservice.Resp{ServerName: "one", Datacenter: s.dc}
|
|
if err := flow.Send(resp); err != nil {
|
|
return err
|
|
}
|
|
time.Sleep(time.Millisecond)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *simple) Something(_ context.Context, _ *testservice.Req) (*testservice.Resp, error) {
|
|
return &testservice.Resp{ServerName: s.name, Datacenter: s.dc}, nil
|
|
}
|
|
|
|
// fakeRPCListener mimics agent/consul.Server.listen to handle the RPCType byte.
|
|
// In the future we should be able to refactor Server and extract this RPC
|
|
// handling logic so that we don't need to use a fake.
|
|
// For now, since this logic is in agent/consul, we can't easily use Server.listen
|
|
// so we fake it.
|
|
type fakeRPCListener struct {
|
|
t *testing.T
|
|
handler *Handler
|
|
}
|
|
|
|
func (f *fakeRPCListener) listen(listener net.Listener) error {
|
|
for {
|
|
conn, err := listener.Accept()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
go f.handleConn(conn)
|
|
}
|
|
}
|
|
|
|
func (f *fakeRPCListener) handleConn(conn net.Conn) {
|
|
buf := make([]byte, 1)
|
|
|
|
if _, err := conn.Read(buf); err != nil {
|
|
if err != io.EOF {
|
|
fmt.Println("ERROR", err.Error())
|
|
}
|
|
conn.Close()
|
|
return
|
|
}
|
|
typ := pool.RPCType(buf[0])
|
|
|
|
if typ == pool.RPCGRPC {
|
|
f.handler.Handle(conn)
|
|
return
|
|
}
|
|
|
|
fmt.Println("ERROR: unexpected byte", typ)
|
|
conn.Close()
|
|
}
|