/
testclient.go
39 lines (36 loc) · 1 KB
/
testclient.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
package grpcutil
import (
"context"
"net"
"testing"
"github.com/pachyderm/pachyderm/v2/src/internal/errors"
"github.com/stretchr/testify/require"
"golang.org/x/sync/errgroup"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
"google.golang.org/grpc/test/bufconn"
)
func NewTestClient(t testing.TB, regFunc func(*grpc.Server)) *grpc.ClientConn {
ctx := context.Background()
eg := errgroup.Group{}
gserv := grpc.NewServer()
listener := bufconn.Listen(1 << 20)
regFunc(gserv)
eg.Go(func() error {
return errors.EnsureStack(gserv.Serve(listener))
})
gconn, err := grpc.DialContext(ctx, "", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
res, err := listener.Dial()
return res, errors.EnsureStack(err)
}), grpc.WithTransportCredentials(insecure.NewCredentials()))
require.NoError(t, err)
t.Cleanup(func() {
gserv.GracefulStop()
err := eg.Wait()
if errors.Is(err, grpc.ErrServerStopped) {
err = nil
}
require.NoError(t, err)
})
return gconn
}