tools/run_tests/xds_k8s_test_driver/framework/rpc/grpc_testing.py

   1 # Copyright 2020 gRPC authors.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 """
  15 This contains helpers for gRPC services defined in
  16 https://github.com/grpc/grpc/blob/master/src/proto/grpc/testing/test.proto
  17 """
  18 import logging
  19 from typing import Optional
  20
  21 import grpc
  22
  23 import framework.rpc
  24 from src.proto.grpc.testing import test_pb2_grpc
  25 from src.proto.grpc.testing import messages_pb2
  26
  27 # Type aliases
  28 _LoadBalancerStatsRequest = messages_pb2.LoadBalancerStatsRequest
  29 LoadBalancerStatsResponse = messages_pb2.LoadBalancerStatsResponse
  30
  31
  32 class LoadBalancerStatsServiceClient(framework.rpc.grpc.GrpcClientHelper):
  33     stub: test_pb2_grpc.LoadBalancerStatsServiceStub
  34     STATS_PARTIAL_RESULTS_TIMEOUT_SEC = 1200
  35
  36     def __init__(self, channel: grpc.Channel):
  37         super().__init__(channel, test_pb2_grpc.LoadBalancerStatsServiceStub)
  38
  39     def get_client_stats(
  40         self,
  41         *,
  42         num_rpcs: int,
  43         timeout_sec: Optional[int] = STATS_PARTIAL_RESULTS_TIMEOUT_SEC,
  44     ) -> LoadBalancerStatsResponse:
  45         if timeout_sec is None:
  46             timeout_sec = self.STATS_PARTIAL_RESULTS_TIMEOUT_SEC
  47
  48         return self.call_unary_with_deadline(rpc='GetClientStats',
  49                                              req=_LoadBalancerStatsRequest(
  50                                                  num_rpcs=num_rpcs,
  51                                                  timeout_sec=timeout_sec),
  52                                              deadline_sec=timeout_sec,
  53                                              log_level=logging.INFO)