github.com/hxx258456/ccgo@v0.0.5-0.20230213014102-48b35f46f66f/grpc/balancer/rls/internal/test/e2e/rls_fakeserver.go (about)

     1  /*
     2   *
     3   * Copyright 2020 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package e2e
    20  
    21  import (
    22  	"context"
    23  	"net"
    24  	"sync"
    25  	"testing"
    26  
    27  	grpc "github.com/hxx258456/ccgo/grpc"
    28  	"github.com/hxx258456/ccgo/grpc/codes"
    29  	rlsgrpc "github.com/hxx258456/ccgo/grpc/internal/proto/grpc_lookup_v1"
    30  	rlspb "github.com/hxx258456/ccgo/grpc/internal/proto/grpc_lookup_v1"
    31  	"github.com/hxx258456/ccgo/grpc/internal/testutils"
    32  	"github.com/hxx258456/ccgo/grpc/status"
    33  )
    34  
    35  // RouteLookupResponse wraps an RLS response and the associated error to be sent
    36  // to a client when the RouteLookup RPC is invoked.
    37  type RouteLookupResponse struct {
    38  	Resp *rlspb.RouteLookupResponse
    39  	Err  error
    40  }
    41  
    42  // FakeRouteLookupServer is a fake implementation of the RouteLookupService.
    43  //
    44  // It is safe for concurrent use.
    45  type FakeRouteLookupServer struct {
    46  	rlsgrpc.UnimplementedRouteLookupServiceServer
    47  	Address string
    48  
    49  	mu     sync.Mutex
    50  	respCb func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse
    51  	reqCb  func(*rlspb.RouteLookupRequest)
    52  }
    53  
    54  // StartFakeRouteLookupServer starts a fake RLS server listening for requests on
    55  // lis. If lis is nil, it creates a new listener on a random local port. The
    56  // returned cancel function should be invoked by the caller upon completion of
    57  // the test.
    58  func StartFakeRouteLookupServer(t *testing.T, lis net.Listener, opts ...grpc.ServerOption) (*FakeRouteLookupServer, func()) {
    59  	t.Helper()
    60  
    61  	if lis == nil {
    62  		var err error
    63  		lis, err = testutils.LocalTCPListener()
    64  		if err != nil {
    65  			t.Fatalf("net.Listen() failed: %v", err)
    66  		}
    67  	}
    68  
    69  	s := &FakeRouteLookupServer{Address: lis.Addr().String()}
    70  	server := grpc.NewServer(opts...)
    71  	rlsgrpc.RegisterRouteLookupServiceServer(server, s)
    72  	go server.Serve(lis)
    73  	return s, func() { server.Stop() }
    74  }
    75  
    76  // RouteLookup implements the RouteLookupService.
    77  func (s *FakeRouteLookupServer) RouteLookup(ctx context.Context, req *rlspb.RouteLookupRequest) (*rlspb.RouteLookupResponse, error) {
    78  	s.mu.Lock()
    79  	defer s.mu.Unlock()
    80  	if s.reqCb != nil {
    81  		s.reqCb(req)
    82  	}
    83  	if err := ctx.Err(); err != nil {
    84  		return nil, status.Error(codes.DeadlineExceeded, err.Error())
    85  	}
    86  	if s.respCb == nil {
    87  		return &rlspb.RouteLookupResponse{}, nil
    88  	}
    89  	resp := s.respCb(ctx, req)
    90  	return resp.Resp, resp.Err
    91  }
    92  
    93  // SetResponseCallback sets a callback to be invoked on every RLS request. If
    94  // this callback is set, the response returned by the fake server depends on the
    95  // value returned by the callback. If this callback is not set, the fake server
    96  // responds with an empty response.
    97  func (s *FakeRouteLookupServer) SetResponseCallback(f func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse) {
    98  	s.mu.Lock()
    99  	s.respCb = f
   100  	s.mu.Unlock()
   101  }
   102  
   103  // SetRequestCallback sets a callback to be invoked on every RLS request. The
   104  // callback is given the incoming request, and tests can use this to verify that
   105  // the request matches its expectations.
   106  func (s *FakeRouteLookupServer) SetRequestCallback(f func(*rlspb.RouteLookupRequest)) {
   107  	s.mu.Lock()
   108  	s.reqCb = f
   109  	s.mu.Unlock()
   110  }