1 files changed, 46 insertions, 31 deletions
diff --git a/grpc/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py b/grpc/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py
index a0013461..838aed07 100644
--- a/grpc/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py
+++ b/grpc/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/compute.py
@@ -28,7 +28,7 @@ logger = logging.getLogger(__name__)
 
 class ComputeV1(gcp.api.GcpProjectApiResource):
     # TODO(sergiitk): move someplace better
-    _WAIT_FOR_BACKEND_SEC = 60 * 5
+    _WAIT_FOR_BACKEND_SEC = 60 * 10
     _WAIT_FOR_OPERATION_SEC = 60 * 5
 
     @dataclasses.dataclass(frozen=True)
@@ -45,22 +45,37 @@ class ComputeV1(gcp.api.GcpProjectApiResource):
 
     class HealthCheckProtocol(enum.Enum):
         TCP = enum.auto()
+        GRPC = enum.auto()
 
     class BackendServiceProtocol(enum.Enum):
         HTTP2 = enum.auto()
         GRPC = enum.auto()
 
-    def create_health_check_tcp(self, name,
-                                use_serving_port=False) -> GcpResource:
+    def create_health_check(self,
+                            name: str,
+                            protocol: HealthCheckProtocol,
+                            *,
+                            port: Optional[int] = None) -> GcpResource:
+        if protocol is self.HealthCheckProtocol.TCP:
+            health_check_field = 'tcpHealthCheck'
+        elif protocol is self.HealthCheckProtocol.GRPC:
+            health_check_field = 'grpcHealthCheck'
+        else:
+            raise TypeError(f'Unexpected Health Check protocol: {protocol}')
+
         health_check_settings = {}
-        if use_serving_port:
+        if port is None:
             health_check_settings['portSpecification'] = 'USE_SERVING_PORT'
+        else:
+            health_check_settings['portSpecification'] = 'USE_FIXED_PORT'
+            health_check_settings['port'] = port
 
-        return self._insert_resource(self.api.healthChecks(), {
-            'name': name,
-            'type': 'TCP',
-            'tcpHealthCheck': health_check_settings,
-        })
+        return self._insert_resource(
+            self.api.healthChecks(), {
+                'name': name,
+                'type': protocol.name,
+                health_check_field: health_check_settings,
+            })
 
     def delete_health_check(self, name):
         self._delete_resource(self.api.healthChecks(), 'healthCheck', name)
@@ -113,12 +128,12 @@ class ComputeV1(gcp.api.GcpProjectApiResource):
                               name)
 
     def create_url_map(
-            self,
-            name: str,
-            matcher_name: str,
-            src_hosts,
-            dst_default_backend_service: GcpResource,
-            dst_host_rule_match_backend_service: Optional[GcpResource] = None,
+        self,
+        name: str,
+        matcher_name: str,
+        src_hosts,
+        dst_default_backend_service: GcpResource,
+        dst_host_rule_match_backend_service: Optional[GcpResource] = None,
     ) -> GcpResource:
         if dst_host_rule_match_backend_service is None:
             dst_host_rule_match_backend_service = dst_default_backend_service
@@ -142,9 +157,9 @@ class ComputeV1(gcp.api.GcpProjectApiResource):
         self._delete_resource(self.api.urlMaps(), 'urlMap', name)
 
     def create_target_grpc_proxy(
-            self,
-            name: str,
-            url_map: GcpResource,
+        self,
+        name: str,
+        url_map: GcpResource,
     ) -> GcpResource:
         return self._insert_resource(self.api.targetGrpcProxies(), {
             'name': name,
@@ -157,9 +172,9 @@ class ComputeV1(gcp.api.GcpProjectApiResource):
                               name)
 
     def create_target_http_proxy(
-            self,
-            name: str,
-            url_map: GcpResource,
+        self,
+        name: str,
+        url_map: GcpResource,
     ) -> GcpResource:
         return self._insert_resource(self.api.targetHttpProxies(), {
             'name': name,
@@ -171,11 +186,11 @@ class ComputeV1(gcp.api.GcpProjectApiResource):
                               name)
 
     def create_forwarding_rule(
-            self,
-            name: str,
-            src_port: int,
-            target_proxy: GcpResource,
-            network_url: str,
+        self,
+        name: str,
+        src_port: int,
+        target_proxy: GcpResource,
+        network_url: str,
     ) -> GcpResource:
         return self._insert_resource(
             self.api.globalForwardingRules(),
@@ -229,11 +244,11 @@ class ComputeV1(gcp.api.GcpProjectApiResource):
         return neg
 
     def wait_for_backends_healthy_status(
-            self,
-            backend_service,
-            backends,
-            timeout_sec=_WAIT_FOR_BACKEND_SEC,
-            wait_sec=4,
+        self,
+        backend_service,
+        backends,
+        timeout_sec=_WAIT_FOR_BACKEND_SEC,
+        wait_sec=4,
     ):
         pending = set(backends)