diff --git a/lib/commands.py b/lib/commands.py index 8f806f11e..741bdacd7 100644 --- a/lib/commands.py +++ b/lib/commands.py @@ -9,7 +9,7 @@ import lib.config as config from lib.netutil import wrap_ip -from typing import TYPE_CHECKING, List, Literal, overload +from typing import TYPE_CHECKING, Generic, List, Literal, TypeVar, overload if TYPE_CHECKING: from lib.common import HostAddress @@ -39,19 +39,21 @@ def __init__(self, returncode: int, stdout: str, cmd: str | list[str]): f'Local command ({cmd}) failed with return code {returncode}{msg_end}' ) -class BaseCmdResult: +ResultOutputT = TypeVar('ResultOutputT', str, bytes) + +class BaseCmdResult(Generic[ResultOutputT]): __slots__ = 'returncode', 'stdout' - def __init__(self, returncode: int, stdout: str | bytes): + def __init__(self, returncode: int, stdout: ResultOutputT): self.returncode = returncode - self.stdout = stdout + self.stdout: ResultOutputT = stdout -class SSHResult(BaseCmdResult): - def __init__(self, returncode: int, stdout: str | bytes): +class SSHResult(BaseCmdResult[ResultOutputT]): + def __init__(self, returncode: int, stdout: ResultOutputT): super(SSHResult, self).__init__(returncode, stdout) -class LocalCommandResult(BaseCmdResult): - def __init__(self, returncode: int, stdout: str | bytes): +class LocalCommandResult(BaseCmdResult[ResultOutputT]): + def __init__(self, returncode: int, stdout: ResultOutputT): super(LocalCommandResult, self).__init__(returncode, stdout) def _ellide_log_lines(log: str) -> str: @@ -77,7 +79,7 @@ def _ssh( decode: bool, options: list[str], multiplexing: bool, -) -> SSHResult | SSHCommandFailed | str | bytes | None: +) -> SSHResult[str] | SSHResult[bytes] | SSHCommandFailed | str | bytes | None: opts = list(options) opts += ['-o', 'BatchMode yes'] opts += ['-o', 'PubkeyAcceptedKeyTypes +ssh-rsa'] @@ -137,22 +139,26 @@ def _ssh( if res.returncode == 255: return SSHCommandFailed(255, "SSH Error: %s" % output_for_errors, cmd) - output: str | bytes = res.stdout + output: bytes = res.stdout if banner_res: if banner_res.returncode == 255: return SSHCommandFailed(255, "SSH Error: %s" % banner_res.stdout.decode(errors='replace'), cmd) output = output[len(banner_res.stdout):] - if decode: - assert isinstance(output, bytes) - output = output.decode() - if res.returncode and check: return SSHCommandFailed(res.returncode, output_for_errors, cmd) - if simple_output: - return output.strip() - return SSHResult(res.returncode, output) + if decode: + output_str = output.decode() + if simple_output: + return output_str.strip() + else: + return SSHResult[str](res.returncode, output_str) + else: + if simple_output: + return output.strip() + else: + return SSHResult[bytes](res.returncode, output) # The actual code is in _ssh(). # This function is kept short for shorter pytest traces upon SSH failures, which are common, @@ -173,7 +179,13 @@ def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, simple_output: Literal[False], suppress_fingerprint_warnings: bool = True, background: Literal[False] = False, - decode: bool = True, options: List[str] = [], multiplexing: bool = True) -> SSHResult: + decode: Literal[True] = True, options: List[str] = [], multiplexing: bool = True) -> SSHResult[str]: + ... +@overload +def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, + simple_output: Literal[False], + suppress_fingerprint_warnings: bool = True, background: Literal[False] = False, + decode: Literal[False], options: List[str] = [], multiplexing: bool = True) -> SSHResult[bytes]: ... @overload def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, @@ -186,12 +198,12 @@ def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, simple_output: bool = True, suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True, options: List[str] = [], multiplexing: bool = True) \ - -> str | bytes | SSHResult | None: + -> str | bytes | SSHResult[str] | SSHResult[bytes] | None: ... def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, simple_output: bool = True, suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True, options: List[str] = [], multiplexing: bool = True) \ - -> str | bytes | SSHResult | None: + -> str | bytes | SSHResult[str] | SSHResult[bytes] | None: result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings, background, decode, options, multiplexing) if isinstance(result_or_exc, SSHCommandFailed): @@ -199,9 +211,21 @@ def ssh(hostname_or_ip: HostAddress, cmd: str, *, check: bool = True, simple_out else: return result_or_exc -def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, suppress_fingerprint_warnings: bool = True, +@overload +def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, *, decode: Literal[True] = True, + suppress_fingerprint_warnings: bool = True, + background: bool = False, options: List[str] = [], + multiplexing: bool = True) -> SSHResult[str]: + ... +@overload +def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, *, decode: Literal[False], + suppress_fingerprint_warnings: bool = True, + background: bool = False, options: List[str] = [], + multiplexing: bool = True) -> SSHResult[bytes]: + ... +def ssh_with_result(hostname_or_ip: HostAddress, cmd: str, *, suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True, options: List[str] = [], - multiplexing: bool = True) -> SSHResult: + multiplexing: bool = True) -> SSHResult[str] | SSHResult[bytes]: result_or_exc = _ssh(hostname_or_ip, cmd, False, False, suppress_fingerprint_warnings, background, decode, options, multiplexing) if isinstance(result_or_exc, SSHCommandFailed): @@ -267,7 +291,15 @@ def sftp( return res -def local_cmd(cmd: List[str], check: bool = True, decode: bool = True) -> LocalCommandResult: +@overload +def local_cmd(cmd: List[str], *, check: bool = True, decode: Literal[True] = True) -> LocalCommandResult[str]: + ... +@overload +def local_cmd(cmd: List[str], *, check: bool = True, decode: Literal[False]) -> LocalCommandResult[bytes]: + ... +def local_cmd( + cmd: List[str], *, check: bool = True, decode: bool = True +) -> LocalCommandResult[str] | LocalCommandResult[bytes]: """ Run a command locally on tester end. """ logging.debug("[local] %s", (cmd,)) res = subprocess.run( @@ -280,10 +312,6 @@ def local_cmd(cmd: List[str], check: bool = True, decode: bool = True) -> LocalC # get a decoded version of the output in any case, replacing potential errors output_for_logs = res.stdout.decode(errors='replace').strip() - output: str | bytes = res.stdout - if decode: - output = res.stdout.decode() - errorcode_msg = "" if res.returncode == 0 else " - Got error code: %s" % res.returncode command = " ".join(cmd) logging.debug(f"[local] {command}{errorcode_msg}{_ellide_log_lines(output_for_logs)}") @@ -291,7 +319,10 @@ def local_cmd(cmd: List[str], check: bool = True, decode: bool = True) -> LocalC if res.returncode and check: raise LocalCommandFailed(res.returncode, output_for_logs, command) - return LocalCommandResult(res.returncode, output) + if decode: + return LocalCommandResult[str](res.returncode, res.stdout.decode()) + else: + return LocalCommandResult[bytes](res.returncode, res.stdout) def encode_powershell_command(cmd: str) -> str: return base64.b64encode(cmd.encode("utf-16-le")).decode("ascii") diff --git a/lib/host.py b/lib/host.py index ced472049..e12e8ad9c 100644 --- a/lib/host.py +++ b/lib/host.py @@ -31,7 +31,7 @@ from lib.vm import VM from lib.xo import xo_cli, xo_object_exists -from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload +from typing import TYPE_CHECKING, Literal, TypedDict, overload if TYPE_CHECKING: from lib.pool import Pool @@ -107,7 +107,13 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[True] = Tr @overload def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False], suppress_fingerprint_warnings: bool = True, background: Literal[False] = False, - decode: bool = True, multiplexing: bool = True) -> commands.SSHResult: + decode: Literal[True] = True, multiplexing: bool = True) -> commands.SSHResult[str]: + ... + + @overload + def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False], + suppress_fingerprint_warnings: bool = True, background: Literal[False] = False, + decode: Literal[False], multiplexing: bool = True) -> commands.SSHResult[bytes]: ... @overload @@ -118,19 +124,18 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, @overload def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, - suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True, - multiplexing: bool = True) \ - -> str | bytes | commands.SSHResult | None: + suppress_fingerprint_warnings: bool = True, background: Literal[False] = False, + decode: Literal[True] = True, multiplexing: bool = True) -> str | commands.SSHResult[str]: ... def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, suppress_fingerprint_warnings: bool = True, background: bool = False, decode: bool = True, - multiplexing: bool = True) -> str | bytes | commands.SSHResult | None: + multiplexing: bool = True) -> str | bytes | commands.SSHResult[str] | commands.SSHResult[bytes] | None: return commands.ssh(self.hostname_or_ip, cmd, check=check, simple_output=simple_output, suppress_fingerprint_warnings=suppress_fingerprint_warnings, background=background, decode=decode, multiplexing=multiplexing) - def ssh_with_result(self, cmd: str) -> commands.SSHResult: + def ssh_with_result(self, cmd: str) -> commands.SSHResult[str]: # doesn't raise if the command's return is nonzero, unless there's a SSH error return commands.ssh_with_result(self.hostname_or_ip, cmd) @@ -148,12 +153,12 @@ def xe(self, action: str, args: dict[str, str | bool | dict[str, str]] = {}, *, @overload def xe(self, action: str, args: dict[str, str | bool | dict[str, str]] = {}, *, check: bool = ..., - simple_output: Literal[False], minimal: bool = ..., force: bool = ...) -> commands.SSHResult: + simple_output: Literal[False], minimal: bool = ..., force: bool = ...) -> commands.SSHResult[str]: ... def xe(self, action: str, args: dict[str, str | bool | dict[str, str]] = {}, *, check: bool = True, simple_output: bool = True, minimal: bool = False, force: bool = False) \ - -> str | commands.SSHResult: + -> str | commands.SSHResult[str]: maybe_param_minimal = '--minimal' if minimal else '' maybe_param_force = '--force' if force else '' @@ -169,14 +174,10 @@ def stringify(key: str, value: str | bool | dict[str, str]) -> str: command: str = f'xe {action} {maybe_param_minimal} {maybe_param_force} ' + \ ' '.join(stringify(key, value) for key, value in args.items()) - result = self.ssh( - command, - check=check, - simple_output=simple_output - ) - assert isinstance(result, (str, commands.SSHResult)) - - return result + if simple_output: + return self.ssh(command, check=check, simple_output=True) + else: + return self.ssh(command, check=check, simple_output=False) @overload def param_get(self, param_name: str, key: str | None = ..., @@ -239,11 +240,11 @@ def execute_script(self, script_contents: str, *, shebang: str = ..., simple_out @overload def execute_script( self, script_contents: str, *, shebang: str = ..., simple_output: Literal[False] - ) -> commands.SSHResult: + ) -> commands.SSHResult[str]: ... def execute_script(self, script_contents: str, shebang: str = 'sh', - simple_output: bool = True) -> str | commands.SSHResult: + simple_output: bool = True) -> str | commands.SSHResult[str]: with tempfile.NamedTemporaryFile('w') as script: os.chmod(script.name, 0o775) script.write('#!/usr/bin/env ' + shebang + '\n') @@ -259,7 +260,7 @@ def execute_script(self, script_contents: str, shebang: str = 'sh', try: logging.debug(f"[{self}] # Will execute this temporary script:\n{script_contents.strip()}") - return cast(str | commands.SSHResult, self.ssh(remote_path, simple_output=simple_output)) + return self.ssh(remote_path, simple_output=simple_output) finally: self.ssh(f'rm -f {remote_path}') diff --git a/lib/vm.py b/lib/vm.py index 0ebaedff4..b4627653d 100644 --- a/lib/vm.py +++ b/lib/vm.py @@ -27,7 +27,7 @@ from lib.vdi import VDI from lib.vif import VIF -from typing import TYPE_CHECKING, Iterable, List, Literal, cast, overload +from typing import TYPE_CHECKING, Iterable, List, Literal, overload if TYPE_CHECKING: from lib.host import Host @@ -116,7 +116,12 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[True] = Tr @overload def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False], - background: Literal[False] = False, decode: bool = True) -> commands.SSHResult: + background: Literal[False] = False, decode: Literal[True] = True) -> commands.SSHResult[str]: + ... + + @overload + def ssh(self, cmd: str, *, check: bool = True, simple_output: Literal[False], + background: Literal[False] = False, decode: Literal[False]) -> commands.SSHResult[bytes]: ... @overload @@ -125,18 +130,18 @@ def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, ... @overload - def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, background: bool = False, - decode: bool = True) -> str | bytes | commands.SSHResult | None: + def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, + background: Literal[False] = False, decode: Literal[True] = True) -> str | commands.SSHResult[str]: ... def ssh(self, cmd: str, *, check: bool = True, simple_output: bool = True, background: bool = False, - decode: bool = True) -> str | bytes | commands.SSHResult | None: + decode: bool = True) -> str | bytes | commands.SSHResult[str] | commands.SSHResult[bytes] | None: # raises by default for any nonzero return code assert self.ip is not None return commands.ssh(self.ip, cmd, check=check, simple_output=simple_output, background=background, decode=decode) - def ssh_with_result(self, cmd: str) -> commands.SSHResult: + def ssh_with_result(self, cmd: str) -> commands.SSHResult[str]: # doesn't raise if the command's return is nonzero, unless there's a SSH error assert self.ip is not None return commands.ssh_with_result(self.ip, cmd) @@ -446,10 +451,10 @@ def execute_script(self, script_contents: str, *, simple_output: Literal[True] = ... @overload - def execute_script(self, script_contents: str, *, simple_output: Literal[False]) -> commands.SSHResult: + def execute_script(self, script_contents: str, *, simple_output: Literal[False]) -> commands.SSHResult[str]: ... - def execute_script(self, script_contents: str, simple_output: bool = True) -> str | commands.SSHResult: + def execute_script(self, script_contents: str, simple_output: bool = True) -> str | commands.SSHResult[str]: with tempfile.NamedTemporaryFile('w') as f: f.write(script_contents) f.flush() @@ -458,7 +463,7 @@ def execute_script(self, script_contents: str, simple_output: bool = True) -> st logging.debug(f"[{self.ip}] # Will execute this temporary script:\n{script_contents.strip()}") # Use bash to run the script, to avoid being hit by differences between shells, for example on FreeBSD # It is a documented requirement that bash is present on all test VMs. - res = cast(str | commands.SSHResult, self.ssh(f'bash {f.name}', simple_output=simple_output)) + res = self.ssh(f'bash {f.name}', simple_output=simple_output) return res finally: self.ssh(f'rm -f {f.name}') @@ -777,26 +782,29 @@ def execute_powershell_script(self, script_contents: str, ... @overload - def execute_powershell_script(self, script_contents: str, - simple_output: Literal[False], - prepend: str = "$ProgressPreference = 'SilentlyContinue';") -> commands.SSHResult: + def execute_powershell_script( + self, + script_contents: str, + simple_output: Literal[False], + prepend: str = "$ProgressPreference = 'SilentlyContinue';", + ) -> commands.SSHResult[str]: ... def execute_powershell_script( self, script_contents: str, simple_output: bool = True, - prepend: str = "$ProgressPreference = 'SilentlyContinue';") -> str | commands.SSHResult: + prepend: str = "$ProgressPreference = 'SilentlyContinue';") -> str | commands.SSHResult[str]: # ProgressPreference is needed to suppress any clixml progress output, # as it's not filtered away from stdout by default, and we're grabbing stdout. assert self.is_windows if prepend is not None: script_contents = prepend + script_contents cmd = commands.encode_powershell_command(script_contents) - return cast(str | commands.SSHResult, self.ssh( + return self.ssh( f"powershell.exe -nologo -noprofile -noninteractive -encodedcommand {cmd}", simple_output=simple_output, - )) + ) def run_powershell_command(self, program: str, args: str) -> int: """ diff --git a/tests/misc/test_access_links.py b/tests/misc/test_access_links.py index daaa40f05..3f18d3098 100644 --- a/tests/misc/test_access_links.py +++ b/tests/misc/test_access_links.py @@ -37,7 +37,6 @@ def test_access_links(host: Host, command_id: str, url_id: str) -> None: # Verify the download worked by comparing with local download # This ensures the content is accessible and identical from both locations local_result = commands.local_cmd(COMMAND) - assert isinstance(local_result.stdout, str) assert local_result.returncode == 0, f"Failed to fetch URL locally: {local_result.stdout}" diff --git a/tests/misc/test_file_server.py b/tests/misc/test_file_server.py index 4d9fe02e6..4a1607d18 100644 --- a/tests/misc/test_file_server.py +++ b/tests/misc/test_file_server.py @@ -20,7 +20,6 @@ def test_fileserver_redirect_https(host: Host) -> None: path = "/path/to/dir/file.txt" ip = wrap_ip(host.hostname_or_ip) res = commands.local_cmd(["curl", "-s", "-i", "http://" + ip + path]) - assert isinstance(res.stdout, str) lines = res.stdout.splitlines() assert lines[0].strip() == "HTTP/1.1 301 Moved Permanently" assert _header_equal(lines[2], "location", "https://" + ip + path) @@ -35,7 +34,6 @@ def __get_header(host: Host) -> list[str]: res = commands.local_cmd( ["curl", "-s", "-XGET", "-k", "-I", "https://" + wrap_ip(host.hostname_or_ip)] ) - assert isinstance(res.stdout, str) return res.stdout.splitlines() def test_fileserver_hsts_default(self, host: Host) -> None: diff --git a/tests/uefi_sb/utils.py b/tests/uefi_sb/utils.py index d087c153d..110c89937 100644 --- a/tests/uefi_sb/utils.py +++ b/tests/uefi_sb/utils.py @@ -198,5 +198,4 @@ def check_vm_cert_md5sum(vm: VM, key: str, reference_file: str) -> None: ) assert res.returncode == 0, f"Cert {key} must be present" reference_md5 = get_md5sum_from_auth(reference_file) - assert isinstance(res.stdout, bytes) assert hashlib.md5(res.stdout).hexdigest() == reference_md5