1 files changed, 108 insertions, 0 deletions
diff --git a/tests/cloud_tests/platforms/azurecloud/image.py b/tests/cloud_tests/platforms/azurecloud/image.py
new file mode 100644
index 00000000..96a946f3
--- /dev/null
+++ b/tests/cloud_tests/platforms/azurecloud/image.py
@@ -0,0 +1,108 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+"""Azure Cloud image Base class."""
+
+from tests.cloud_tests import LOG
+
+from ..images import Image
+from .snapshot import AzureCloudSnapshot
+
+
+class AzureCloudImage(Image):
+    """Azure Cloud backed image."""
+
+    platform_name = 'azurecloud'
+
+    def __init__(self, platform, config, image_id):
+        """Set up image.
+
+        @param platform: platform object
+        @param config: image configuration
+        @param image_id: image id used to boot instance
+        """
+        super(AzureCloudImage, self).__init__(platform, config)
+        self.image_id = image_id
+        self._img_instance = None
+
+    @property
+    def _instance(self):
+        """Internal use only, returns a running instance"""
+        LOG.debug('creating instance')
+        if not self._img_instance:
+            self._img_instance = self.platform.create_instance(
+                self.properties, self.config, self.features,
+                self.image_id, user_data=None)
+        return self._img_instance
+
+    def destroy(self):
+        """Delete the instance used to create a custom image."""
+        LOG.debug('deleting VM that was used to create image')
+        if self._img_instance:
+            LOG.debug('Deleting instance %s', self._img_instance.name)
+            delete_vm = self.platform.compute_client.virtual_machines.delete(
+                self.platform.resource_group.name, self.image_id)
+            delete_vm.wait()
+
+        super(AzureCloudImage, self).destroy()
+
+    def _execute(self, *args, **kwargs):
+        """Execute command in image, modifying image."""
+        LOG.debug('executing commands on image')
+        self._instance.start()
+        return self._instance._execute(*args, **kwargs)
+
+    def push_file(self, local_path, remote_path):
+        """Copy file at 'local_path' to instance at 'remote_path'."""
+        LOG.debug('pushing file to image')
+        return self._instance.push_file(local_path, remote_path)
+
+    def run_script(self, *args, **kwargs):
+        """Run script in image, modifying image.
+
+        @return_value: script output
+        """
+        LOG.debug('running script on image')
+        self._instance.start()
+        return self._instance.run_script(*args, **kwargs)
+
+    def snapshot(self):
+        """ Create snapshot (image) of instance, wait until done.
+
+        If no instance has been booted, base image is returned.
+        Otherwise runs the clean script, deallocates, generalizes
+        and creates custom image from instance.
+        """
+        LOG.debug('creating image from VM')
+        if not self._img_instance:
+            return AzureCloudSnapshot(self.platform, self.properties,
+                                      self.config, self.features,
+                                      self.image_id, delete_on_destroy=False)
+
+        if self.config.get('boot_clean_script'):
+            self._img_instance.run_script(self.config.get('boot_clean_script'))
+
+        deallocate = self.platform.compute_client.virtual_machines.deallocate(
+            self.platform.resource_group.name, self.image_id)
+        deallocate.wait()
+
+        self.platform.compute_client.virtual_machines.generalize(
+            self.platform.resource_group.name, self.image_id)
+
+        image_params = {
+            "location": self.platform.location,
+            "properties": {
+                "sourceVirtualMachine": {
+                    "id": self._img_instance.instance.id
+                }
+            }
+        }
+        self.platform.compute_client.images.create_or_update(
+            self.platform.resource_group.name, self.image_id,
+            image_params)
+
+        self.destroy()
+
+        return AzureCloudSnapshot(self.platform, self.properties, self.config,
+                                  self.features, self.image_id)
+
+# vi: ts=4 expandtab