Allow interop tests to configure locality_lb_policies (#29688)

* Allow tests to configure locality_lb_policies
* Custom LB test

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':

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
 

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

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)