diff --git a/endpoints/base b/endpoints/base index 8dec7f65..fb5644b2 100755 --- a/endpoints/base +++ b/endpoints/base @@ -758,7 +758,7 @@ function process_bench_roadblocks() { echo "Initializing data structures:" while read -u 9 line; do - iter_samp=`echo ${line} | awk '{print $1}'` + iter_samp="${line}" iter_id=`echo ${iter_samp} | awk -F- '{print $1}'` samp_id=`echo ${iter_samp} | awk -F- '{print $2}'` iter_array_idx=${total_tests} @@ -772,7 +772,7 @@ function process_bench_roadblocks() { sample_data_attempt_fail[${iter_array_idx}]=0 (( total_tests += 1 )) - done 9< "${engine_bench_cmds_dir}/client/1/start" + done 9< <(xzcat "${engine_bench_cmds_dir}/client/1/start.json.xz" | jq -r '.[] | .test') echo "Total tests: ${total_tests}" @@ -788,7 +788,7 @@ function process_bench_roadblocks() { (( current_test += 1 )) - iter_samp=`echo ${line} | awk '{print $1}'` + iter_samp="${line}" iter_id=`echo ${iter_samp} | awk -F- '{print $1}'` samp_id=`echo ${iter_samp} | awk -F- '{print $2}'` let iter_array_idx=${current_test}-1 @@ -978,7 +978,7 @@ function process_bench_roadblocks() { ${max_sample_failures} \ ${sample_result} done - done 9< "$engine_bench_cmds_dir/client/1/start" + done 9< <(xzcat "$engine_bench_cmds_dir/client/1/start.json.xz" | jq -r '.[] | .test') } function process_roadblocks() { diff --git a/endpoints/endpoints.py b/endpoints/endpoints.py index c3fdaaee..12f69fab 100644 --- a/endpoints/endpoints.py +++ b/endpoints/endpoints.py @@ -1229,26 +1229,28 @@ def process_bench_roadblocks(callbacks = None, roadblock_id = None, endpoint_lab iteration_sample_data = [] logger.info("Initializing data structures") - with open(engine_commands_dir + "/client/1/start") as bench_cmds_fp: - for line in bench_cmds_fp: - split = line.split(" ") - iteration_sample = split[0] - split = iteration_sample.split("-") - iteration_id = int(split[0]) - sample_id = int(split[1]) - - logger.info("iteration_sample=%s iteration_id=%s sample_id=%s" % (iteration_sample, iteration_id, sample_id)) - - obj = { - "iteration-sample": iteration_sample, - "iteration-id": iteration_id, - "sample-id": sample_id, - "failures": 0, - "complete": False, - "attempt-num": 0, - "attempt-fail": 0 - } - iteration_sample_data.append(obj) + bench_cmds, err = load_json_file(engine_commands_dir + "/client/1/start.json.xz", uselzma = True) + if bench_cmds is None: + logger.error("Failed to load bench commands from %s/client/1/start.json.xz: %s" % (engine_commands_dir, err)) + return 1 + for entry in bench_cmds: + iteration_sample = entry["test"] + split = iteration_sample.split("-") + iteration_id = int(split[0]) + sample_id = int(split[1]) + + logger.info("iteration_sample=%s iteration_id=%s sample_id=%s" % (iteration_sample, iteration_id, sample_id)) + + obj = { + "iteration-sample": iteration_sample, + "iteration-id": iteration_id, + "sample-id": sample_id, + "failures": 0, + "complete": False, + "attempt-num": 0, + "attempt-fail": 0 + } + iteration_sample_data.append(obj) logger.info("Total tests: %d" % (len(iteration_sample_data))) diff --git a/engine/engine_lib.py b/engine/engine_lib.py index d9d8195f..d344162a 100644 --- a/engine/engine_lib.py +++ b/engine/engine_lib.py @@ -7,6 +7,7 @@ import logging import os import re +import shlex import shutil import sys import tempfile @@ -533,16 +534,16 @@ def get_data(self): "bench-cmds", self.cs_type, self.cs_id, - "start", + "start.json.xz", ), - "bench-start-cmds", + "bench-start-cmds.json.xz", ) else: self.scp_from_controller( os.path.join( - self.engine_config_dir, "bench-cmds", "client", "1", "start" + self.engine_config_dir, "bench-cmds", "client", "1", "start.json.xz" ), - "bench-start-cmds", + "bench-start-cmds.json.xz", ) if self.cs_type == "client": @@ -552,9 +553,9 @@ def get_data(self): "bench-cmds", self.cs_type, self.cs_id, - "infra", + "infra.json.xz", ), - "bench-infra-cmds", + "bench-infra-cmds.json.xz", ) if self.cs_id == "1": self.scp_from_controller( @@ -563,9 +564,9 @@ def get_data(self): "bench-cmds", self.cs_type, self.cs_id, - "runtime", + "runtime.json.xz", ), - "bench-runtime-cmds", + "bench-runtime-cmds.json.xz", ) elif self.cs_type == "server": self.scp_from_controller( @@ -574,9 +575,9 @@ def get_data(self): "bench-cmds", self.cs_type, self.cs_id, - "stop", + "stop.json.xz", ), - "bench-stop-cmds", + "bench-stop-cmds.json.xz", ) if self.cs_type in ("client", "server"): @@ -625,7 +626,7 @@ def _parse_tool_commands(self, cmds_file): raise EngineError( "Failed to load tool commands from %s: %s" % (cmds_file, err) ) - return [(t["name"], t["command"]) for t in data.get("tools", [])] + return [(t["name"], t["argv"]) for t in data.get("tools", [])] def start_tools(self, one_tool=None): logger.info("Starting tools") @@ -645,7 +646,7 @@ def start_tools(self, one_tool=None): tools = self._parse_tool_commands(self.tool_start_cmds) total = 0 - for tool_name, tool_command in tools: + for tool_name, tool_argv in tools: if one_tool and one_tool != tool_name: logger.info( "Skipping tool '%s' (engine runs '%s' only)", @@ -658,14 +659,14 @@ def start_tools(self, one_tool=None): tool_dir = os.path.join("tool-data", tool_name) os.makedirs(tool_dir, exist_ok=True) logger.info("Starting tool '%s'", tool_name) - run_command("cd %s && %s" % (tool_dir, tool_command)) + run_command("cd %s && %s" % (tool_dir, shlex.join(tool_argv))) if total == 0: logger.info("No tools configured for this engine") # ---- Benchmark execution --------------------------------------------- - def run_bench_cmd(self, matching_type, cmd_type, cmd, force=False): + def run_bench_cmd(self, matching_type, cmd_type, argv, force=False): if self.cs_type != matching_type: return 0 if not force and (self.abort or self.quit): @@ -676,18 +677,21 @@ def run_bench_cmd(self, matching_type, cmd_type, cmd, force=False): self.quit, ) return 0 - if not cmd: + if not argv: logger.info("No %s command to run", cmd_type) return 0 logger.info("Running %s command", cmd_type) - result = run_command(cmd) + result = run_command(shlex.join(argv)) return result.return_code def _load_bench_cmds(self, filename): if not os.path.exists(filename): return [] - with open(filename) as fp: - return [line.strip() for line in fp if line.strip()] + data, err = load_json_file(filename, uselzma=True) + if data is None: + logger.error("Failed to load bench commands from %s: %s", filename, err) + return [] + return data def _roadblock_and_evaluate(self, rb_name, timeout, iter_idx, sample_data, msgs_file=None, do_abort=False): @@ -731,10 +735,10 @@ def process_bench_roadblocks(self): rc = self.do_roadblock("setup-bench-begin", self.default_timeout) self.roadblock_exit_on_error(rc) - bench_start_cmds = self._load_bench_cmds("bench-start-cmds") - bench_infra_cmds = self._load_bench_cmds("bench-infra-cmds") - bench_runtime_cmds = self._load_bench_cmds("bench-runtime-cmds") - bench_stop_cmds = self._load_bench_cmds("bench-stop-cmds") + bench_start_cmds = self._load_bench_cmds("bench-start-cmds.json.xz") + bench_infra_cmds = self._load_bench_cmds("bench-infra-cmds.json.xz") + bench_runtime_cmds = self._load_bench_cmds("bench-runtime-cmds.json.xz") + bench_stop_cmds = self._load_bench_cmds("bench-stop-cmds.json.xz") if not bench_start_cmds: self.abort_error("bench-start-cmds not found", "setup-bench-end") @@ -742,8 +746,8 @@ def process_bench_roadblocks(self): total_tests = len(bench_start_cmds) sample_data = [] - for i, line in enumerate(bench_start_cmds): - iter_samp = line.split()[0] + for i, entry in enumerate(bench_start_cmds): + iter_samp = entry["test"] parts = iter_samp.split("-") iter_id = int(parts[0]) samp_id = int(parts[1]) @@ -772,7 +776,7 @@ def process_bench_roadblocks(self): sd = sample_data[i] iter_id = sd["iteration-id"] samp_id = sd["sample-id"] - iter_samp = bench_start_cmds[i].split()[0] + iter_samp = bench_start_cmds[i]["test"] iter_samp_dir = os.path.join( self.cs_dir, @@ -784,17 +788,17 @@ def process_bench_roadblocks(self): sd["rx-msgs-dir"] = cs_rx_msgs_dir if self.cs_type == "client": - start_cmd = bench_start_cmds[i].split(None, 1)[1] if len(bench_start_cmds[i].split()) > 1 else "" - runtime_cmd = bench_runtime_cmds[i].split(None, 1)[1] if i < len(bench_runtime_cmds) and len(bench_runtime_cmds[i].split()) > 1 else "" - infra_cmd = bench_infra_cmds[i].split(None, 1)[1] if i < len(bench_infra_cmds) and len(bench_infra_cmds[i].split()) > 1 else "" - stop_cmd = "" + start_argv = bench_start_cmds[i].get("argv", []) + runtime_argv = bench_runtime_cmds[i].get("argv", []) if i < len(bench_runtime_cmds) else [] + infra_argv = bench_infra_cmds[i].get("argv", []) if i < len(bench_infra_cmds) else [] + stop_argv = [] elif self.cs_type == "server": - start_cmd = bench_start_cmds[i].split(None, 1)[1] if len(bench_start_cmds[i].split()) > 1 else "" - stop_cmd = bench_stop_cmds[i].split(None, 1)[1] if i < len(bench_stop_cmds) and len(bench_stop_cmds[i].split()) > 1 else "" - runtime_cmd = "" - infra_cmd = "" + start_argv = bench_start_cmds[i].get("argv", []) + stop_argv = bench_stop_cmds[i].get("argv", []) if i < len(bench_stop_cmds) else [] + runtime_argv = [] + infra_argv = [] else: - start_cmd = runtime_cmd = infra_cmd = stop_cmd = "" + start_argv = runtime_argv = infra_argv = stop_argv = [] self.abort = False @@ -844,7 +848,7 @@ def process_bench_roadblocks(self): ) self._roadblock_and_evaluate(rb_name, timeout, i, sample_data, msgs_file) - abort_rc = self.run_bench_cmd("client", "infra", infra_cmd) + abort_rc = self.run_bench_cmd("client", "infra", infra_argv) do_abort_arg = abort_rc != 0 if do_abort_arg: self.abort = True @@ -865,7 +869,7 @@ def process_bench_roadblocks(self): ) self._roadblock_and_evaluate(rb_name, timeout, i, sample_data, msgs_file) - abort_rc = self.run_bench_cmd("server", "server", start_cmd) + abort_rc = self.run_bench_cmd("server", "server", start_argv) do_abort_arg = abort_rc != 0 if do_abort_arg: self.abort = True @@ -900,8 +904,8 @@ def process_bench_roadblocks(self): and not self.quit and self.cs_type == "client" and self.cs_id == "1" - and runtime_cmd): - result = run_command(runtime_cmd) + and runtime_argv): + result = run_command(shlex.join(runtime_argv)) runtime_output = result.stdout.strip() if result.return_code == 0 and runtime_output: @@ -959,11 +963,11 @@ def process_bench_roadblocks(self): msgs_file = prepare_user_msgs_file( cs_tx_msgs_dir, iter_samp_dir, rb_name, default_recipients ) - wait_for_cmd = ( - "python3 /usr/local/bin/engine_lib.py" - " run_bench_cmd '%s' 'client' 'client' '%s' '%s' '0' '%s'" - % (self.cs_type, self.abort, self.quit, start_cmd) - ) + wait_for_cmd = [ + "python3", "/usr/local/bin/engine_lib.py", "run_bench_cmd", + self.cs_type, "client", "client", + str(self.abort), str(self.quit), "0", + ] + start_argv rc = self.do_roadblock( rb_name, timeout, messages=msgs_file, wait_for=wait_for_cmd, @@ -978,7 +982,7 @@ def process_bench_roadblocks(self): if result["is_abort"]: self.abort = True else: - abort_rc = self.run_bench_cmd("client", "client", start_cmd) + abort_rc = self.run_bench_cmd("client", "client", start_argv) do_abort_arg = abort_rc != 0 if do_abort_arg: self.abort = True @@ -1030,7 +1034,7 @@ def process_bench_roadblocks(self): self._roadblock_and_evaluate(rb_name, timeout, i, sample_data, msgs_file) self.run_bench_cmd( - "server", "server", stop_cmd, force=force_server_stop + "server", "server", stop_argv, force=force_server_stop ) abort_rc = 0 do_abort_arg = self.abort @@ -1169,7 +1173,7 @@ def cli_stop_tools(working_dir, tool_cmds_file, disabled, one_tool=""): for tool in data.get("tools", []): tool_name = tool["name"] - tool_command = tool["command"] + tool_argv = tool["argv"] if one_tool and one_tool != tool_name: log.info("Skipping tool '%s' (engine runs '%s' only)", tool_name, one_tool) @@ -1180,7 +1184,7 @@ def cli_stop_tools(working_dir, tool_cmds_file, disabled, one_tool=""): continue log.info("Stopping tool '%s'", tool_name) - run_command("cd %s && %s" % (tool_dir, tool_command)) + run_command("cd %s && %s" % (tool_dir, shlex.join(tool_argv))) if os.path.isfile(env_file): shutil.copy2(env_file, tool_dir) @@ -1230,7 +1234,7 @@ def cli_send_data(ssh_id_file, src_dir, dest_host, dest_path): def cli_run_bench_cmd(cs_type, matching_type, cmd_type, abort_str, quit_str, - force_str, cmd): + force_str, *cmd_argv): """Standalone entry point for roadblock wait-for: run bench command.""" logging.basicConfig(level=logging.INFO, format="%(message)s") if cs_type != matching_type: @@ -1240,9 +1244,9 @@ def cli_run_bench_cmd(cs_type, matching_type, cmd_type, abort_str, quit_str, force = force_str not in ("0", "False", "false", "") if not force and (abort or quit_flag): return - if not cmd: + if not cmd_argv: return - result = run_command(cmd) + result = run_command(shlex.join(cmd_argv)) sys.exit(result.return_code) diff --git a/rickshaw-run.py b/rickshaw-run.py index bc7ca193..307e9167 100755 --- a/rickshaw-run.py +++ b/rickshaw-run.py @@ -12,6 +12,7 @@ import platform import random import re +import shlex import shutil import signal import subprocess @@ -106,6 +107,19 @@ def dir_entries(dirpath, pattern=None): return entries +def render_param(arg, val): + """Render one {arg, val} pair as a single '--arg' or '--arg=value' token, + quoting the value only if the shell would otherwise misinterpret it + (e.g. embedded spaces). Shared by benchmark and tool command rendering + so both use the same grammar for param_regex fixups to target. Only + None/"" render as a bare flag -- any other value (including 0/False) is + stringified and rendered, never silently dropped via Python truthiness.""" + val = "" if val is None else str(val) + if val: + return f"--{arg}={shlex.quote(val)}" + return f"--{arg}" + + def dump_params(params, cs_id, engine, ids_to_benchmark): default_role = "client" benchmark = ids_to_benchmark.get(str(cs_id)) if cs_id is not None else None @@ -113,7 +127,7 @@ def dump_params(params, cs_id, engine, ids_to_benchmark): for param in params: arg = param.get("arg", "") - val = param.get("val", "") + val = param.get("val") bench = param.get("benchmark", "") role = param.get("role", default_role) param_id = param.get("id") @@ -125,12 +139,9 @@ def dump_params(params, cs_id, engine, ids_to_benchmark): if role != engine and role != "all": continue - if val: - if cs_id is not None: - val = val.replace("%client-id%", str(cs_id)) - params_str += f" --{arg}={val}" - else: - params_str += f" --{arg}" + if val is not None and cs_id is not None: + val = str(val).replace("%client-id%", str(cs_id)) + params_str += f" {render_param(arg, val)}" return params_str.lstrip() @@ -147,6 +158,24 @@ def perl_s_regex(cmd, regex_str): return re.sub(pattern, replacement, cmd, count=count, flags=re_flags) +def apply_param_regex_and_split(cmd, param_regex_list, context_desc): + """Apply param_regex fixups to a fully-rendered command string, then + split it into a clean argv list. A param_regex pattern anchored on + non-whitespace (e.g. \\S+) can leave an unbalanced quote if it ever + matches into a value that needed shlex.quote() protection -- fail + loudly with the offending text rather than silently mis-tokenizing.""" + for r in param_regex_list: + cmd = perl_s_regex(cmd, r) + try: + return shlex.split(cmd) + except ValueError as exc: + logger.error( + "[ERROR] Could not parse rendered command for %s: %s -- rendered text: [%s]", + context_desc, exc, cmd, + ) + sys.exit(1) + + class RunState: def __init__(self): self.run = {} @@ -1468,7 +1497,7 @@ def build_tool_cmd(self, tool_entry, start_stop, endpoint_type=None): tool = { "name": tool_id, - "command": "declare -a ARGS=(", + "argv": [], "deployment": "auto", "opt-tag": None, } @@ -1478,14 +1507,17 @@ def build_tool_cmd(self, tool_entry, start_stop, endpoint_type=None): if "opt-tag" in tool_entry: tool["opt-tag"] = tool_entry["opt-tag"] + tokens = [self.tools_configs[tool_name]["collector"][start_stop]] for tool_param in tool_entry.get("params", []): if tool_param.get("enabled") == "no": continue - tool["command"] += f"'--{tool_param['arg']}' '{tool_param.get('val', '')}' " + tokens.append(render_param(tool_param["arg"], tool_param.get("val", ""))) + cmd = " ".join(tokens) - tool["command"] = tool["command"].rstrip() - tool["command"] += ") && " - tool["command"] += self.tools_configs[tool_name]["collector"][start_stop] + ' "${ARGS[@]}"' + param_regex_list = self.tools_configs[tool_name].get("collector", {}).get("param_regex", []) + tool["argv"] = apply_param_regex_and_split( + cmd, param_regex_list, "tool '%s' (%s)" % (tool_id, start_stop) + ) return tool @@ -1649,29 +1681,31 @@ def prepare_bench_tool_engines(self): for cmd_type in cmd_type_files: if cmd_type == "runtime" and int(cs_id) > 1: continue - this_cmd_file = os.path.join(this_cmds_dir, cmd_type) - with open(this_cmd_file, "w") as fh: - for test_ref in self.tests: - test_iter = test_ref["iteration-id"] - test_samp = test_ref["sample-id"] - iter_array_idx = test_iter - 1 - benchmark = self.ids_to_benchmark.get(str(cs_id)) - bench_cfg = self.bench_configs.get(benchmark, {}) - cmd_template = bench_cfg.get(cs_type, {}).get(cmd_type, "") - if cmd_template: - params_str = dump_params( - self.run["iterations"][iter_array_idx].get("params", []), - cs_id, cs_type, self.ids_to_benchmark - ) - cmd = f"{cmd_template} {params_str}" - param_regex_list = bench_cfg.get(cs_type, {}).get("param_regex", []) - for r in param_regex_list: - cmd = perl_s_regex(cmd, r) - fh.write(f"{test_iter}-{test_samp} {cmd}\n") - elif cmd_type != "infra": - logger.error("[ERROR] Could not find %s in bench_config", cmd_type) - sys.exit(1) - os.chmod(this_cmd_file, 0o755) + this_cmd_file = os.path.join(this_cmds_dir, f"{cmd_type}.json") + entries = [] + for test_ref in self.tests: + test_iter = test_ref["iteration-id"] + test_samp = test_ref["sample-id"] + iter_array_idx = test_iter - 1 + benchmark = self.ids_to_benchmark.get(str(cs_id)) + bench_cfg = self.bench_configs.get(benchmark, {}) + cmd_template = bench_cfg.get(cs_type, {}).get(cmd_type, "") + if cmd_template: + params_str = dump_params( + self.run["iterations"][iter_array_idx].get("params", []), + cs_id, cs_type, self.ids_to_benchmark + ) + cmd = f"{cmd_template} {params_str}" + param_regex_list = bench_cfg.get(cs_type, {}).get("param_regex", []) + argv = apply_param_regex_and_split( + cmd, param_regex_list, + "benchmark '%s' (%s/%s)" % (benchmark, cs_type, cmd_type) + ) + entries.append({"test": f"{test_iter}-{test_samp}", "argv": argv}) + elif cmd_type != "infra": + logger.error("[ERROR] Could not find %s in bench_config", cmd_type) + sys.exit(1) + save_json_file(this_cmd_file, entries) for cs_type in list(self.clients_servers.keys()) + all_collector_types: if cs_type in ("client", "server"): diff --git a/schema/benchmark.json b/schema/benchmark.json index a3da5546..4042c898 100644 --- a/schema/benchmark.json +++ b/schema/benchmark.json @@ -108,7 +108,7 @@ "pattern": "^.+$" }, "param_regex": { - "description": "Array of regex patterns used to validate benchmark parameters. Each pattern defines valid parameter formats.", + "description": "Array of sed-style 's/pattern/replacement/[flags]' substitutions applied to the fully-rendered client command string (params rendered as '--arg=value' tokens) before it is split back into an argv list. Lets a param declare a placeholder value (e.g. 'ON'/'OFF') for multiplex validation purposes while being rendered as a bare flag or omitted entirely, or lets the benchmark rewrite/strip params outright (e.g. swapping a controller-supplied path for a fixed in-container filename). Shares the same grammar as tool.json's collector 'param_regex'.", "type": "array", "items": { "type": "string", @@ -169,7 +169,7 @@ "pattern": "^.+$" }, "param_regex": { - "description": "Array of regex patterns used to validate server-side benchmark parameters.", + "description": "Array of sed-style 's/pattern/replacement/[flags]' substitutions applied to the fully-rendered server command string (params rendered as '--arg=value' tokens) before it is split back into an argv list. Lets a param declare a placeholder value (e.g. 'ON'/'OFF') for multiplex validation purposes while being rendered as a bare flag or omitted entirely, or lets the benchmark rewrite/strip params outright. Shares the same grammar as tool.json's collector 'param_regex'.", "type": "array", "items": { "type": "string", diff --git a/schema/tool.json b/schema/tool.json index 43294d64..be5d3531 100644 --- a/schema/tool.json +++ b/schema/tool.json @@ -173,6 +173,15 @@ "description": "The script that terminates data collection on the collector engine and prepares the collected data for transfer back to the controller.", "type": "string", "pattern": "^.+$" + }, + "param_regex": { + "description": "Array of sed-style 's/pattern/replacement/[flags]' substitutions applied to the fully-rendered start/stop command string (params rendered as '--arg=value' tokens) before it is split back into an argv list. Lets a param declare a placeholder value (e.g. 'ON'/'OFF') for multiplex validation purposes while being rendered as a bare flag or omitted entirely, for scripts whose getopt parsing doesn't accept a real value. Shares the same grammar as benchmark.json's client/server 'param_regex'.", + "type": "array", + "items": { + "type": "string", + "pattern": "^.+$" + }, + "additionalItems": false } }, "required": [ diff --git a/tests/test_build_tool_cmd.py b/tests/test_build_tool_cmd.py new file mode 100644 index 00000000..2c4acf3c --- /dev/null +++ b/tests/test_build_tool_cmd.py @@ -0,0 +1,261 @@ +#!/usr/bin/env python3 +# -*- mode: python; indent-tabs-mode: nil; python-indent-level: 4 -*- +# vim: autoindent tabstop=4 shiftwidth=4 expandtab softtabstop=4 filetype=python + +"""Unit tests for rickshaw-run.py's RunState.build_tool_cmd(). + +Tools render params the same way benchmarks do -- '--arg=value' tokens via +render_param(), quoted only when needed -- rather than the old bash-syntax +'declare -a ARGS=(...)' construction. "collector.param_regex" (rickshaw#867) +lets a tool param declare a placeholder value (e.g. "ON"/"OFF") for +multiplex validation purposes while build_tool_cmd() rewrites the rendered +command via sed-style regex before splitting it back into an argv list -- +sharing the same grammar as benchmark.json's client/server "param_regex". + +toolbox is mocked out rather than required, since rickshaw-run.py imports +from it at module scope and CI does not check toolbox out for this job. +""" + +import importlib.machinery +import importlib.util +import logging +import sys +import types +import unittest + + +def import_rickshaw_run(): + """Load rickshaw-run.py as a module with toolbox mocked out.""" + mock_fileio = types.ModuleType("toolbox.fileio") + mock_fileio.open_write_text_file = lambda *a, **k: None + + mock_json = types.ModuleType("toolbox.json") + mock_json.load_json_file = lambda *a, **k: (None, "not used in this test") + mock_json.save_json_file = lambda *a, **k: None + mock_json.validate_schema = lambda *a, **k: (True, None) + + mock_jsonsettings = types.ModuleType("toolbox.jsonsettings") + mock_jsonsettings.get_json_setting = lambda *a, **k: None + + mock_logging_mod = types.ModuleType("toolbox.logging") + mock_logging_mod.setup_logging = lambda *a, **k: None + + mock_roadblock = types.ModuleType("toolbox.roadblock") + mock_roadblock.do_roadblock = lambda *a, **k: (0, None) + mock_roadblock.ROADBLOCK_EXITS = { + "success": 0, "input": 2, "timeout": 3, + "abort": 4, "heartbeat_timeout": 5, "abort_waiting": 6, + } + + mock_run = types.ModuleType("toolbox.run") + mock_run.run_cmd = lambda *a, **k: ("cmd", "", 0) + + mock_toolbox = types.ModuleType("toolbox") + mock_toolbox.fileio = mock_fileio + mock_toolbox.json = mock_json + mock_toolbox.jsonsettings = mock_jsonsettings + mock_toolbox.logging = mock_logging_mod + mock_toolbox.roadblock = mock_roadblock + mock_toolbox.run = mock_run + + mod_name = "rickshaw_run_under_test_build_tool_cmd" + sys.modules.pop(mod_name, None) + + mocks = { + "toolbox": mock_toolbox, + "toolbox.fileio": mock_fileio, + "toolbox.json": mock_json, + "toolbox.jsonsettings": mock_jsonsettings, + "toolbox.logging": mock_logging_mod, + "toolbox.roadblock": mock_roadblock, + "toolbox.run": mock_run, + } + saved = {key: sys.modules.get(key) for key in mocks} + sys.modules.update(mocks) + + try: + import os + script_path = os.path.join(os.path.dirname(__file__), "..", "rickshaw-run.py") + loader = importlib.machinery.SourceFileLoader(mod_name, script_path) + spec = importlib.util.spec_from_loader(mod_name, loader) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + finally: + for key, val in saved.items(): + if val is None: + sys.modules.pop(key, None) + else: + sys.modules[key] = val + + mod.logger = logging.getLogger("test_build_tool_cmd") + return mod + + +class TestBuildToolCmd(unittest.TestCase): + def setUp(self): + self.mod = import_rickshaw_run() + self.state = self.mod.RunState() + + def _collector(self, **overrides): + collector = {"start": "kerneltools-start", "stop": "kerneltools-stop"} + collector.update(overrides) + return collector + + def test_no_param_regex_configured_renders_normally(self): + self.state.tools_configs = {"kernel": {"collector": self._collector()}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [{"arg": "interval", "val": "5"}], + } + tool = self.state.build_tool_cmd(tool_entry, "start") + self.assertEqual(tool["argv"], ["kerneltools-start", "--interval=5"]) + + def test_bare_flag_with_no_value(self): + self.state.tools_configs = {"kernel": {"collector": self._collector()}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [{"arg": "some-flag", "val": ""}], + } + tool = self.state.build_tool_cmd(tool_entry, "start") + self.assertEqual(tool["argv"], ["kerneltools-start", "--some-flag"]) + + def test_zero_value_is_not_silently_dropped(self): + # a code-review catch: an int/JSON-number val of 0 is "falsy" in + # Python but is a real, meaningful value -- must not collapse to a + # bare flag + self.state.tools_configs = {"kernel": {"collector": self._collector()}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [{"arg": "interval", "val": 0}], + } + tool = self.state.build_tool_cmd(tool_entry, "start") + self.assertEqual(tool["argv"], ["kerneltools-start", "--interval=0"]) + + def test_disabled_param_is_skipped(self): + self.state.tools_configs = {"kernel": {"collector": self._collector()}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [ + {"arg": "interval", "val": "5"}, + {"arg": "skip-me", "val": "1", "enabled": "no"}, + ], + } + tool = self.state.build_tool_cmd(tool_entry, "start") + self.assertEqual(tool["argv"], ["kerneltools-start", "--interval=5"]) + + def test_value_with_space_survives_round_trip(self): + self.state.tools_configs = {"kernel": {"collector": self._collector()}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [{"arg": "record-opts", "val": "-e cycles -c 1000"}], + } + tool = self.state.build_tool_cmd(tool_entry, "start") + self.assertEqual( + tool["argv"], ["kerneltools-start", "--record-opts=-e cycles -c 1000"] + ) + + def test_param_regex_on_strips_value_keeps_bare_flag(self): + # converged grammar: identical text to benchmarks' shipped ON/OFF + # patterns (fio, trafficgen), now that both sides use the same + # '--arg=value' rendering + collector = self._collector(param_regex=[ + r"s/(\s--[^\s]+)=ON/$1/g", + r"s/\s--[^\s]+=OFF//g", + ]) + self.state.tools_configs = {"kernel": {"collector": collector}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [ + {"arg": "subtools", "val": "perf"}, + {"arg": "perf-gen-local-report", "val": "ON"}, + ], + } + tool = self.state.build_tool_cmd(tool_entry, "stop") + self.assertEqual( + tool["argv"], + ["kerneltools-stop", "--subtools=perf", "--perf-gen-local-report"], + ) + + def test_param_regex_off_removes_flag_and_value(self): + collector = self._collector(param_regex=[ + r"s/(\s--[^\s]+)=ON/$1/g", + r"s/\s--[^\s]+=OFF//g", + ]) + self.state.tools_configs = {"kernel": {"collector": collector}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [ + {"arg": "subtools", "val": "perf"}, + {"arg": "perf-gen-local-report", "val": "OFF"}, + ], + } + tool = self.state.build_tool_cmd(tool_entry, "stop") + self.assertEqual(tool["argv"], ["kerneltools-stop", "--subtools=perf"]) + + def test_param_regex_is_generic_across_multiple_flags(self): + # the shipped patterns match any '--flag' name, not just one hardcoded + # param, so adding a second ON/OFF-style tool param never requires a + # rickshaw.json param_regex edit + collector = self._collector(param_regex=[ + r"s/(\s--[^\s]+)=ON/$1/g", + r"s/\s--[^\s]+=OFF//g", + ]) + self.state.tools_configs = {"kernel": {"collector": collector}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [ + {"arg": "foo-flag", "val": "ON"}, + {"arg": "subtools", "val": "perf"}, + {"arg": "bar-flag", "val": "OFF"}, + {"arg": "perf-gen-local-report", "val": "ON"}, + ], + } + tool = self.state.build_tool_cmd(tool_entry, "stop") + self.assertEqual( + tool["argv"], + [ + "kerneltools-stop", + "--foo-flag", + "--subtools=perf", + "--perf-gen-local-report", + ], + ) + + def test_param_regex_leaving_unbalanced_quote_raises(self): + # a \S+-anchored pattern (the fio 'jobfile'-class fragility found + # during design review) can leave a dangling unbalanced quote if it + # ever matches into a quoted (space-containing) value -- this must + # fail loudly via sys.exit, not silently mis-tokenize + collector = self._collector(param_regex=[ + r"s/\s--record-opts=(\S+)//", + ]) + self.state.tools_configs = {"kernel": {"collector": collector}} + tool_entry = { + "tool": "kernel", + "tool-id": "kernel", + "params": [{"arg": "record-opts", "val": "-e cycles -c 1000"}], + } + with self.assertRaises(SystemExit): + self.state.build_tool_cmd(tool_entry, "start") + + def test_blacklisted_endpoint_returns_none(self): + collector = self._collector(blacklist=[ + {"endpoint": "remotehosts", "collector-types": ["client", "server"]}, + ]) + self.state.tools_configs = {"kernel": {"collector": collector}} + tool_entry = {"tool": "kernel", "tool-id": "kernel", "params": []} + tool = self.state.build_tool_cmd(tool_entry, "start", endpoint_type="remotehosts") + self.assertIsNone(tool) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_dump_params.py b/tests/test_dump_params.py new file mode 100644 index 00000000..1fc9b5fb --- /dev/null +++ b/tests/test_dump_params.py @@ -0,0 +1,183 @@ +#!/usr/bin/env python3 +# -*- mode: python; indent-tabs-mode: nil; python-indent-level: 4 -*- +# vim: autoindent tabstop=4 shiftwidth=4 expandtab softtabstop=4 filetype=python + +"""Unit tests for rickshaw-run.py's dump_params()/render_param(). + +Covers rickshaw#867: a param value containing a space must be rendered with +enough quoting to survive both shell execution and the shlex.split() round +trip used to rebuild the argv list shipped to the engine, while every +existing (space-free) value renders byte-for-byte identically to before. + +toolbox is mocked out rather than required, since rickshaw-run.py imports +from it at module scope and CI does not check toolbox out for this job. +""" + +import importlib.machinery +import importlib.util +import logging +import shlex +import sys +import types +import unittest + + +def import_rickshaw_run(): + """Load rickshaw-run.py as a module with toolbox mocked out.""" + mock_fileio = types.ModuleType("toolbox.fileio") + mock_fileio.open_write_text_file = lambda *a, **k: None + + mock_json = types.ModuleType("toolbox.json") + mock_json.load_json_file = lambda *a, **k: (None, "not used in this test") + mock_json.save_json_file = lambda *a, **k: None + mock_json.validate_schema = lambda *a, **k: (True, None) + + mock_jsonsettings = types.ModuleType("toolbox.jsonsettings") + mock_jsonsettings.get_json_setting = lambda *a, **k: None + + mock_logging_mod = types.ModuleType("toolbox.logging") + mock_logging_mod.setup_logging = lambda *a, **k: None + + mock_roadblock = types.ModuleType("toolbox.roadblock") + mock_roadblock.do_roadblock = lambda *a, **k: (0, None) + mock_roadblock.ROADBLOCK_EXITS = { + "success": 0, "input": 2, "timeout": 3, + "abort": 4, "heartbeat_timeout": 5, "abort_waiting": 6, + } + + mock_run = types.ModuleType("toolbox.run") + mock_run.run_cmd = lambda *a, **k: ("cmd", "", 0) + + mock_toolbox = types.ModuleType("toolbox") + mock_toolbox.fileio = mock_fileio + mock_toolbox.json = mock_json + mock_toolbox.jsonsettings = mock_jsonsettings + mock_toolbox.logging = mock_logging_mod + mock_toolbox.roadblock = mock_roadblock + mock_toolbox.run = mock_run + + mod_name = "rickshaw_run_under_test_dump_params" + sys.modules.pop(mod_name, None) + + mocks = { + "toolbox": mock_toolbox, + "toolbox.fileio": mock_fileio, + "toolbox.json": mock_json, + "toolbox.jsonsettings": mock_jsonsettings, + "toolbox.logging": mock_logging_mod, + "toolbox.roadblock": mock_roadblock, + "toolbox.run": mock_run, + } + saved = {key: sys.modules.get(key) for key in mocks} + sys.modules.update(mocks) + + try: + import os + script_path = os.path.join(os.path.dirname(__file__), "..", "rickshaw-run.py") + loader = importlib.machinery.SourceFileLoader(mod_name, script_path) + spec = importlib.util.spec_from_loader(mod_name, loader) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + finally: + for key, val in saved.items(): + if val is None: + sys.modules.pop(key, None) + else: + sys.modules[key] = val + + mod.logger = logging.getLogger("test_dump_params") + return mod + + +class TestRenderParam(unittest.TestCase): + def setUp(self): + self.mod = import_rickshaw_run() + + def test_bare_flag_with_no_value(self): + self.assertEqual(self.mod.render_param("some-flag", ""), "--some-flag") + + def test_simple_value_unquoted(self): + # every currently-shipped value (ON/OFF, paths, profile names) must + # render byte-for-byte identically to before this change + self.assertEqual(self.mod.render_param("rw", "read"), "--rw=read") + self.assertEqual(self.mod.render_param("iodepth", "4"), "--iodepth=4") + + def test_space_containing_value_is_quoted(self): + token = self.mod.render_param("exec_prerun", "echo starting test") + self.assertEqual(token, "--exec_prerun='echo starting test'") + + def test_quoted_value_round_trips_through_shlex_split(self): + token = self.mod.render_param("exec_prerun", "echo starting test") + self.assertEqual(shlex.split(token), ["--exec_prerun=echo starting test"]) + + def test_none_renders_as_bare_flag(self): + self.assertEqual(self.mod.render_param("some-flag", None), "--some-flag") + + def test_falsy_non_string_values_are_not_silently_dropped(self): + # a code-review catch: val=0/False are "falsy" in Python but are + # real, meaningful values a param author may explicitly want to + # send -- only None/"" mean "no value" + self.assertEqual(self.mod.render_param("interval", 0), "--interval=0") + self.assertEqual(self.mod.render_param("enabled", False), "--enabled=False") + + +class TestDumpParams(unittest.TestCase): + def setUp(self): + self.mod = import_rickshaw_run() + + def test_simple_params_render_unchanged(self): + params = [ + {"arg": "iodepth", "val": "4", "role": "client"}, + {"arg": "rw", "val": "read", "role": "client"}, + ] + result = self.mod.dump_params(params, "1", "client", {}) + self.assertEqual(result, "--iodepth=4 --rw=read") + + def test_null_value_renders_as_bare_flag_not_the_string_none(self): + # a code-review catch: a present-but-null JSON "val" must render as + # a bare flag, matching the pre-existing (falsy) behavior -- not get + # stringified into the literal text "None" before being checked + params = [{"arg": "foo", "val": None, "role": "client"}] + result = self.mod.dump_params(params, "1", "client", {}) + self.assertEqual(result, "--foo") + + def test_zero_value_is_not_silently_dropped(self): + params = [{"arg": "foo", "val": 0, "role": "client"}] + result = self.mod.dump_params(params, "1", "client", {}) + self.assertEqual(result, "--foo=0") + + def test_space_containing_value_survives_the_pipeline(self): + # the concrete rickshaw#867 regression case: exec_prerun's value must + # come back out exactly as it went in after render -> shlex.split + params = [ + {"arg": "exec_prerun", "val": "echo starting test", "role": "client"}, + ] + rendered = self.mod.dump_params(params, "1", "client", {}) + argv = shlex.split(rendered) + self.assertEqual(argv, ["--exec_prerun=echo starting test"]) + + def test_client_id_substitution_still_applies(self): + params = [{"arg": "id", "val": "client-%client-id%", "role": "client"}] + result = self.mod.dump_params(params, "3", "client", {}) + self.assertEqual(result, "--id=client-3") + + def test_role_filtering_excludes_non_matching_role(self): + params = [ + {"arg": "client-only", "val": "1", "role": "client"}, + {"arg": "server-only", "val": "1", "role": "server"}, + ] + result = self.mod.dump_params(params, "1", "client", {}) + self.assertEqual(result, "--client-only=1") + + def test_benchmark_filtering_excludes_non_matching_benchmark(self): + params = [ + {"arg": "fio-only", "val": "1", "role": "client", "benchmark": "fio"}, + {"arg": "uperf-only", "val": "1", "role": "client", "benchmark": "uperf"}, + ] + result = self.mod.dump_params(params, "1", "client", {"1": "fio"}) + self.assertEqual(result, "--fio-only=1") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_endpoints_process_bench_roadblocks.py b/tests/test_endpoints_process_bench_roadblocks.py new file mode 100644 index 00000000..28088076 --- /dev/null +++ b/tests/test_endpoints_process_bench_roadblocks.py @@ -0,0 +1,210 @@ +#!/usr/bin/env python3 +# -*- mode: python; indent-tabs-mode: nil; python-indent-level: 4 -*- +# vim: autoindent tabstop=4 shiftwidth=4 expandtab softtabstop=4 filetype=python + +"""Unit tests for endpoints.py's process_bench_roadblocks() data-structure +initialization (rickshaw#867's third consumer). + +This is a separate, endpoint-side copy of "figure out how many iterations +and samples there are" from engine_lib.py's own process_bench_roadblocks() +-- it only needs the "test" id per entry, not the argv, but it independently +read the old flat-text bench-cmds file directly (bypassing engine_lib.py +entirely) and was missed during the initial argv/JSON conversion, causing a +live end-to-end run to crash with FileNotFoundError against the old +extensionless filename. This test locks in the fix: reading the same +'start.json.xz' file tool/bench commands now use everywhere else. + +toolbox, roadblock, and the fabric/invoke/paramiko third-party deps are all +mocked out rather than required, since endpoints.py imports from them at +module scope and CI does not check any of them out or install them for this +test job. +""" + +import importlib.machinery +import importlib.util +import json +import lzma +import os +import shutil +import sys +import tempfile +import types +import unittest +from unittest.mock import MagicMock + + +def import_endpoints(): + """Load endpoints.py as a module with all of its external deps mocked out.""" + mock_fabric = types.ModuleType("fabric") + mock_fabric.Connection = MagicMock + + mock_invoke = types.ModuleType("invoke") + mock_invoke.run = MagicMock() + + mock_ssh_exception = types.ModuleType("paramiko.ssh_exception") + mock_ssh_exception.AuthenticationException = Exception + mock_ssh_exception.NoValidConnectionsError = Exception + + mock_paramiko = types.ModuleType("paramiko") + mock_paramiko.ssh_exception = mock_ssh_exception + + def _fake_load_json_file(json_file, uselzma=False): + try: + opener = lzma.open if uselzma else open + with opener(json_file, "rt") as fh: + return json.load(fh), None + except (OSError, ValueError) as exc: + return None, str(exc) + + mock_toolbox_json = types.ModuleType("toolbox.json") + mock_toolbox_json.load_json_file = _fake_load_json_file + + mock_toolbox_messages = types.ModuleType("toolbox.messages") + mock_toolbox_messages.create_roadblock_msg = lambda *a, **k: None + mock_toolbox_messages.prepare_user_msgs_file = lambda *a, **k: None + mock_toolbox_messages.evaluate_roadblock_result = lambda *a, **k: None + mock_toolbox_messages.save_received_messages = lambda *a, **k: None + mock_toolbox_messages.ROADBLOCK_EXITS = { + "success": 0, + "input": 2, + "timeout": 3, + "abort": 4, + "heartbeat_timeout": 5, + "abort_waiting": 6, + } + + mock_toolbox_roadblock = types.ModuleType("toolbox.roadblock") + mock_toolbox_roadblock.do_roadblock = MagicMock(return_value=(0, None)) + + mock_toolbox = types.ModuleType("toolbox") + mock_toolbox.json = mock_toolbox_json + mock_toolbox.messages = mock_toolbox_messages + mock_toolbox.roadblock = mock_toolbox_roadblock + + mock_roadblock_engine_mod = types.ModuleType("roadblock") + mock_roadblock_engine_mod.VERBOSE_DEBUG_LEVEL = 5 + + mod_name = "endpoints_under_test_process_bench_roadblocks" + sys.modules.pop(mod_name, None) + + mocks = { + "fabric": mock_fabric, + "invoke": mock_invoke, + "paramiko": mock_paramiko, + "paramiko.ssh_exception": mock_ssh_exception, + "toolbox": mock_toolbox, + "toolbox.json": mock_toolbox_json, + "toolbox.messages": mock_toolbox_messages, + "toolbox.roadblock": mock_toolbox_roadblock, + "roadblock": mock_roadblock_engine_mod, + } + saved = {key: sys.modules.get(key) for key in mocks} + sys.modules.update(mocks) + + with tempfile.TemporaryDirectory() as tmp_home: + toolbox_python_dir = os.path.join(tmp_home, "python") + os.makedirs(toolbox_python_dir) + roadblock_dir = os.path.join(tmp_home, "roadblock") + os.makedirs(roadblock_dir) + open(os.path.join(roadblock_dir, "roadblock.py"), "w").close() + + saved_env = { + "TOOLBOX_HOME": os.environ.get("TOOLBOX_HOME"), + "ROADBLOCK_HOME": os.environ.get("ROADBLOCK_HOME"), + } + os.environ["TOOLBOX_HOME"] = tmp_home + os.environ["ROADBLOCK_HOME"] = roadblock_dir + + try: + script_path = os.path.join( + os.path.dirname(__file__), "..", "endpoints", "endpoints.py" + ) + loader = importlib.machinery.SourceFileLoader(mod_name, script_path) + spec = importlib.util.spec_from_loader(mod_name, loader) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + finally: + for key, val in saved.items(): + if val is None: + sys.modules.pop(key, None) + else: + sys.modules[key] = val + for key, val in saved_env.items(): + if val is None: + os.environ.pop(key, None) + else: + os.environ[key] = val + + return mod + + +class TestProcessBenchRoadblocksInit(unittest.TestCase): + def setUp(self): + self.mod = import_endpoints() + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir, ignore_errors=True) + self.msgs_dir = self.tmpdir + + def _write_start_json(self, entries): + client_dir = os.path.join(self.tmpdir, "client", "1") + os.makedirs(client_dir, exist_ok=True) + path = os.path.join(client_dir, "start.json.xz") + with lzma.open(path, "wt") as fh: + json.dump(entries, fh) + + def test_reads_new_json_format_and_stops_before_iteration_loop(self): + self._write_start_json([ + {"test": "1-1", "argv": ["fio", "--iodepth=4"]}, + {"test": "2-1", "argv": ["fio", "--iodepth=8"]}, + ]) + + # 0 for setup-bench-begin, then a non-zero rc for setup-bench-end so + # the function returns immediately after the fix's own logic runs, + # without needing to fake the rest of the (unrelated, unchanged) + # iteration loop + self.mod.do_roadblock = MagicMock(side_effect=[0, 1]) + + with self.assertLogs(logger=self.mod.logger, level="INFO") as cm: + rc = self.mod.process_bench_roadblocks( + callbacks={}, + roadblock_id="run-1", + endpoint_label="kube-1", + roadblock_password="secret", + max_sample_failures=1, + roadblock_messages_dir=self.msgs_dir, + roadblock_timeouts={"default": 60}, + engine_commands_dir=self.tmpdir, + endpoint_dir=self.tmpdir, + roadblock_connection_watchdog=True, + ) + + self.assertEqual(rc, 1) + joined = "\n".join(cm.output) + self.assertIn("iteration_sample=1-1 iteration_id=1 sample_id=1", joined) + self.assertIn("iteration_sample=2-1 iteration_id=2 sample_id=1", joined) + self.assertIn("Total tests: 2", joined) + + def test_missing_file_logs_error_and_returns_nonzero(self): + # no start.json.xz written at all + self.mod.do_roadblock = MagicMock(return_value=0) + + with self.assertLogs(logger=self.mod.logger, level="ERROR"): + rc = self.mod.process_bench_roadblocks( + callbacks={}, + roadblock_id="run-1", + endpoint_label="kube-1", + roadblock_password="secret", + max_sample_failures=1, + roadblock_messages_dir=self.msgs_dir, + roadblock_timeouts={"default": 60}, + engine_commands_dir=self.tmpdir, + endpoint_dir=self.tmpdir, + roadblock_connection_watchdog=True, + ) + + self.assertEqual(rc, 1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_engine_lib.py b/tests/test_engine_lib.py new file mode 100644 index 00000000..2a7d7321 --- /dev/null +++ b/tests/test_engine_lib.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +# -*- mode: python; indent-tabs-mode: nil; python-indent-level: 4 -*- +# vim: autoindent tabstop=4 shiftwidth=4 expandtab softtabstop=4 filetype=python + +"""Unit tests for engine/engine_lib.py's argv-based command handling +(rickshaw#867's engine-side half): tool/bench commands are shipped as plain +JSON argv lists rather than pre-rendered shell strings, and shlex.join() is +applied exactly once, immediately before each run_command() call. + +Also covers the wait_for double-quoting hazard found while designing this: +process_bench_roadblocks()'s unbounded-timeout path must hand roadblock a +literal Python list (which toolbox.roadblock passes straight to +subprocess.Popen with no re-parsing), not a pre-rendered string that would +need re-splitting. + +fabric, invoke, and toolbox are mocked out rather than required, since +engine_lib.py imports from them at module scope and CI does not have real +engine dependencies available for this job. +""" + +import importlib.machinery +import importlib.util +import json +import logging +import os +import shutil +import sys +import tempfile +import types +import unittest +from unittest import mock + + +def import_engine_lib(): + """Load engine/engine_lib.py as a module with its dependencies mocked out.""" + mock_fabric = types.ModuleType("fabric") + mock_fabric.Connection = mock.MagicMock() + + mock_invoke = types.ModuleType("invoke") + mock_invoke.run = mock.MagicMock() + + def _fake_load_json_file(json_file, uselzma=False): + import lzma + try: + opener = lzma.open if uselzma else open + with opener(json_file, "rt") as fh: + return json.load(fh), None + except (OSError, ValueError) as exc: + return None, str(exc) + + mock_json = types.ModuleType("toolbox.json") + mock_json.load_json_file = _fake_load_json_file + + mock_jsonsettings = types.ModuleType("toolbox.jsonsettings") + mock_jsonsettings.get_json_setting = lambda *a, **k: (None, 0) + + mock_messages = types.ModuleType("toolbox.messages") + mock_messages.ROADBLOCK_EXITS = { + "success": 0, "input": 2, "timeout": 3, + "abort": 4, "heartbeat_timeout": 5, "abort_waiting": 6, + } + mock_messages.evaluate_roadblock_result = lambda *a, **k: { + "is_timeout": False, "is_abort": False, "messages": None, + } + mock_messages.prepare_user_msgs_file = lambda *a, **k: None + mock_messages.resolve_svc_messages = lambda *a, **k: None + mock_messages.save_received_messages = lambda *a, **k: None + + mock_roadblock = types.ModuleType("toolbox.roadblock") + mock_roadblock.do_roadblock = lambda *a, **k: (0, None) + + mock_toolbox = types.ModuleType("toolbox") + mock_toolbox.json = mock_json + mock_toolbox.jsonsettings = mock_jsonsettings + mock_toolbox.messages = mock_messages + mock_toolbox.roadblock = mock_roadblock + + mod_name = "engine_lib_under_test" + sys.modules.pop(mod_name, None) + + mocks = { + "fabric": mock_fabric, + "invoke": mock_invoke, + "toolbox": mock_toolbox, + "toolbox.json": mock_json, + "toolbox.jsonsettings": mock_jsonsettings, + "toolbox.messages": mock_messages, + "toolbox.roadblock": mock_roadblock, + } + saved = {key: sys.modules.get(key) for key in mocks} + sys.modules.update(mocks) + + try: + script_path = os.path.join( + os.path.dirname(__file__), "..", "engine", "engine_lib.py" + ) + loader = importlib.machinery.SourceFileLoader(mod_name, script_path) + spec = importlib.util.spec_from_loader(mod_name, loader) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + finally: + for key, val in saved.items(): + if val is None: + sys.modules.pop(key, None) + else: + sys.modules[key] = val + + mod.logger = logging.getLogger("test_engine_lib") + return mod + + +class TestParseToolCommands(unittest.TestCase): + def setUp(self): + self.mod = import_engine_lib() + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir, ignore_errors=True) + + def _write_json_xz(self, name, data): + import lzma + path = os.path.join(self.tmpdir, name) + with lzma.open(path, "wt") as fh: + json.dump(data, fh) + return path + + def test_returns_name_argv_tuples(self): + path = self._write_json_xz("start.json.xz", { + "tools": [ + {"name": "kernel-1", "argv": ["kerneltools-start", "--interval=5"], + "deployment": "auto", "opt-tag": None}, + ] + }) + engine = self.mod.Engine() + result = engine._parse_tool_commands(path) + self.assertEqual( + result, [("kernel-1", ["kerneltools-start", "--interval=5"])] + ) + + def test_missing_file_raises_engine_error(self): + engine = self.mod.Engine() + with self.assertRaises(self.mod.EngineError): + engine._parse_tool_commands(os.path.join(self.tmpdir, "missing.json.xz")) + + +class TestStartTools(unittest.TestCase): + def setUp(self): + self.mod = import_engine_lib() + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir, ignore_errors=True) + self.orig_cwd = os.getcwd() + os.chdir(self.tmpdir) + self.addCleanup(os.chdir, self.orig_cwd) + + def test_argv_is_shlex_joined_before_invocation(self): + engine = self.mod.Engine() + engine.disable_tools = False + engine.tool_start_cmds = os.path.join(self.tmpdir, "start.json.xz") + engine.tool_stop_cmds = os.path.join(self.tmpdir, "stop.json.xz") + for f in (engine.tool_start_cmds, engine.tool_stop_cmds): + open(f, "a").close() + + argv = ["kerneltools-start", "--record-opts=-e cycles -c 1000"] + engine._parse_tool_commands = mock.MagicMock( + return_value=[("kernel-1", argv)] + ) + + with mock.patch.object(self.mod, "run_command") as mock_run: + engine.start_tools() + + mock_run.assert_called_once() + (cmd_str,), _ = mock_run.call_args + self.assertIn("'--record-opts=-e cycles -c 1000'", cmd_str) + self.assertTrue(cmd_str.startswith("cd tool-data/kernel-1 && ")) + + +class TestLoadBenchCmds(unittest.TestCase): + def setUp(self): + self.mod = import_engine_lib() + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir, ignore_errors=True) + + def test_reads_json_entries(self): + import lzma + path = os.path.join(self.tmpdir, "start.json.xz") + entries = [{"test": "1-1", "argv": ["fio", "--iodepth=4"]}] + with lzma.open(path, "wt") as fh: + json.dump(entries, fh) + + engine = self.mod.Engine() + result = engine._load_bench_cmds(path) + self.assertEqual(result, entries) + + def test_missing_file_returns_empty_list(self): + engine = self.mod.Engine() + result = engine._load_bench_cmds(os.path.join(self.tmpdir, "missing.json.xz")) + self.assertEqual(result, []) + + +class TestRunBenchCmd(unittest.TestCase): + def setUp(self): + self.mod = import_engine_lib() + + def test_shlex_joins_argv_before_running(self): + engine = self.mod.Engine() + engine.cs_type = "client" + engine.abort = False + engine.quit = False + with mock.patch.object(self.mod, "run_command") as mock_run: + mock_run.return_value.return_code = 0 + engine.run_bench_cmd("client", "client", ["fio", "--exec_prerun=echo starting test"]) + mock_run.assert_called_once_with("fio '--exec_prerun=echo starting test'") + + def test_empty_argv_is_a_no_op(self): + engine = self.mod.Engine() + engine.cs_type = "client" + engine.abort = False + engine.quit = False + with mock.patch.object(self.mod, "run_command") as mock_run: + rc = engine.run_bench_cmd("client", "client", []) + mock_run.assert_not_called() + self.assertEqual(rc, 0) + + def test_non_matching_type_is_a_no_op(self): + engine = self.mod.Engine() + engine.cs_type = "server" + with mock.patch.object(self.mod, "run_command") as mock_run: + rc = engine.run_bench_cmd("client", "client", ["fio"]) + mock_run.assert_not_called() + self.assertEqual(rc, 0) + + +class TestCliRunBenchCmd(unittest.TestCase): + def setUp(self): + self.mod = import_engine_lib() + + def test_trailing_args_collected_as_argv_and_joined(self): + with mock.patch.object(self.mod, "run_command") as mock_run: + mock_run.return_value.return_code = 0 + with self.assertRaises(SystemExit) as ctx: + self.mod.cli_run_bench_cmd( + "client", "client", "client", "False", "False", "0", + "fio", "--exec_prerun=echo starting test", + ) + mock_run.assert_called_once_with("fio '--exec_prerun=echo starting test'") + self.assertEqual(ctx.exception.code, 0) + + def test_no_trailing_args_is_a_no_op(self): + with mock.patch.object(self.mod, "run_command") as mock_run: + self.mod.cli_run_bench_cmd("client", "client", "client", "False", "False", "0") + mock_run.assert_not_called() + + def test_dispatch_splat_matches_sys_argv_shape(self): + # sys.argv[2:] arrives as separate literal elements (no shell, since + # process_bench_roadblocks() now passes wait_for as a real list to + # subprocess.Popen) -- confirm cli_run_bench_cmd(*args) still accepts + # an arbitrary number of trailing argv tokens via the same splat the + # dispatch table already uses + args = ["client", "client", "client", "False", "False", "0", + "fio", "--iodepth=4", "--rw=read"] + with mock.patch.object(self.mod, "run_command") as mock_run: + mock_run.return_value.return_code = 0 + with self.assertRaises(SystemExit): + self.mod.cli_run_bench_cmd(*args) + mock_run.assert_called_once_with("fio --iodepth=4 --rw=read") + + +class TestWaitForUsesArgvList(unittest.TestCase): + """The core regression test for the wait_for double-quoting hazard: + once start_argv can legitimately contain a shlex.quote()-protected, + single-quote-containing token, wait_for must be built as a literal list + handed straight to roadblock -- never re-rendered into a string that + would need re-splitting.""" + + def setUp(self): + self.mod = import_engine_lib() + self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir, ignore_errors=True) + self.orig_cwd = os.getcwd() + self.addCleanup(os.chdir, self.orig_cwd) + + def test_wait_for_is_a_list_containing_start_argv_verbatim(self): + engine = self.mod.Engine() + engine.cs_type = "client" + engine.cs_id = "1" + engine.cs_label = "client-1" + engine.cs_buddy = "server-1" + engine.cs_dir = self.tmpdir + engine.roadblock_msgs_dir = os.path.join(self.tmpdir, "rb-msgs") + os.makedirs(engine.roadblock_msgs_dir, exist_ok=True) + engine.default_timeout = 300 + engine.max_sample_failures = 3 + engine.abort = False + engine.quit = False + engine.endpoint_label = "" + engine._copy_files_to_sample_dir = lambda *_a, **_k: None + + # a value that needed shlex.quote()-protection -- exactly the shape + # that broke the old single-quote-wrapped wait_for string + start_argv = ["fio", "--exec_prerun=echo starting test"] + + bench_cmds = { + "bench-start-cmds.json.xz": [{"test": "1-1", "argv": start_argv}], + "bench-infra-cmds.json.xz": [{"test": "1-1", "argv": []}], + "bench-runtime-cmds.json.xz": [{"test": "1-1", "argv": ["probe"]}], + "bench-stop-cmds.json.xz": [{"test": "1-1", "argv": []}], + } + engine._load_bench_cmds = lambda filename: bench_cmds[filename] + + do_roadblock_calls = [] + + def fake_do_roadblock(label, timeout, messages=None, wait_for=None, do_abort=False): + do_roadblock_calls.append({"label": label, "wait_for": wait_for}) + return 0 + + engine.do_roadblock = fake_do_roadblock + + fake_runtime_result = mock.MagicMock(return_code=0, stdout="-1") + with mock.patch.object(self.mod, "run_command", return_value=fake_runtime_result): + engine.process_bench_roadblocks() + + wait_for_calls = [c for c in do_roadblock_calls if c["wait_for"] is not None] + self.assertEqual(len(wait_for_calls), 1) + wait_for_cmd = wait_for_calls[0]["wait_for"] + + self.assertIsInstance(wait_for_cmd, list) + self.assertNotIsInstance(wait_for_cmd, str) + self.assertEqual(wait_for_cmd[:6], [ + "python3", "/usr/local/bin/engine_lib.py", "run_bench_cmd", + "client", "client", "client", + ]) + # start_argv's tokens appear verbatim, with their embedded quote + # character intact -- never re-parsed through shlex.split() + self.assertEqual(wait_for_cmd[-len(start_argv):], start_argv) + + +if __name__ == "__main__": + unittest.main()