1# Copyright 2021 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.
14import logging
15from typing import List
16
17from absl import flags
18from absl.testing import absltest
19
20from framework import xds_k8s_testcase
21
22logger = logging.getLogger(__name__)
23flags.adopt_module_key_flags(xds_k8s_testcase)
24
25# Type aliases
26_XdsTestServer = xds_k8s_testcase.XdsTestServer
27_XdsTestClient = xds_k8s_testcase.XdsTestClient
28
29
30class RoundRobinTest(xds_k8s_testcase.RegularXdsKubernetesTestCase):
31
32    def test_round_robin(self) -> None:
33        REPLICA_COUNT = 2
34
35        with self.subTest('00_create_health_check'):
36            self.td.create_health_check()
37
38        with self.subTest('01_create_backend_services'):
39            self.td.create_backend_service()
40
41        with self.subTest('02_create_url_map'):
42            self.td.create_url_map(self.server_xds_host, self.server_xds_port)
43
44        with self.subTest('03_create_target_proxy'):
45            self.td.create_target_proxy()
46
47        with self.subTest('04_create_forwarding_rule'):
48            self.td.create_forwarding_rule(self.server_xds_port)
49
50        test_servers: List[_XdsTestServer]
51        with self.subTest('05_start_test_servers'):
52            test_servers = self.startTestServers(replica_count=REPLICA_COUNT)
53
54        with self.subTest('06_add_server_backends_to_backend_services'):
55            self.setupServerBackends()
56
57        test_client: _XdsTestClient
58        with self.subTest('07_start_test_client'):
59            test_client = self.startTestClient(test_servers[0])
60
61        with self.subTest('08_test_client_xds_config_exists'):
62            self.assertXdsConfigExists(test_client)
63
64        with self.subTest('09_test_server_received_rpcs_from_test_client'):
65            self.assertSuccessfulRpcs(test_client)
66
67        with self.subTest('10_round_robin'):
68            num_rpcs = 100
69            expected_rpcs_per_replica = num_rpcs / REPLICA_COUNT
70
71            rpcs_by_peer = self.getClientRpcStats(test_client,
72                                                  num_rpcs).rpcs_by_peer
73            total_requests_received = sum(rpcs_by_peer[x] for x in rpcs_by_peer)
74            self.assertEqual(total_requests_received, num_rpcs,
75                             'Wrong number of RPCS')
76            for server in test_servers:
77                hostname = server.hostname
78                self.assertIn(hostname, rpcs_by_peer,
79                              f'Server {hostname} did not receive RPCs')
80                self.assertLessEqual(
81                    abs(rpcs_by_peer[hostname] - expected_rpcs_per_replica), 1,
82                    f'Wrong number of RPCs for server {hostname}')
83
84
85if __name__ == '__main__':
86    absltest.main(failfast=True)
87