lib/vm.py: add vdis related function

Add a list of vdis object to a VM
Add get_dom0_vm in host.py
Add function to connect/disconnect a VDI from a VM
Add a path to `Snapshot.revert()` to reset the VDIs list of the VM
object that created it since a revert would change the VDI UUIDs of the
VM.

Signed-off-by: Damien Thenot <[email protected]>

Author: Nambrok

lib/basevm.py

@@ -1,6 +1,6 @@
 import logging
 
-from typing import TYPE_CHECKING, Any, Literal, Optional, overload
+from typing import TYPE_CHECKING, Any, List, Literal, Optional, overload
 
 if TYPE_CHECKING:
     import lib.host
@@ -59,7 +59,7 @@ def name(self) -> str:
     def _disk_list(self):
         raise NotImplementedError()
 
-    def vdi_uuids(self, sr_uuid=None):
+    def vdi_uuids(self, sr_uuid: Optional[str] = None) -> List[str]:
         output = self._disk_list()
         if output == '':
             return []

lib/host.py

@@ -67,6 +67,7 @@ def __init__(self, pool: 'lib.pool.Pool', hostname_or_ip):
         self.uuid = self.inventory['INSTALLATION_UUID']
         self.xcp_version = version.parse(self.inventory['PRODUCT_VERSION'])
         self.xcp_version_short = f"{self.xcp_version.major}.{self.xcp_version.minor}"
+        self._dom0: Optional[VM] = None
 
     def __str__(self):
         return self.hostname_or_ip
@@ -711,3 +712,20 @@ def enable_hsts_header(self):
     def disable_hsts_header(self):
         self.ssh(['rm', '-f', f'{XAPI_CONF_DIR}/00-XCP-ng-tests-enable-hsts-header.conf'])
         self.restart_toolstack(verify=True)
+
+    def get_dom0_uuid(self):
+        return self.inventory["CONTROL_DOMAIN_UUID"]
+
+    def get_dom0_vm(self) -> VM:
+        if not self._dom0:
+            self._dom0 = VM(self.get_dom0_uuid(), self)
+        return self._dom0
+
+    def get_sr_from_vdi_uuid(self, vdi_uuid: str) -> Optional[SR]:
+        sr_uuid = self.xe("vdi-param-get", {
+            "param-name": "sr-uuid",
+            "uuid": vdi_uuid,
+        })
+        if not sr_uuid:
+            return None
+        return SR(sr_uuid, self.pool)

lib/snapshot.py

@@ -2,7 +2,18 @@
 
 from lib.basevm import BaseVM
 
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from lib.host import Host
+    from lib.vm import VM
+
+
 class Snapshot(BaseVM):
+    def __init__(self, uuid: str, host: "Host", vm: "VM"):
+        self.basevm = vm
+        super(Snapshot, self).__init__(uuid, host)
+
     def _disk_list(self):
         return self.host.xe('snapshot-disk-list', {'uuid': self.uuid, 'vbd-params': ''},
                             minimal=True)
@@ -21,3 +32,4 @@ def exists(self):
     def revert(self):
         logging.info("Revert to snapshot %s", self.uuid)
         self.host.xe('snapshot-revert', {'uuid': self.uuid})
+        self.basevm.create_vdis_list() # We reset the base VM object VDIs list because it changed following the revert

lib/sr.py

@@ -1,11 +1,18 @@
 import logging
 import time
-from typing import Optional
 
 import lib.commands as commands
-from lib.common import prefix_object_name, safe_split, strtobool, wait_for, wait_for_not
+from lib.common import (
+    prefix_object_name,
+    safe_split,
+    strtobool,
+    wait_for,
+    wait_for_not,
+)
 from lib.vdi import VDI
 
+from typing import Optional
+
 class SR:
     def __init__(self, uuid, pool):
         self.uuid = uuid

lib/vdi.py

@@ -25,11 +25,7 @@ def __init__(self, uuid, *, host=None, sr=None):
         # TODO: use a different approach when migration is possible
         if sr is None:
             assert host
-            sr_uuid = host.pool.get_vdi_sr_uuid(uuid)
-            # avoid circular import
-            # FIXME should get it from Host instead
-            from lib.sr import SR
-            self.sr = SR(sr_uuid, host.pool)
+            self.sr = host.get_sr_from_vdi_uuid(self.uuid)
         else:
             self.sr = sr
 

lib/vm.py

@@ -19,6 +19,7 @@
 )
 from lib.snapshot import Snapshot
 from lib.vbd import VBD
+from lib.vdi import VDI
 from lib.vif import VIF
 
 from typing import TYPE_CHECKING, List, Literal, Optional, Union, overload
@@ -33,6 +34,7 @@ def __init__(self, uuid: str, host: 'Host'):
         self.previous_host = None # previous host when migrated or being migrated
         self.is_windows = self.param_get('platform', 'device_id', accept_unknown_key=True) == '0002'
         self.is_uefi = self.param_get('HVM-boot-params', 'firmware', accept_unknown_key=True) == 'uefi'
+        self.create_vdis_list()
 
     def power_state(self) -> str:
         return self.param_get('power-state')
@@ -273,19 +275,74 @@ def migrate(self, target_host, sr=None, network=None):
 
         self.previous_host = self.host
         self.host = target_host
+        self.create_vdis_list()
 
     def snapshot(self, ignore_vdis=None):
         logging.info("Snapshot VM")
         args = {'uuid': self.uuid, 'new-name-label': 'Snapshot of %s' % self.uuid}
         if ignore_vdis:
             args['ignore-vdi-uuids'] = ','.join(ignore_vdis)
-        return Snapshot(self.host.xe('vm-snapshot', args), self.host)
+        snap_uuid = self.host.xe('vm-snapshot', args)
+        return Snapshot(snap_uuid, self.host, 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}),
-                        self.host)
+                        self.host, self)
+
+    def connect_vdi(self, vdi: VDI, device: str = "autodetect") -> str:
+        logging.info(f">> Plugging VDI {vdi.uuid} on VM {self.uuid}")
+        vbd_uuid = self.host.xe("vbd-create", {
+            "vdi-uuid": vdi.uuid,
+            "vm-uuid": self.uuid,
+            "device": device,
+        })
+        try:
+            self.host.xe("vbd-plug", {"uuid": vbd_uuid})
+        except commands.SSHCommandFailed:
+            self.host.xe("vbd-destroy", {"uuid": vbd_uuid})
+            raise
+
+        self.vdis.append(vdi)
+
+        return vbd_uuid
+
+    def disconnect_vdi(self, vdi: VDI):
+        logging.info(f"<< Unplugging VDI {vdi.uuid} from VM {self.uuid}")
+        assert vdi in self.vdis, "VDI {vdi.uuid} not in VM {self.uuid} VDI list"
+        vbd_uuid = self.host.xe("vbd-list", {
+            "vdi-uuid": vdi.uuid,
+            "vm-uuid": self.uuid
+        }, minimal=True)
+        try:
+            self.host.xe("vbd-unplug", {"uuid": vbd_uuid})
+        except commands.SSHCommandFailed as e:
+            if e.stdout == f"The device is not currently attached\ndevice: {vbd_uuid}":
+                logging.info(f"VBD {vbd_uuid} already unplugged")
+            else:
+                raise
+        self.host.xe("vbd-destroy", {"uuid": vbd_uuid})
+        self.vdis.remove(vdi)
+
+    def destroy_vdi(self, vdi_uuid: str) -> None:
+        for vdi in self.vdis:
+            if vdi.uuid == vdi_uuid:
+                self.vdis.remove(vdi)
+                super().destroy_vdi(vdi_uuid)
+                break
+
+    def create_vdis_list(self) -> None:
+        """ Used to redo the VDIs list of the VM when reverting a snapshot. """
+        try:
+            self.vdis = [VDI(vdi_uuid, host=self.host) for vdi_uuid in self.vdi_uuids()]
+        except commands.SSHCommandFailed as e:
+            # Doesn't work with Dom0 since `vm-disk-list` doesn't work on it so we create empty list
+            if e.stdout == "Error: No matching VMs found":
+                logging.info("Couldn't get disks list. We are Dom0. Continuing...")
+                self.vdis = []
+            else:
+                raise
 
     def vifs(self):
         _vifs = []