diff --git a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py index 1f9a3d4a4b5..faffeb477e2 100644 --- a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py +++ b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py @@ -120,7 +120,8 @@ class ComputeV1(gcp.api.GcpProjectApiResource): # pylint: disable=too-many-publ health_check: 'GcpResource', affinity_header: Optional[str] = None, protocol: Optional[BackendServiceProtocol] = None, - subset_size: Optional[int] = None) -> 'GcpResource': + subset_size: Optional[int] = None, + locality_lb_policies: Optional[List[dict]] = None) -> 'GcpResource': if not isinstance(protocol, self.BackendServiceProtocol): raise TypeError(f'Unexpected Backend Service protocol: {protocol}') body = { @@ -142,6 +143,8 @@ class ComputeV1(gcp.api.GcpProjectApiResource): # pylint: disable=too-many-publ 'policy': 'CONSISTENT_HASH_SUBSETTING', 'subsetSize': subset_size } + if locality_lb_policies: + body['localityLbPolicies'] = locality_lb_policies return self._insert_resource(self.api.backendServices(), body) def get_backend_service_traffic_director(self, name: str) -> 'GcpResource': diff --git a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py index ff395301d33..431f8f58aad 100644 --- a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py +++ b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/traffic_director.py @@ -195,7 +195,8 @@ class TrafficDirectorManager: # pylint: disable=too-many-public-methods self, protocol: Optional[BackendServiceProtocol] = _BackendGRPC, subset_size: Optional[int] = None, - affinity_header: Optional[str] = None): + affinity_header: Optional[str] = None, + locality_lb_policies: Optional[List[dict]] = None): if protocol is None: protocol = _BackendGRPC @@ -206,7 +207,8 @@ class TrafficDirectorManager: # pylint: disable=too-many-public-methods health_check=self.health_check, protocol=protocol, subset_size=subset_size, - affinity_header=affinity_header) + affinity_header=affinity_header, + locality_lb_policies=locality_lb_policies) self.backend_service = resource self.backend_service_protocol = protocol diff --git a/tools/run_tests/xds_k8s_test_driver/kubernetes-manifests/client.deployment.yaml b/tools/run_tests/xds_k8s_test_driver/kubernetes-manifests/client.deployment.yaml index bc79d2f1be1..3246f914660 100644 --- a/tools/run_tests/xds_k8s_test_driver/kubernetes-manifests/client.deployment.yaml +++ b/tools/run_tests/xds_k8s_test_driver/kubernetes-manifests/client.deployment.yaml @@ -42,6 +42,8 @@ spec: value: "true" - name: GRPC_XDS_EXPERIMENTAL_ENABLE_RETRY value: "true" + - name: GRPC_EXPERIMENTAL_XDS_CUSTOM_LB_CONFIG + value: "true" volumeMounts: - mountPath: /tmp/grpc-xds/ name: grpc-td-conf diff --git a/tools/run_tests/xds_k8s_test_driver/tests/custom_lb_test.py b/tools/run_tests/xds_k8s_test_driver/tests/custom_lb_test.py new file mode 100644 index 00000000000..475c9661846 --- /dev/null +++ b/tools/run_tests/xds_k8s_test_driver/tests/custom_lb_test.py @@ -0,0 +1,81 @@ +# Copyright 2022 gRPC authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import datetime +import logging + +from absl import flags +from absl.testing import absltest +import grpc + +from framework import xds_k8s_testcase + +logger = logging.getLogger(__name__) +flags.adopt_module_key_flags(xds_k8s_testcase) + +# Type aliases +_XdsTestServer = xds_k8s_testcase.XdsTestServer +_XdsTestClient = xds_k8s_testcase.XdsTestClient + +_EXPECTED_STATUS = grpc.StatusCode.DATA_LOSS + + +class CustomLbTest(xds_k8s_testcase.RegularXdsKubernetesTestCase): + + def test_custom_lb_config(self): + with self.subTest('0_create_health_check'): + self.td.create_health_check() + + # Configures a custom, test LB on the client to instruct the servers + # to always respond with a specific error code. + with self.subTest('1_create_backend_service'): + self.td.create_backend_service(locality_lb_policies=[{ + 'customPolicy': { + 'name': + 'test.RpcBehaviorLoadBalancer', + 'data': + f'{{ "rpcBehavior": "error-code-{_EXPECTED_STATUS.value[0]}" }}' + } + }]) + + with self.subTest('2_create_url_map'): + self.td.create_url_map(self.server_xds_host, self.server_xds_port) + + with self.subTest('3_create_target_proxy'): + self.td.create_target_proxy() + + with self.subTest('4_create_forwarding_rule'): + self.td.create_forwarding_rule(self.server_xds_port) + + with self.subTest('5_start_test_server'): + test_server: _XdsTestServer = self.startTestServers()[0] + + with self.subTest('6_add_server_backends_to_backend_service'): + self.setupServerBackends() + + with self.subTest('7_start_test_client'): + test_client: _XdsTestClient = self.startTestClient(test_server) + + with self.subTest('8_test_client_xds_config_exists'): + self.assertXdsConfigExists(test_client) + + # Verify status codes from the servers have the configured one. + with self.subTest('9_test_server_returned_configured_status_code'): + self.assertRpcStatusCodes(test_client, + status_code=_EXPECTED_STATUS, + duration=datetime.timedelta(seconds=10), + method='UNARY_CALL') + + +if __name__ == '__main__': + absltest.main(failfast=True)