-
Notifications
You must be signed in to change notification settings - Fork 4.3k
/
fake_rls_server.go
134 lines (120 loc) · 4.35 KB
/
fake_rls_server.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
134
/*
*
* Copyright 2022 gRPC authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*
*/
// Package rls contains utilities for RouteLookupService e2e tests.
package rls
import (
"context"
"net"
"sync"
"testing"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
rlsgrpc "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
"google.golang.org/grpc/internal/testutils"
"google.golang.org/grpc/status"
)
// RouteLookupResponse wraps an RLS response and the associated error to be sent
// to a client when the RouteLookup RPC is invoked.
type RouteLookupResponse struct {
Resp *rlspb.RouteLookupResponse
Err error
}
// SetupFakeRLSServer starts and returns a fake RouteLookupService server
// listening on the given listener or on a random local port. Also returns a
// channel for tests to get notified whenever the RouteLookup RPC is invoked on
// the fake server.
//
// This function sets up the fake server to respond with an empty response for
// the RouteLookup RPCs. Tests can override this by calling the
// SetResponseCallback() method on the returned fake server.
func SetupFakeRLSServer(t *testing.T, lis net.Listener, opts ...grpc.ServerOption) (*FakeRouteLookupServer, chan struct{}) {
s, cancel := StartFakeRouteLookupServer(t, lis, opts...)
t.Logf("Started fake RLS server at %q", s.Address)
ch := make(chan struct{}, 1)
s.SetRequestCallback(func(request *rlspb.RouteLookupRequest) {
select {
case ch <- struct{}{}:
default:
}
})
t.Cleanup(cancel)
return s, ch
}
// FakeRouteLookupServer is a fake implementation of the RouteLookupService.
//
// It is safe for concurrent use.
type FakeRouteLookupServer struct {
rlsgrpc.UnimplementedRouteLookupServiceServer
Address string
mu sync.Mutex
respCb func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse
reqCb func(*rlspb.RouteLookupRequest)
}
// StartFakeRouteLookupServer starts a fake RLS server listening for requests on
// lis. If lis is nil, it creates a new listener on a random local port. The
// returned cancel function should be invoked by the caller upon completion of
// the test.
func StartFakeRouteLookupServer(t *testing.T, lis net.Listener, opts ...grpc.ServerOption) (*FakeRouteLookupServer, func()) {
t.Helper()
if lis == nil {
var err error
lis, err = testutils.LocalTCPListener()
if err != nil {
t.Fatalf("net.Listen() failed: %v", err)
}
}
s := &FakeRouteLookupServer{Address: lis.Addr().String()}
server := grpc.NewServer(opts...)
rlsgrpc.RegisterRouteLookupServiceServer(server, s)
go server.Serve(lis)
return s, func() { server.Stop() }
}
// RouteLookup implements the RouteLookupService.
func (s *FakeRouteLookupServer) RouteLookup(ctx context.Context, req *rlspb.RouteLookupRequest) (*rlspb.RouteLookupResponse, error) {
s.mu.Lock()
defer s.mu.Unlock()
if s.reqCb != nil {
s.reqCb(req)
}
if err := ctx.Err(); err != nil {
return nil, status.Error(codes.DeadlineExceeded, err.Error())
}
if s.respCb == nil {
return &rlspb.RouteLookupResponse{}, nil
}
resp := s.respCb(ctx, req)
return resp.Resp, resp.Err
}
// SetResponseCallback sets a callback to be invoked on every RLS request. If
// this callback is set, the response returned by the fake server depends on the
// value returned by the callback. If this callback is not set, the fake server
// responds with an empty response.
func (s *FakeRouteLookupServer) SetResponseCallback(f func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse) {
s.mu.Lock()
s.respCb = f
s.mu.Unlock()
}
// SetRequestCallback sets a callback to be invoked on every RLS request. The
// callback is given the incoming request, and tests can use this to verify that
// the request matches its expectations.
func (s *FakeRouteLookupServer) SetRequestCallback(f func(*rlspb.RouteLookupRequest)) {
s.mu.Lock()
s.reqCb = f
s.mu.Unlock()
}