diff --git a/tools/run_tests/xds_k8s_test_driver/bin/cleanup.sh b/tools/run_tests/xds_k8s_test_driver/bin/cleanup.sh index ba53491c3fc..3b4bf270aed 100755 --- a/tools/run_tests/xds_k8s_test_driver/bin/cleanup.sh +++ b/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 diff --git a/tools/run_tests/xds_k8s_test_driver/config/grpc-testing.cfg b/tools/run_tests/xds_k8s_test_driver/config/grpc-testing.cfg index 93c8a6d9e54..2e89c7ebca7 100644 --- a/tools/run_tests/xds_k8s_test_driver/config/grpc-testing.cfg +++ b/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 diff --git a/tools/run_tests/xds_k8s_test_driver/framework/helpers/highlighter.py b/tools/run_tests/xds_k8s_test_driver/framework/helpers/highlighter.py new file mode 100644 index 00000000000..0c9f0112614 --- /dev/null +++ b/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) diff --git a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/api.py b/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/api.py index 2c384a3f320..c5184a84d59 100644 --- a/tools/run_tests/xds_k8s_test_driver/framework/infrastructure/gcp/api.py +++ b/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' diff --git a/tools/run_tests/xds_k8s_test_driver/framework/test_app/base_runner.py b/tools/run_tests/xds_k8s_test_driver/framework/test_app/base_runner.py index b1fee03e6f8..5c9e32c88f7 100644 --- a/tools/run_tests/xds_k8s_test_driver/framework/test_app/base_runner.py +++ b/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) diff --git a/tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py b/tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py index 6d13a79d4ba..4232f37c8c2 100644 --- a/tools/run_tests/xds_k8s_test_driver/framework/xds_flags.py +++ b/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. diff --git a/tools/run_tests/xds_k8s_test_driver/requirements.txt b/tools/run_tests/xds_k8s_test_driver/requirements.txt index 8f573ca22a7..7783cc32365 100644 --- a/tools/run_tests/xds_k8s_test_driver/requirements.txt +++ b/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