google.golang.org/grpc@v1.62.1/internal/testutils/rls/fake_rls_server.go (about)

     1  /*
     2   *
     3   * Copyright 2022 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 rls contains utilities for RouteLookupService e2e tests.
    20  package rls
    21  
    22  import (
    23  	"context"
    24  	"net"
    25  	"sync"
    26  	"testing"
    27  
    28  	"google.golang.org/grpc"
    29  	"google.golang.org/grpc/codes"
    30  	rlsgrpc "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
    31  	rlspb "google.golang.org/grpc/internal/proto/grpc_lookup_v1"
    32  	"google.golang.org/grpc/internal/testutils"
    33  	"google.golang.org/grpc/status"
    34  )
    35  
    36  // RouteLookupResponse wraps an RLS response and the associated error to be sent
    37  // to a client when the RouteLookup RPC is invoked.
    38  type RouteLookupResponse struct {
    39  	Resp *rlspb.RouteLookupResponse
    40  	Err  error
    41  }
    42  
    43  // SetupFakeRLSServer starts and returns a fake RouteLookupService server
    44  // listening on the given listener or on a random local port. Also returns a
    45  // channel for tests to get notified whenever the RouteLookup RPC is invoked on
    46  // the fake server.
    47  //
    48  // This function sets up the fake server to respond with an empty response for
    49  // the RouteLookup RPCs. Tests can override this by calling the
    50  // SetResponseCallback() method on the returned fake server.
    51  func SetupFakeRLSServer(t *testing.T, lis net.Listener, opts ...grpc.ServerOption) (*FakeRouteLookupServer, chan struct{}) {
    52  	s, cancel := StartFakeRouteLookupServer(t, lis, opts...)
    53  	t.Logf("Started fake RLS server at %q", s.Address)
    54  
    55  	ch := make(chan struct{}, 1)
    56  	s.SetRequestCallback(func(request *rlspb.RouteLookupRequest) {
    57  		select {
    58  		case ch <- struct{}{}:
    59  		default:
    60  		}
    61  	})
    62  	t.Cleanup(cancel)
    63  	return s, ch
    64  }
    65  
    66  // FakeRouteLookupServer is a fake implementation of the RouteLookupService.
    67  //
    68  // It is safe for concurrent use.
    69  type FakeRouteLookupServer struct {
    70  	rlsgrpc.UnimplementedRouteLookupServiceServer
    71  	Address string
    72  
    73  	mu     sync.Mutex
    74  	respCb func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse
    75  	reqCb  func(*rlspb.RouteLookupRequest)
    76  }
    77  
    78  // StartFakeRouteLookupServer starts a fake RLS server listening for requests on
    79  // lis. If lis is nil, it creates a new listener on a random local port. The
    80  // returned cancel function should be invoked by the caller upon completion of
    81  // the test.
    82  func StartFakeRouteLookupServer(t *testing.T, lis net.Listener, opts ...grpc.ServerOption) (*FakeRouteLookupServer, func()) {
    83  	t.Helper()
    84  
    85  	if lis == nil {
    86  		var err error
    87  		lis, err = testutils.LocalTCPListener()
    88  		if err != nil {
    89  			t.Fatalf("net.Listen() failed: %v", err)
    90  		}
    91  	}
    92  
    93  	s := &FakeRouteLookupServer{Address: lis.Addr().String()}
    94  	server := grpc.NewServer(opts...)
    95  	rlsgrpc.RegisterRouteLookupServiceServer(server, s)
    96  	go server.Serve(lis)
    97  	return s, func() { server.Stop() }
    98  }
    99  
   100  // RouteLookup implements the RouteLookupService.
   101  func (s *FakeRouteLookupServer) RouteLookup(ctx context.Context, req *rlspb.RouteLookupRequest) (*rlspb.RouteLookupResponse, error) {
   102  	s.mu.Lock()
   103  	defer s.mu.Unlock()
   104  	if s.reqCb != nil {
   105  		s.reqCb(req)
   106  	}
   107  	if err := ctx.Err(); err != nil {
   108  		return nil, status.Error(codes.DeadlineExceeded, err.Error())
   109  	}
   110  	if s.respCb == nil {
   111  		return &rlspb.RouteLookupResponse{}, nil
   112  	}
   113  	resp := s.respCb(ctx, req)
   114  	return resp.Resp, resp.Err
   115  }
   116  
   117  // SetResponseCallback sets a callback to be invoked on every RLS request. If
   118  // this callback is set, the response returned by the fake server depends on the
   119  // value returned by the callback. If this callback is not set, the fake server
   120  // responds with an empty response.
   121  func (s *FakeRouteLookupServer) SetResponseCallback(f func(context.Context, *rlspb.RouteLookupRequest) *RouteLookupResponse) {
   122  	s.mu.Lock()
   123  	s.respCb = f
   124  	s.mu.Unlock()
   125  }
   126  
   127  // SetRequestCallback sets a callback to be invoked on every RLS request. The
   128  // callback is given the incoming request, and tests can use this to verify that
   129  // the request matches its expectations.
   130  func (s *FakeRouteLookupServer) SetRequestCallback(f func(*rlspb.RouteLookupRequest)) {
   131  	s.mu.Lock()
   132  	s.reqCb = f
   133  	s.mu.Unlock()
   134  }