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"""
15This contains helpers for gRPC services defined in
16https://github.com/grpc/grpc/blob/master/src/proto/grpc/testing/test.proto
17"""
18import logging
19from typing import Optional
20
21import grpc
22
23import framework.rpc
24from src.proto.grpc.testing import test_pb2_grpc
25from src.proto.grpc.testing import messages_pb2
26
27# Type aliases
28_LoadBalancerStatsRequest = messages_pb2.LoadBalancerStatsRequest
29LoadBalancerStatsResponse = messages_pb2.LoadBalancerStatsResponse
30
31
32class 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                                             wait_for_ready_sec=timeout_sec,
53                                             log_level=logging.INFO)
54