diff options
Diffstat (limited to 'tools/run_tests/stress_test/run_on_gke.py')
| -rwxr-xr-x | tools/run_tests/stress_test/run_on_gke.py | 13 |
1 files changed, 9 insertions, 4 deletions
diff --git a/tools/run_tests/stress_test/run_on_gke.py b/tools/run_tests/stress_test/run_on_gke.py index d4f1c4ad3d..583e58316f 100755 --- a/tools/run_tests/stress_test/run_on_gke.py +++ b/tools/run_tests/stress_test/run_on_gke.py @@ -69,7 +69,7 @@ class ClientTemplate: def __init__(self, name, stress_client_cmd, metrics_client_cmd, metrics_port, wrapper_script_path, poll_interval_secs, client_args_dict, - metrics_args_dict, will_run_forever): + metrics_args_dict, will_run_forever, env_dict): self.name = name self.stress_client_cmd = stress_client_cmd self.metrics_client_cmd = metrics_client_cmd @@ -79,19 +79,21 @@ class ClientTemplate: self.client_args_dict = client_args_dict self.metrics_args_dict = metrics_args_dict self.will_run_forever = will_run_forever + self.env_dict = env_dict class ServerTemplate: """ Contains all the common settings used by a stress server """ def __init__(self, name, server_cmd, wrapper_script_path, server_port, - server_args_dict, will_run_forever): + server_args_dict, will_run_forever, env_dict): self.name = name self.server_cmd = server_cmd self.wrapper_script_path = wrapper_script_path self.server_port = server_port self.server_args_dict = server_args_dict self.will_run_forever = will_run_forever + self.env_dict = env_dict class DockerImage: @@ -240,6 +242,7 @@ class Gke: # server_pod_spec.template.wrapper_script_path) are are injected into the # container via environment variables server_env = self.gke_env.copy() + server_env.update(server_pod_spec.template.env_dict) server_env.update({ 'STRESS_TEST_IMAGE_TYPE': 'SERVER', 'STRESS_TEST_CMD': server_pod_spec.template.server_cmd, @@ -283,6 +286,7 @@ class Gke: # client_pod_spec.template.wrapper_script_path) are are injected into the # container via environment variables client_env = self.gke_env.copy() + client_env.update(client_pod_spec.template.env_dict) client_env.update({ 'STRESS_TEST_IMAGE_TYPE': 'CLIENT', 'STRESS_TEST_CMD': client_pod_spec.template.stress_client_cmd, @@ -425,7 +429,8 @@ class Config: template_name, stress_client_cmd, metrics_client_cmd, temp_dict['metricsPort'], temp_dict['wrapperScriptPath'], temp_dict['pollIntervalSecs'], temp_dict['clientArgs'].copy(), - temp_dict['metricsArgs'].copy(), temp_dict.get('willRunForever', 1)) + temp_dict['metricsArgs'].copy(), temp_dict.get('willRunForever', 1), + temp_dict.get('env', {}).copy()) return client_templates_dict @@ -461,7 +466,7 @@ class Config: server_templates_dict[template_name] = ServerTemplate( template_name, stress_server_cmd, temp_dict['wrapperScriptPath'], temp_dict['serverPort'], temp_dict['serverArgs'].copy(), - temp_dict.get('willRunForever', 1)) + temp_dict.get('willRunForever', 1), temp_dict.get('env', {}).copy()) return server_templates_dict |