From 8c9f3c2b30ae548de0c498cd073a3cd925716fa0 Mon Sep 17 00:00:00 2001
From: Paulo Castello da Costa <6579971+paulosjca@users.noreply.github.com>
Date: Mon, 21 Jun 2021 12:06:30 -0700
Subject: [PATCH] Update tools in preparation for load test timeouts. (#26505)

- Support multiline strings in template generation (as already supported in config generation, supporting roundtrip).
- Ignore substitution of variables that are set by the controller at runtime (`DRIVER_PORT`, `KILL_AFTER`, `POD_TIMEOUT`).
---
 tools/run_tests/performance/README.md            |  5 +++++
 tools/run_tests/performance/loadtest_config.py   | 12 +++++++++++-
 tools/run_tests/performance/loadtest_template.py |  9 +++++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/tools/run_tests/performance/README.md b/tools/run_tests/performance/README.md
index b87b9f53724..8f2b0ffed4a 100644
--- a/tools/run_tests/performance/README.md
+++ b/tools/run_tests/performance/README.md
@@ -257,6 +257,11 @@ The script `loadtest_config.py` takes the following options:
 - `-t`, `--template`<br> Template file. A template is a configuration file that
   may contain multiple client and server configuration, and may also include
   substitution keys.
+- `-s`, `--substitution` Substitution keys, in the format `key=value`. These
+  keys are substituted while processing the template. Environment variables that
+  are set by the load test controller at runtime are ignored by default
+  (`DRIVER_PORT`, `KILL_AFTER`, `POD_TIMEOUT`). The user can override this
+  behavior by specifying these variables as keys.
 - `-p`, `--prefix`<br> Test names consist of a prefix_joined with a uuid with a
   dash. Test names are stored in `metadata.name`. The prefix is also added as
   the `prefix` label in `metadata.labels`. The prefix defaults to the user name
diff --git a/tools/run_tests/performance/loadtest_config.py b/tools/run_tests/performance/loadtest_config.py
index a25ed07628f..0e48daaa832 100755
--- a/tools/run_tests/performance/loadtest_config.py
+++ b/tools/run_tests/performance/loadtest_config.py
@@ -383,7 +383,17 @@ def main() -> None:
     if args.runs_per_test < 1:
         argp.error('runs_per_test must be greater than zero.')
 
-    substitutions = parse_key_value_args(args.substitutions)
+    # Config generation ignores environment variables that are passed by the
+    # controller at runtime.
+    substitutions = {
+        'DRIVER_PORT': '${DRIVER_PORT}',
+        'KILL_AFTER': '${KILL_AFTER}',
+        'POD_TIMEOUT': '${POD_TIMEOUT}',
+    }
+
+    # The user can override the ignored variables above by passing them in as
+    # substitution keys.
+    substitutions.update(parse_key_value_args(args.substitutions))
 
     uniquifier_elements = args.uniquifier_elements
     if args.d:
diff --git a/tools/run_tests/performance/loadtest_template.py b/tools/run_tests/performance/loadtest_template.py
index e2ed4a451b2..58145526d35 100755
--- a/tools/run_tests/performance/loadtest_template.py
+++ b/tools/run_tests/performance/loadtest_template.py
@@ -177,6 +177,15 @@ def template_dumper(header_comment: str) -> Type[yaml.SafeDumper]:
                 self.write_indent()
                 self.write_indicator(header_comment, need_whitespace=False)
 
+    def str_presenter(dumper, data):
+        if '\n' in data:
+            return dumper.represent_scalar('tag:yaml.org,2002:str',
+                                           data,
+                                           style='|')
+        return dumper.represent_scalar('tag:yaml.org,2002:str', data)
+
+    TemplateDumper.add_representer(str, str_presenter)
+
     return TemplateDumper