add some type hints

Signed-off-by: Gaëtan Lehmann <[email protected]>

Author: glehmann

lib/basevm.py

@@ -1,6 +1,10 @@
 import logging
 
+from typing import TYPE_CHECKING
+
 import lib.commands as commands
+if TYPE_CHECKING:
+    import lib.host
 
 from lib.common import _param_add, _param_clear, _param_get, _param_remove, _param_set
 from lib.sr import SR
@@ -10,7 +14,7 @@ class BaseVM:
 
     xe_prefix = "vm"
 
-    def __init__(self, uuid, host):
+    def __init__(self, uuid, host: 'lib.host.Host'):
         logging.info("New %s: %s", type(self).__name__, uuid)
         self.uuid = uuid
         self.host = host
@@ -36,7 +40,9 @@ def param_clear(self, param_name):
                      param_name)
 
     def name(self):
-        return self.param_get('name-label')
+        n = self.param_get('name-label')
+        assert isinstance(n, str)
+        return n
 
     def vdi_uuids(self, sr_uuid=None):
         output = self._disk_list()

lib/commands.py

@@ -2,11 +2,14 @@
 import logging
 import shlex
 import subprocess
+from typing import Union
 
-import lib.config as config
+from _pytest.fixtures import _teardown_yield_fixture
 
+import lib.config as config
 from lib.netutil import wrap_ip
 
+
 class BaseCommandFailed(Exception):
     __slots__ = 'returncode', 'stdout', 'cmd'
 
@@ -61,7 +64,7 @@ def _ellide_log_lines(log):
     return "\n{}".format("\n".join(reduced_message))
 
 def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
-         background, target_os, decode, options):
+         background, target_os, decode, options) -> Union[SSHResult, SSHCommandFailed, str, bytes, subprocess.Popen]:
     opts = list(options)
     opts.append('-o "BatchMode yes"')
     if suppress_fingerprint_warnings:
@@ -86,6 +89,7 @@ def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warning
 
     windows_background = background and target_os == "windows"
     # Fetch banner and remove it to avoid stdout/stderr pollution.
+    banner_res = None
     if config.ignore_ssh_banner and not windows_background:
         banner_res = subprocess.run(
             "ssh root@%s %s '%s'" % (hostname_or_ip, ' '.join(opts), '\n'),
@@ -103,9 +107,10 @@ def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warning
     )
     logging.debug(f"[{hostname_or_ip}] {command}")
     if windows_background:
-        return True, process
+        return process
 
     stdout = []
+    assert process.stdout is not None
     for line in iter(process.stdout.readline, b''):
         readable_line = line.decode(errors='replace').strip()
         stdout.append(line)
@@ -118,34 +123,56 @@ def _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warning
 
     # Even if check is False, we still raise in case of return code 255, which means a SSH error.
     if res.returncode == 255:
-        return False, SSHCommandFailed(255, "SSH Error: %s" % output_for_errors, command)
+        return SSHCommandFailed(255, "SSH Error: %s" % output_for_errors, command)
 
-    output = res.stdout
-    if config.ignore_ssh_banner:
+    output: Union[bytes, str] = res.stdout
+    if banner_res:
         if banner_res.returncode == 255:
-            return False, SSHCommandFailed(255, "SSH Error: %s" % banner_res.stdout.decode(errors='replace'), command)
+            return SSHCommandFailed(255, "SSH Error: %s" % banner_res.stdout.decode(errors='replace'), command)
         output = output[len(banner_res.stdout):]
 
     if decode:
+        assert isinstance(output, bytes)
         output = output.decode()
 
     if res.returncode and check:
-        return False, SSHCommandFailed(res.returncode, output_for_errors, command)
+        return SSHCommandFailed(res.returncode, output_for_errors, command)
 
     if simple_output:
-        return True, output.strip()
-    return True, SSHResult(res.returncode, output)
+        return output.strip()
+    return SSHResult(res.returncode, output)
 
 # The actual code is in _ssh().
 # This function is kept short for shorter pytest traces upon SSH failures, which are common,
 # as pytest prints the whole function definition that raised the SSHCommandFailed exception
-def ssh(hostname_or_ip, cmd, check=True, simple_output=True, suppress_fingerprint_warnings=True,
-        background=False, target_os='linux', decode=True, options=[]):
-    success, result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
-                                  background, target_os, decode, options)
-    if not success:
+def ssh(hostname_or_ip, cmd, check=True, simple_output=True, suppress_fingerprint_warnings=True, background=False,
+        target_os='linux', decode=True, options=[]) -> Union[SSHResult, str, bytes, subprocess.Popen]:
+    result_or_exc = _ssh(hostname_or_ip, cmd, check, simple_output, suppress_fingerprint_warnings,
+                         background, target_os, decode, options)
+    if isinstance(result_or_exc, SSHCommandFailed):
+        raise result_or_exc
+    else:
+        return result_or_exc
+
+def ssh_str(hostname_or_ip, cmd, check=True, suppress_fingerprint_warnings=True,
+            background=False, target_os='linux', options=[]) -> str:
+    result_or_exc = _ssh(hostname_or_ip, cmd, check, True, suppress_fingerprint_warnings,
+                         background, target_os, True, options)
+    if isinstance(result_or_exc, SSHCommandFailed):
+        raise result_or_exc
+    elif isinstance(result_or_exc, str):
+        return result_or_exc
+    assert False, "unexpected type"
+
+def ssh_with_result(hostname_or_ip, cmd, suppress_fingerprint_warnings=True,
+                    background=False, target_os='linux', decode=True, options=[]) -> SSHResult:
+    result_or_exc = _ssh(hostname_or_ip, cmd, False, False, suppress_fingerprint_warnings,
+                         background, target_os, decode, options)
+    if isinstance(result_or_exc, SSHCommandFailed):
         raise result_or_exc
-    return result_or_exc
+    elif isinstance(result_or_exc, SSHResult):
+        return result_or_exc
+    assert False, "unexpected type"
 
 def scp(hostname_or_ip, src, dest, check=True, suppress_fingerprint_warnings=True, local_dest=False):
     opts = '-o "BatchMode yes"'
@@ -179,6 +206,7 @@ def scp(hostname_or_ip, src, dest, check=True, suppress_fingerprint_warnings=Tru
     return res
 
 def sftp(hostname_or_ip, cmds, check=True, suppress_fingerprint_warnings=True):
+    opts = ''
     if suppress_fingerprint_warnings:
         # Suppress warnings and questions related to host key fingerprints
         # because on a test network IPs get reused, VMs are reinstalled, etc.

lib/host.py

@@ -6,9 +6,12 @@
 import uuid
 
 from packaging import version
+from typing import TYPE_CHECKING
 
 import lib.commands as commands
 import lib.pif as pif
+if TYPE_CHECKING:
+    import lib.pool
 
 from lib.common import _param_add, _param_clear, _param_get, _param_remove, _param_set
 from lib.common import safe_split, strip_suffix, to_xapi_bool, wait_for, wait_for_not
@@ -34,7 +37,7 @@ def host_data(hostname_or_ip):
 class Host:
     xe_prefix = "host"
 
-    def __init__(self, pool, hostname_or_ip):
+    def __init__(self, pool: 'lib.pool.Pool', hostname_or_ip):
         self.pool = pool
         self.hostname_or_ip = hostname_or_ip
         self.inventory = None
@@ -62,9 +65,13 @@ def ssh(self, cmd, check=True, simple_output=True, suppress_fingerprint_warnings
                             suppress_fingerprint_warnings=suppress_fingerprint_warnings, background=background,
                             decode=decode)
 
-    def ssh_with_result(self, cmd):
+    def ssh_str(self, cmd, check=True, background=False) -> str:
+        # raises by default for any nonzero return code
+        return commands.ssh_str(self.hostname_or_ip, cmd, check=check, background=background)
+
+    def ssh_with_result(self, cmd) -> commands.SSHResult:
         # doesn't raise if the command's return is nonzero, unless there's a SSH error
-        return self.ssh(cmd, check=False, simple_output=False)
+        return commands.ssh_with_result(self.hostname_or_ip, cmd)
 
     def scp(self, src, dest, check=True, suppress_fingerprint_warnings=True, local_dest=False):
         return commands.scp(

lib/pool.py

@@ -1,14 +1,15 @@
 import logging
 import traceback
+from typing import Optional
 
 from packaging import version
 
 import lib.commands as commands
-
-from lib.common import safe_split, wait_for, wait_for_not, _param_get, _param_set
+from lib.common import _param_get, _param_set, safe_split, wait_for, wait_for_not
 from lib.host import Host
 from lib.sr import SR
 
+
 class Pool:
     xe_prefix = "pool"
 
@@ -20,8 +21,8 @@ def __init__(self, master_hostname_or_ip):
 
         # wait for XAPI startup to be done, or we can get "Connection
         # refused (calling connect )" when calling self.hosts_uuids()
-        wait_for(lambda: commands.ssh(master_hostname_or_ip, ['xapi-wait-init-complete', '60'],
-                                      check=False, simple_output=False).returncode == 0,
+        wait_for(lambda: commands.ssh_with_result(master_hostname_or_ip,
+                                                  ['xapi-wait-init-complete', '60']).returncode == 0,
                  f"Wait for XAPI init to be complete on {master_hostname_or_ip}",
                  timeout_secs=30 * 60)
 
@@ -108,7 +109,7 @@ def first_host_that_isnt(self, host):
                 return h
         return None
 
-    def first_shared_sr(self):
+    def first_shared_sr(self) -> Optional[SR]:
         uuids = safe_split(self.master.xe('sr-list', {'shared': True, 'content-type': 'user'}, minimal=True))
         if len(uuids) > 0:
             return SR(uuids[0], self)
@@ -138,7 +139,7 @@ def save_uefi_certs(self):
         assert self.master.xcp_version < version.parse("8.3"), "this function should only be needed on XCP-ng 8.2"
         logging.info('Saving pool UEFI certificates')
 
-        if int(self.master.ssh(["secureboot-certs", "--version"]).split(".")[0]) < 1:
+        if int(self.master.ssh_str(["secureboot-certs", "--version"]).split(".")[0]) < 1:
             raise RuntimeError("The host must have secureboot-certs version >= 1.0.0")
 
         saved_certs = {

lib/vm.py

@@ -88,9 +88,15 @@ def ssh(self, cmd, check=True, simple_output=True, background=False, decode=True
         return commands.ssh(self.ip, cmd, check=check, simple_output=simple_output, background=background,
                             target_os=target_os, decode=decode)
 
-    def ssh_with_result(self, cmd):
+    def ssh_str(self, cmd, check=True, background=False) -> str:
+        # raises by default for any nonzero return code
+        target_os = "windows" if self.is_windows else "linux"
+        return commands.ssh_str(self.ip, cmd, check=check, background=background,
+                                target_os=target_os)
+
+    def ssh_with_result(self, cmd) -> commands.SSHResult:
         # doesn't raise if the command's return is nonzero, unless there's a SSH error
-        return self.ssh(cmd, check=False, simple_output=False)
+        return commands.ssh_with_result(self.ip, cmd)
 
     def scp(self, src, dest, check=True, suppress_fingerprint_warnings=True, local_dest=False):
         # Stop execution if scp() is used on Windows VMs as some OpenSSH releases for Windows don't
@@ -235,7 +241,7 @@ def snapshot(self, ignore_vdis=None):
             args['ignore-vdi-uuids'] = ','.join(ignore_vdis)
         return Snapshot(self.host.xe('vm-snapshot', args), self.host)
 
-    def checkpoint(self):
+    def checkpoint(self) -> Snapshot:
         logging.info("Checkpoint VM")
         return Snapshot(self.host.xe('vm-checkpoint', {'uuid': self.uuid,
                                                        'new-name-label': 'Checkpoint of %s' % self.uuid}),
@@ -255,7 +261,7 @@ def get_residence_host(self):
         host_uuid = self.param_get('resident-on')
         return self.host.pool.get_host_by_uuid(host_uuid)
 
-    def start_background_process(self, cmd):
+    def start_background_process(self, cmd) -> str:
         script = "/tmp/bg_process.sh"
         pidfile = "/tmp/bg_process.pid"
         with tempfile.NamedTemporaryFile('w') as f:
@@ -276,7 +282,7 @@ def start_background_process(self, cmd):
             self.ssh(['bash', script], background=True)
             wait_for(lambda: self.ssh_with_result(['test', '-f', pidfile]),
                      "wait for pid file %s to exist" % pidfile)
-            pid = self.ssh(['cat', pidfile])
+            pid = self.ssh_str(['cat', pidfile])
             self.ssh(['rm', '-f', script])
             self.ssh(['rm', '-f', pidfile])
             return pid

tests/misc/test_basic_without_ssh.py

@@ -2,6 +2,9 @@
 import pytest
 
 from lib.common import wait_for
+from lib.host import Host
+from lib.sr import SR
+from lib.vm import VM
 
 # These tests are basic tests meant to be run to check that a VM performs
 # well, without obvious issues.
@@ -19,14 +22,14 @@
 #   Note however that an existing VM will be left on a different SR after the tests.
 
 @pytest.fixture(scope='session')
-def existing_shared_sr(host):
+def existing_shared_sr(host: Host) -> SR:
     sr = host.pool.first_shared_sr()
     assert sr is not None, "A shared SR on the pool is required"
     return sr
 
 @pytest.mark.multi_vms # run them on a variety of VMs
 @pytest.mark.big_vm # and also on a really big VM ideally
-def test_vm_start_stop(imported_vm):
+def test_vm_start_stop(imported_vm: VM):
     vm = imported_vm
     # if VM already running, stop it
     if (vm.is_running()):
@@ -43,19 +46,19 @@ def test_vm_start_stop(imported_vm):
 @pytest.mark.big_vm # and also on a really big VM ideally
 @pytest.mark.usefixtures("started_vm")
 class TestBasicNoSSH:
-    def test_pause(self, imported_vm):
+    def test_pause(self, imported_vm: VM):
         vm = imported_vm
         vm.pause(verify=True)
         vm.unpause()
         vm.wait_for_os_booted()
 
-    def test_suspend(self, imported_vm):
+    def test_suspend(self, imported_vm: VM):
         vm = imported_vm
         vm.suspend(verify=True)
         vm.resume()
         vm.wait_for_os_booted()
 
-    def test_snapshot(self, imported_vm):
+    def test_snapshot(self, imported_vm: VM):
         vm = imported_vm
         snapshot = vm.snapshot()
         try:
@@ -65,7 +68,7 @@ def test_snapshot(self, imported_vm):
         finally:
             snapshot.destroy(verify=True)
 
-    def test_checkpoint(self, imported_vm):
+    def test_checkpoint(self, imported_vm: VM):
         vm = imported_vm
         snapshot = vm.checkpoint()
         try:
@@ -79,7 +82,7 @@ def test_checkpoint(self, imported_vm):
     # We want to test storage migration (memory+disks) and live migration without storage migration (memory only).
     # The order will depend on the initial location of the VM: a local SR or a shared SR.
     @pytest.mark.usefixtures("hostA2")
-    def test_live_migrate(self, imported_vm, existing_shared_sr):
+    def test_live_migrate(self, imported_vm: VM, existing_shared_sr: SR):
         def live_migrate(vm, dest_host, dest_sr, check_vdis=False):
             vm.migrate(dest_host, dest_sr)
             if check_vdis:

tests/misc/test_vm_basic_operations.py

@@ -2,29 +2,30 @@
 import pytest
 
 from lib.common import wait_for_not
+from lib.vm import VM
 
 @pytest.mark.multi_vms
 class Test:
-    def test_pause(self, running_vm):
+    def test_pause(self, running_vm: VM):
         vm = running_vm
         vm.pause(verify=True)
         vm.unpause()
         vm.wait_for_vm_running_and_ssh_up()
 
-    def test_suspend(self, running_vm):
+    def test_suspend(self, running_vm: VM):
         vm = running_vm
         vm.suspend(verify=True)
         vm.resume()
         vm.wait_for_vm_running_and_ssh_up()
 
-    def test_snapshot(self, running_vm):
+    def test_snapshot(self, running_vm: VM):
         vm = running_vm
         vm.test_snapshot_on_running_vm()
 
     # When using a windows VM the background ssh process is never terminated
     # This results in a ResourceWarning
     @pytest.mark.filterwarnings("ignore::ResourceWarning")
-    def test_checkpoint(self, running_vm):
+    def test_checkpoint(self, running_vm: VM):
         vm = running_vm
         logging.info("Start a 'sleep' process on VM through SSH")
         pid = vm.start_background_process('sleep 10000')