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 Any, Optional, 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,33 +14,36 @@ 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
 
-    def param_get(self, param_name, key=None, accept_unknown_key=False):
+    def param_get(self, param_name: str, key: Optional[str] = None,
+                  accept_unknown_key: bool = False) -> Optional[str]:
         return _param_get(self.host, self.xe_prefix, self.uuid,
                           param_name, key, accept_unknown_key)
 
-    def param_set(self, param_name, value, key=None):
+    def param_set(self, param_name: str, value: Any, key: Optional[str] = None) -> None:
         _param_set(self.host, self.xe_prefix, self.uuid,
                    param_name, value, key)
 
-    def param_remove(self, param_name, key, accept_unknown_key=False):
+    def param_remove(self, param_name: str, key: str, accept_unknown_key=False) -> None:
         _param_remove(self.host, self.xe_prefix, self.uuid,
                       param_name, key, accept_unknown_key)
 
-    def param_add(self, param_name, value, key=None):
+    def param_add(self, param_name: str, value: str, key=None) -> None:
         _param_add(self.host, self.xe_prefix, self.uuid,
                    param_name, value, key)
 
-    def param_clear(self, param_name):
+    def param_clear(self, param_name: str) -> None:
         _param_clear(self.host, self.xe_prefix, self.uuid,
                      param_name)
 
-    def name(self):
-        return self.param_get('name-label')
+    def name(self) -> str:
+        n = self.param_get('name-label')
+        assert isinstance(n, str)
+        return n
 
     def vdi_uuids(self, sr_uuid=None):
         output = self._disk_list()
@@ -54,7 +61,7 @@ def vdi_uuids(self, sr_uuid=None):
                 vdis_on_sr.append(vdi)
         return vdis_on_sr
 
-    def destroy_vdi(self, vdi_uuid):
+    def destroy_vdi(self, vdi_uuid: str) -> None:
         self.host.xe('vdi-destroy', {'uuid': vdi_uuid})
 
     # FIXME: move this method and the above back to class VM if not useful in Snapshot class?
@@ -70,7 +77,7 @@ def all_vdis_on_host(self, host):
                 return False
         return True
 
-    def all_vdis_on_sr(self, sr):
+    def all_vdis_on_sr(self, sr) -> bool:
         for vdi_uuid in self.vdi_uuids():
             if self.host.pool.get_vdi_sr_uuid(vdi_uuid) != sr.uuid:
                 return False
@@ -84,7 +91,7 @@ def get_sr(self):
         assert sr.attached_to_host(self.host)
         return sr
 
-    def export(self, filepath, compress='none'):
+    def export(self, filepath, compress='none') -> None:
         logging.info("Export VM %s to %s with compress=%s" % (self.uuid, filepath, compress))
         params = {
             'uuid': self.uuid,

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,46 @@ 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_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 +196,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,9 @@ 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_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,27 +1,28 @@
 import logging
 import traceback
+from typing import Any, Dict, 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"
 
-    def __init__(self, master_hostname_or_ip):
+    def __init__(self, master_hostname_or_ip: str) -> None:
         master = Host(self, master_hostname_or_ip)
         assert master.is_master(), f"Host {master_hostname_or_ip} is not a master host. Aborting."
         self.master = master
         self.hosts = [master]
 
         # 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)
 
@@ -30,7 +31,7 @@ def __init__(self, master_hostname_or_ip):
                 host = Host(self, self.host_ip(host_uuid))
                 self.hosts.append(host)
         self.uuid = self.master.xe('pool-list', minimal=True)
-        self.saved_uefi_certs = None
+        self.saved_uefi_certs: Optional[Dict[str, Any]] = None
 
     def param_get(self, param_name, key=None, accept_unknown_key=False):
         return _param_get(self.master, Pool.xe_prefix, self.uuid, param_name, key, accept_unknown_key)
@@ -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)
@@ -117,7 +118,7 @@ def first_shared_sr(self):
     def get_vdi_sr_uuid(self, vdi_uuid):
         return self.master.xe('vdi-param-get', {'uuid': vdi_uuid, 'param-name': 'sr-uuid'})
 
-    def save_uefi_certs(self):
+    def save_uefi_certs(self) -> None:
         """
         Save UEFI certificates in order to restore them later. XCP-ng 8.2 only.
 

lib/vm.py

@@ -88,9 +88,9 @@ 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_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 +235,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 +255,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:

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: