mirror of https://github.com/grpc/grpc-go.git
135 lines
4.3 KiB
Go
135 lines
4.3 KiB
Go
/*
|
|
*
|
|
* 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(*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()
|
|
}
|