github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/rpc/context_testutils.go (about)

     1  // Copyright 2019 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package rpc
    12  
    13  import (
    14  	"github.com/cockroachdb/cockroach/pkg/base"
    15  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    16  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    17  	"github.com/cockroachdb/cockroach/pkg/util/log"
    18  	"github.com/cockroachdb/cockroach/pkg/util/stop"
    19  	"github.com/cockroachdb/cockroach/pkg/util/tracing"
    20  	"github.com/cockroachdb/cockroach/pkg/util/uuid"
    21  	"google.golang.org/grpc"
    22  )
    23  
    24  // ContextTestingKnobs provides hooks to aid in testing the system. The testing
    25  // knob functions are called at various points in the Context life cycle if they
    26  // are non-nil.
    27  type ContextTestingKnobs struct {
    28  
    29  	// UnaryClientInterceptor if non-nil will be called at dial time to provide
    30  	// the base unary interceptor for client connections.
    31  	// This function may return a nil interceptor to avoid injecting behavior
    32  	// for a given target and class.
    33  	UnaryClientInterceptor func(target string, class ConnectionClass) grpc.UnaryClientInterceptor
    34  
    35  	// StreamClient if non-nil will be called at dial time to provide
    36  	// the base stream interceptor for client connections.
    37  	// This function may return a nil interceptor to avoid injecting behavior
    38  	// for a given target and class.
    39  	StreamClientInterceptor func(target string, class ConnectionClass) grpc.StreamClientInterceptor
    40  
    41  	// ArtificialLatencyMap if non-nil contains a map from target address
    42  	// (server.RPCServingAddr() of a remote node) to artificial latency in
    43  	// milliseconds to inject. Setting this will cause the server to pause for
    44  	// the given amount of milliseconds on every network write.
    45  	ArtificialLatencyMap map[string]int
    46  
    47  	// ClusterID initializes the Context's ClusterID container to this value if
    48  	// non-nil at construction time.
    49  	ClusterID *uuid.UUID
    50  }
    51  
    52  // NewInsecureTestingContext creates an insecure rpc Context suitable for tests.
    53  func NewInsecureTestingContext(clock *hlc.Clock, stopper *stop.Stopper) *Context {
    54  	clusterID := uuid.MakeV4()
    55  	return NewInsecureTestingContextWithClusterID(clock, stopper, clusterID)
    56  }
    57  
    58  // NewInsecureTestingContextWithClusterID creates an insecure rpc Context
    59  // suitable for tests. The context is given the provided cluster ID.
    60  func NewInsecureTestingContextWithClusterID(
    61  	clock *hlc.Clock, stopper *stop.Stopper, clusterID uuid.UUID,
    62  ) *Context {
    63  	return NewInsecureTestingContextWithKnobs(clock, stopper, ContextTestingKnobs{
    64  		ClusterID: &clusterID,
    65  	})
    66  }
    67  
    68  // NewInsecureTestingContextWithKnobs creates an insecure rpc Context
    69  // suitable for tests configured with the provided knobs.
    70  func NewInsecureTestingContextWithKnobs(
    71  	clock *hlc.Clock, stopper *stop.Stopper, knobs ContextTestingKnobs,
    72  ) *Context {
    73  	return NewContextWithTestingKnobs(
    74  		log.AmbientContext{Tracer: tracing.NewTracer()},
    75  		&base.Config{Insecure: true},
    76  		clock,
    77  		stopper,
    78  		cluster.MakeTestingClusterSettings(),
    79  		knobs,
    80  	)
    81  }