xds-k8s: Colorize all the things (#26845)

tools/run_tests/xds_k8s_test_driver/bin/cleanup.sh

@@ -50,10 +50,10 @@ cd "${XDS_K8S_DRIVER_DIR}"
 if [[ "$1" == "--nosecure" ]]; then
   shift
   ./run.sh bin/run_td_setup.py --cmd=cleanup "$@" && \
-  ./run.sh bin/run_test_client.py --cmd=cleanup "$@" && \
+  ./run.sh bin/run_test_client.py --cmd=cleanup --cleanup_namespace "$@" && \
   ./run.sh bin/run_test_server.py --cmd=cleanup --cleanup_namespace "$@"
 else
   ./run.sh bin/run_td_setup.py --cmd=cleanup --security=mtls "$@" && \
-  ./run.sh bin/run_test_client.py --cmd=cleanup --secure "$@" && \
+  ./run.sh bin/run_test_client.py --cmd=cleanup --cleanup_namespace --secure "$@" && \
   ./run.sh bin/run_test_server.py --cmd=cleanup --cleanup_namespace --secure "$@"
 fi

tools/run_tests/xds_k8s_test_driver/config/grpc-testing.cfg

@@ -5,3 +5,5 @@
 --private_api_key_secret_name=projects/830293263384/secrets/xds-interop-tests-private-api-access-key
 # Randomize xds port.
 --server_xds_port=0
+# ResultStore UI doesn't support 256 colors.
+--color_style=ansi16

tools/run_tests/xds_k8s_test_driver/framework/helpers/highlighter.py

@@ -0,0 +1,99 @@
+# Copyright 2021 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.
+"""The module contains helpers to enable color output in terminals.
+
+Use this to log resources dumped as a structured document (f.e. YAML),
+and enable colorful syntax highlighting.
+
+TODO(sergiitk): This can be used to output protobuf responses formatted as JSON.
+"""
+import logging
+from typing import Optional
+
+from absl import flags
+import pygments
+import pygments.formatter
+import pygments.formatters.other
+import pygments.formatters.terminal
+import pygments.formatters.terminal256
+import pygments.lexer
+import pygments.lexers.data
+import pygments.styles
+
+# The style for terminals supporting 8/16 colors.
+STYLE_ANSI_16 = 'ansi16'
+# Join with pygments styles for terminals supporting 88/256 colors.
+ALL_COLOR_STYLES = [STYLE_ANSI_16] + list(pygments.styles.get_all_styles())
+
+# Flags.
+COLOR = flags.DEFINE_bool("color", default=True, help='Colorize the output')
+COLOR_STYLE = flags.DEFINE_enum(
+    "color_style",
+    default='material',
+    enum_values=ALL_COLOR_STYLES,
+    help=('Color styles for terminals supporting 256 colors. '
+          f'Use {STYLE_ANSI_16} style for terminals supporting 8/16 colors'))
+
+logger = logging.getLogger(__name__)
+
+# Type aliases.
+Lexer = pygments.lexer.Lexer
+YamlLexer = pygments.lexers.data.YamlLexer
+Formatter = pygments.formatter.Formatter
+NullFormatter = pygments.formatters.other.NullFormatter
+TerminalFormatter = pygments.formatters.terminal.TerminalFormatter
+Terminal256Formatter = pygments.formatters.terminal256.Terminal256Formatter
+
+
+class Highlighter:
+    formatter: Formatter
+    lexer: Lexer
+    color: bool
+    color_style: Optional[str] = None
+
+    def __init__(self,
+                 *,
+                 lexer: Lexer,
+                 color: Optional[bool] = None,
+                 color_style: Optional[str] = None):
+        self.lexer = lexer
+        self.color = color if color is not None else COLOR.value
+
+        if self.color:
+            color_style = color_style if color_style else COLOR_STYLE.value
+            if color_style not in ALL_COLOR_STYLES:
+                raise ValueError(f'Unrecognized color style {color_style}, '
+                                 f'valid styles: {ALL_COLOR_STYLES}')
+            if color_style == STYLE_ANSI_16:
+                # 8/16 colors support only.
+                self.formatter = TerminalFormatter()
+            else:
+                # 88/256 colors.
+                self.formatter = Terminal256Formatter(style=color_style)
+        else:
+            self.formatter = NullFormatter()
+
+    def highlight(self, code: str) -> str:
+        return pygments.highlight(code, self.lexer, self.formatter)
+
+
+class HighlighterYaml(Highlighter):
+
+    def __init__(self,
+                 *,
+                 color: Optional[bool] = None,
+                 color_style: Optional[str] = None):
+        super().__init__(lexer=YamlLexer(encoding='utf-8'),
+                         color=color,
+                         color_style=color_style)

tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/api.py

@@ -28,6 +28,8 @@ import googleapiclient.http
 import tenacity
 import yaml
 
+import framework.helpers.highlighter
+
 logger = logging.getLogger(__name__)
 PRIVATE_API_KEY_SECRET_NAME = flags.DEFINE_string(
     "private_api_key_secret_name",
@@ -51,6 +53,7 @@ GCP_UI_URL = flags.DEFINE_string("gcp_ui_url",
 # Type aliases
 _HttpError = googleapiclient.errors.HttpError
 _HttpLib2Error = googleapiclient.http.httplib2.HttpLib2Error
+_HighlighterYaml = framework.helpers.highlighter.HighlighterYaml
 Operation = operations_pb2.Operation
 HttpRequest = googleapiclient.http.HttpRequest
 
@@ -273,6 +276,7 @@ class GcpProjectApiResource:
     def __init__(self, api: discovery.Resource, project: str):
         self.api: discovery.Resource = api
         self.project: str = project
+        self._highlighter = _HighlighterYaml()
 
     # TODO(sergiitk): in upcoming GCP refactoring, differentiate between
     #   _execute for LRO (Long Running Operations), and immediate operations.
@@ -299,6 +303,11 @@ class GcpProjectApiResource:
         except _HttpLib2Error as error:
             raise TransportError(error)
 
+    def resource_pretty_format(self, body: dict) -> str:
+        """Return a string with pretty-printed resource body."""
+        yaml_out: str = yaml.dump(body, explicit_start=True, explicit_end=True)
+        return self._highlighter.highlight(yaml_out)
+
     @staticmethod
     def wait_for_operation(operation_request,
                            test_success_fn,
@@ -313,11 +322,6 @@ class GcpProjectApiResource:
             reraise=True)
         return retryer(operation_request.execute)
 
-    @staticmethod
-    def resource_pretty_format(body: dict) -> str:
-        """Return a string with pretty-printed resource body."""
-        return yaml.dump(body, explicit_start=True, explicit_end=True)
-
 
 class GcpStandardCloudApiResource(GcpProjectApiResource, metaclass=abc.ABCMeta):
     GLOBAL_LOCATION = 'global'

tools/run_tests/xds_k8s_test_driver/framework/test_app/base_runner.py

@@ -22,12 +22,14 @@ import mako.template
 import yaml
 
 import framework.helpers.datetime
+import framework.helpers.highlighter
 from framework.infrastructure import gcp
 from framework.infrastructure import k8s
 
 logger = logging.getLogger(__name__)
 
 # Type aliases
+_HighlighterYaml = framework.helpers.highlighter.HighlighterYaml
 _helper_datetime = framework.helpers.datetime
 timedelta = datetime.timedelta
 
@@ -57,6 +59,8 @@ class KubernetesBaseRunner:
                  k8s_namespace,
                  namespace_template=None,
                  reuse_namespace=False):
+        self._highlighter = _HighlighterYaml()
+
         # Kubernetes namespaced resources manager
         self.k8s_namespace: k8s.KubernetesNamespace = k8s_namespace
         self.reuse_namespace = reuse_namespace
@@ -107,7 +111,7 @@ class KubernetesBaseRunner:
 
         yaml_doc = self._render_template(template_file, **kwargs)
         logger.info("Rendered template %s/%s:\n%s", self.TEMPLATE_DIR_NAME,
-                    template_name, yaml_doc)
+                    template_name, self._highlighter.highlight(yaml_doc))
 
         manifests = self._manifests_from_str(yaml_doc)
         manifest = next(manifests)

tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from absl import flags
-import googleapiclient.discovery
+
+from framework.helpers import highlighter
 
 # GCP
 PROJECT = flags.DEFINE_string("project",
@@ -114,6 +115,8 @@ CLIENT_PORT = flags.DEFINE_integer(
         "XdsStats, XdsUpdateClientConfigure, and ProtoReflection (optional).\n"
         "Doesn't have to be within --firewall_allowed_ports."))
 
+flags.adopt_module_key_flags(highlighter)
+
 flags.mark_flags_as_required([
     "project",
     # TODO(sergiitk): Make required when --namespace is removed.

tools/run_tests/xds_k8s_test_driver/requirements.txt

@@ -12,5 +12,6 @@ kubernetes~=12.0
 # Context: https://github.com/grpc/grpc/pull/24983#discussion_r543017022
 retrying~=1.3
 tenacity~=6.2
+Pygments~=2.9
 protobuf~=3.14
 xds-protos~=0.0.8