diff --git a/psh/tools/psh.py b/psh/tools/psh.py index add264b7c..2a66f70c1 100644 --- a/psh/tools/psh.py +++ b/psh/tools/psh.py @@ -50,15 +50,15 @@ def _check_result(pexpect_proc, result): 'Please choose between susccess/fail/dont-check'))) -def assert_cmd(pexpect_proc, cmd, *, expected='', result='success', msg='', is_regex=False): - ''' Sends specified command and asserts that it's displayed correctly - with optional expected output and next prompt. Exit status is asserted depending on `result`''' +def assert_cmd(pexpect_proc, cmd, *, expected="", result="success", msg="", is_regex=False, timeout=30): + """Sends specified command and asserts that it's displayed correctly + with optional expected output and next prompt. Exit status is asserted depending on `result`""" pexpect_proc.sendline(cmd) cmd = re.escape(cmd) - exp_regex = '' + exp_regex = "" if is_regex: exp_regex = expected - elif expected != '': + elif expected != "": if not (isinstance(expected, tuple) or isinstance(expected, list)): expected = tuple((expected,)) for line in expected: @@ -67,9 +67,9 @@ def assert_cmd(pexpect_proc, cmd, *, expected='', result='success', msg='', is_r exp_regex = cmd + EOL + exp_regex + PROMPT exp_readable = _readable(exp_regex) - msg = f'Expected output regex was: \n---\n{exp_readable}\n---\n' + msg + msg = f"Expected output regex was: \n---\n{exp_readable}\n---\n" + msg - assert pexpect_proc.expect([exp_regex, pexpect.TIMEOUT, pexpect.EOF]) == 0, msg + assert pexpect_proc.expect([exp_regex, pexpect.TIMEOUT, pexpect.EOF], timeout=timeout) == 0, msg _check_result(pexpect_proc, result)