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

     1  /*
     2   *
     3   * Copyright 2023 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 testutils
    20  
    21  import (
    22  	"context"
    23  
    24  	"google.golang.org/grpc/stats"
    25  )
    26  
    27  // StubStatsHandler is a stats handler that is easy to customize within
    28  // individual test cases. It is a stubbable implementation of
    29  // google.golang.org/grpc/stats.Handler for testing purposes.
    30  type StubStatsHandler struct {
    31  	TagRPCF     func(ctx context.Context, info *stats.RPCTagInfo) context.Context
    32  	HandleRPCF  func(ctx context.Context, info stats.RPCStats)
    33  	TagConnF    func(ctx context.Context, info *stats.ConnTagInfo) context.Context
    34  	HandleConnF func(ctx context.Context, info stats.ConnStats)
    35  }
    36  
    37  // TagRPC calls the StubStatsHandler's TagRPCF, if set.
    38  func (ssh *StubStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
    39  	if ssh.TagRPCF != nil {
    40  		return ssh.TagRPCF(ctx, info)
    41  	}
    42  	return ctx
    43  }
    44  
    45  // HandleRPC calls the StubStatsHandler's HandleRPCF, if set.
    46  func (ssh *StubStatsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
    47  	if ssh.HandleRPCF != nil {
    48  		ssh.HandleRPCF(ctx, rs)
    49  	}
    50  }
    51  
    52  // TagConn calls the StubStatsHandler's TagConnF, if set.
    53  func (ssh *StubStatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
    54  	if ssh.TagConnF != nil {
    55  		return ssh.TagConnF(ctx, info)
    56  	}
    57  	return ctx
    58  }
    59  
    60  // HandleConn calls the StubStatsHandler's HandleConnF, if set.
    61  func (ssh *StubStatsHandler) HandleConn(ctx context.Context, cs stats.ConnStats) {
    62  	if ssh.HandleConnF != nil {
    63  		ssh.HandleConnF(ctx, cs)
    64  	}
    65  }