diff options
Diffstat (limited to 'tools/testing/selftests/net/lib/py/utils.py')
| -rw-r--r-- | tools/testing/selftests/net/lib/py/utils.py | 197 |
1 files changed, 163 insertions, 34 deletions
diff --git a/tools/testing/selftests/net/lib/py/utils.py b/tools/testing/selftests/net/lib/py/utils.py index 34470d65d871..6c44a3d2bbf7 100644 --- a/tools/testing/selftests/net/lib/py/utils.py +++ b/tools/testing/selftests/net/lib/py/utils.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: GPL-2.0 -import errno import json as _json import os -import random import re import select import socket @@ -11,9 +9,17 @@ import subprocess import time +class CmdInitFailure(Exception): + """ Command failed to start. Only raised by bkg(). """ + def __init__(self, msg, cmd_obj): + super().__init__(msg + "\n" + repr(cmd_obj)) + self.cmd = cmd_obj + + class CmdExitFailure(Exception): + """ Command failed (returned non-zero exit code). """ def __init__(self, msg, cmd_obj): - super().__init__(msg) + super().__init__(msg + "\n" + repr(cmd_obj)) self.cmd = cmd_obj @@ -21,18 +27,20 @@ def fd_read_timeout(fd, timeout): rlist, _, _ = select.select([fd], [], [], timeout) if rlist: return os.read(fd, 1024) - else: - raise TimeoutError("Timeout waiting for fd read") + raise TimeoutError("Timeout waiting for fd read") class cmd: """ Execute a command on local or remote host. + @shell defaults to false, and class will try to split @comm into a list + if it's a string with spaces. + Use bkg() instead to run a command in the background. """ - def __init__(self, comm, shell=True, fail=True, ns=None, background=False, - host=None, timeout=5, ksft_wait=None): + def __init__(self, comm, shell=None, fail=True, ns=None, background=False, + host=None, timeout=5, ksft_ready=None, ksft_wait=None): if ns: comm = f'ip netns exec {ns} ' + comm @@ -41,55 +49,93 @@ class cmd: self.ret = None self.ksft_term_fd = None + self.host = host self.comm = comm + if host: self.proc = host.cmd(comm) else: + # If user doesn't explicitly request shell try to avoid it. + if shell is None and isinstance(comm, str) and ' ' in comm: + comm = comm.split() + # ksft_wait lets us wait for the background process to fully start, # we pass an FD to the child process, and wait for it to write back. # Similarly term_fd tells child it's time to exit. - pass_fds = () + pass_fds = [] env = os.environ.copy() if ksft_wait is not None: - rfd, ready_fd = os.pipe() wait_fd, self.ksft_term_fd = os.pipe() - pass_fds = (ready_fd, wait_fd, ) - env["KSFT_READY_FD"] = str(ready_fd) + pass_fds.append(wait_fd) env["KSFT_WAIT_FD"] = str(wait_fd) + ksft_ready = True # ksft_wait implies ready + if ksft_ready is not None: + rfd, ready_fd = os.pipe() + pass_fds.append(ready_fd) + env["KSFT_READY_FD"] = str(ready_fd) self.proc = subprocess.Popen(comm, shell=shell, stdout=subprocess.PIPE, stderr=subprocess.PIPE, pass_fds=pass_fds, env=env) if ksft_wait is not None: - os.close(ready_fd) os.close(wait_fd) + if ksft_ready is not None: + os.close(ready_fd) msg = fd_read_timeout(rfd, ksft_wait) os.close(rfd) if not msg: - raise Exception("Did not receive ready message") + terminate = self.proc.poll() is None + self._process_terminate(terminate=terminate, timeout=1) + raise CmdInitFailure("Did not receive ready message", self) if not background: self.process(terminate=False, fail=fail, timeout=timeout) - def process(self, terminate=True, fail=None, timeout=5): - if fail is None: - fail = not terminate - - if self.ksft_term_fd: - os.write(self.ksft_term_fd, b"1") + def _process_terminate(self, terminate, timeout): if terminate: self.proc.terminate() - stdout, stderr = self.proc.communicate(timeout) + stdout, stderr = self.proc.communicate(timeout=timeout) self.stdout = stdout.decode("utf-8") self.stderr = stderr.decode("utf-8") self.proc.stdout.close() self.proc.stderr.close() self.ret = self.proc.returncode + return stdout, stderr + + def process(self, terminate=True, fail=None, timeout=5): + if fail is None: + fail = not terminate + + if self.ksft_term_fd: + os.write(self.ksft_term_fd, b"1") + + stdout, stderr = self._process_terminate(terminate=terminate, + timeout=timeout) if self.proc.returncode != 0 and fail: if len(stderr) > 0 and stderr[-1] == "\n": stderr = stderr[:-1] - raise CmdExitFailure("Command failed: %s\nSTDOUT: %s\nSTDERR: %s" % - (self.proc.args, stdout, stderr), self) + raise CmdExitFailure("Command failed", self) + + def __repr__(self): + def str_fmt(name, s): + name += ': ' + return (name + s.strip().replace('\n', '\n' + ' ' * len(name))) + + ret = "CMD" + if self.host: + ret += "[remote]" + if self.ret is None: + ret += f" (unterminated): {self.comm}\n" + elif self.ret == 0: + ret += f" (success): {self.comm}\n" + else: + ret += f": {self.comm}\n" + ret += f" EXIT: {self.ret}\n" + if self.stdout: + ret += str_fmt(" STDOUT", self.stdout) + "\n" + if self.stderr: + ret += str_fmt(" STDERR", self.stderr) + "\n" + return ret.strip() class bkg(cmd): @@ -111,12 +157,13 @@ class bkg(cmd): with bkg("my_binary", ksft_wait=5): """ - def __init__(self, comm, shell=True, fail=None, ns=None, host=None, - exit_wait=False, ksft_wait=None): + def __init__(self, comm, shell=None, fail=None, ns=None, host=None, + exit_wait=False, ksft_ready=None, ksft_wait=None): super().__init__(comm, background=True, shell=shell, fail=fail, ns=ns, host=host, - ksft_wait=ksft_wait) + ksft_ready=ksft_ready, ksft_wait=ksft_wait) self.terminate = not exit_wait and not ksft_wait + self._exit_wait = exit_wait self.check_fail = fail if shell and self.terminate: @@ -127,16 +174,20 @@ class bkg(cmd): return self def __exit__(self, ex_type, ex_value, ex_tb): - return self.process(terminate=self.terminate, fail=self.check_fail) + terminate = self.terminate + # Force termination on exception, but only if bkg() didn't already exit + # since forcing termination silences failures with fail=None + if self.proc.poll() is None: + terminate = terminate or (self._exit_wait and ex_type is not None) + return self.process(terminate=terminate, fail=self.check_fail) -global_defer_queue = [] +GLOBAL_DEFER_QUEUE = [] +GLOBAL_DEFER_ARMED = False class defer: def __init__(self, func, *args, **kwargs): - global global_defer_queue - if not callable(func): raise Exception("defer created with un-callable object, did you call the function instead of passing its name?") @@ -144,7 +195,9 @@ class defer: self.args = args self.kwargs = kwargs - self._queue = global_defer_queue + if not GLOBAL_DEFER_ARMED: + raise Exception("defer queue not armed, did you use defer() outside of a test case?") + self._queue = GLOBAL_DEFER_QUEUE self._queue.append(self) def __enter__(self): @@ -175,6 +228,10 @@ def tool(name, args, json=None, ns=None, host=None): return cmd_obj +def bpftool(args, json=None, ns=None, host=None): + return tool('bpftool', args, json=json, ns=ns, host=host) + + def ip(args, json=None, ns=None, host=None): if ns: args = f'-netns {ns} ' + args @@ -185,13 +242,67 @@ def ethtool(args, json=None, ns=None, host=None): return tool('ethtool', args, json=json, ns=ns, host=host) -def rand_port(type=socket.SOCK_STREAM): +def bpftrace(expr, json=None, ns=None, host=None, timeout=None): + """ + Run bpftrace and return map data (if json=True). + The output of bpftrace is inconvenient, so the helper converts + to a dict indexed by map name, e.g.: + { + "@": { ... }, + "@map2": { ... }, + } + """ + cmd_arr = ['bpftrace'] + # Throw in --quiet if json, otherwise the output has two objects + if json: + cmd_arr += ['-f', 'json', '-q'] + if timeout: + expr += ' interval:s:' + str(timeout) + ' { exit(); }' + timeout += 20 + cmd_arr += ['-e', expr] + cmd_obj = cmd(cmd_arr, ns=ns, host=host, shell=False, timeout=timeout) + if json: + # bpftrace prints objects as lines + ret = {} + for l in cmd_obj.stdout.split('\n'): + if not l.strip(): + continue + one = _json.loads(l) + if one.get('type') != 'map': + continue + for k, v in one["data"].items(): + if k.startswith('@'): + k = k.lstrip('@') + ret[k] = v + return ret + return cmd_obj + + +def rand_port(stype=socket.SOCK_STREAM): """ Get a random unprivileged port. """ - with socket.socket(socket.AF_INET6, type) as s: - s.bind(("", 0)) - return s.getsockname()[1] + return rand_ports(1, stype)[0] + + +def rand_ports(count, stype=socket.SOCK_STREAM): + """ + Get a unique set of random unprivileged ports. + """ + sockets = [] + ports = [] + + try: + for _ in range(count): + s = socket.socket(socket.AF_INET6, stype) + sockets.append(s) + s.bind(("", 0)) + ports.append(s.getsockname()[1]) + finally: + for s in sockets: + s.close() + + return ports def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadline=5): @@ -210,3 +321,21 @@ def wait_port_listen(port, proto="tcp", ns=None, host=None, sleep=0.005, deadlin if time.monotonic() > end: raise Exception("Waiting for port listen timed out") time.sleep(sleep) + + +def wait_file(fname, test_fn, sleep=0.005, deadline=5, encoding='utf-8'): + """ + Wait for file contents on the local system to satisfy a condition. + test_fn() should take one argument (file contents) and return whether + condition is met. + """ + end = time.monotonic() + deadline + + with open(fname, "r", encoding=encoding) as fp: + while True: + if test_fn(fp.read()): + break + fp.seek(0) + if time.monotonic() > end: + raise TimeoutError("Wait for file contents failed", fname) + time.sleep(sleep) |