testsuites/vstf/vstf_scripts/vstf/agent/env/basic/image_manager.py

   1 ##############################################################################
   2 # Copyright (c) 2015 Huawei Technologies Co.,Ltd and others.
   3 #
   4 # All rights reserved. This program and the accompanying materials
   5 # are made available under the terms of the Apache License, Version 2.0
   6 # which accompanies this distribution, and is available at
   7 # http://www.apache.org/licenses/LICENSE-2.0
   8 ##############################################################################
   9
  10 from vstf.common.utils import check_call
  11 import os
  12 import logging
  13
  14 LOG = logging.getLogger(__name__)
  15
  16
  17 class _ImageManager(object):
  18     """
  19     A qemu-img wrapper to create qcow2 child image from a parent image.
  20
  21     """
  22
  23     def __init__(self, parent_image_path, child_image_dir):
  24         """
  25         :param parent_image_path    str: the parent image path.
  26         :param child_image_dir      str: the destination path to put child images.
  27         """
  28         self._create_child_str = 'qemu-img create -f %(image_type)s %(child_path)s -o backing_file=%(parent_path)s'
  29         self._convert_str = "qemu-img convert -O %(image_type)s %(parent_path)s %(child_path)s"
  30         self.child_image_dir = child_image_dir
  31         self.parent_image_path = parent_image_path
  32         assert os.path.isfile(self.parent_image_path)
  33         assert os.path.isdir(self.child_image_dir)
  34
  35     def create_child_image(
  36             self,
  37             child_name,
  38             full_clone=False,
  39             image_type='qcow2'):
  40         """
  41         create a child image and put it in self.child_image_dir.
  42
  43         :param child_name:  the image name to be created..
  44         :return: return the path of child image.
  45         """
  46
  47         image_path = os.path.join(
  48             self.child_image_dir,
  49             child_name) + '.' + image_type
  50         if full_clone:
  51             cmd = self._convert_str % {
  52                 'image_type': image_type,
  53                 'child_path': image_path,
  54                 'parent_path': self.parent_image_path}
  55         else:
  56             cmd = self._create_child_str % {
  57                 'child_path': image_path,
  58                 'parent_path': self.parent_image_path,
  59                 'image_type': image_type}
  60         check_call(cmd.split())
  61         return image_path
  62
  63
  64 class ImageManager(object):
  65
  66     def __init__(self, cfg):
  67         """
  68         ImageManager creates images from configuration context.
  69
  70         :param cfg: dict, example:
  71         {
  72             'parent_image': "/mnt/sdb/ubuntu_salt_master.img",
  73             'dst_location': "/mnt/sdb",
  74             'full_clone':False,
  75             'type': "qcow2",
  76             'names': ['vm1','vm2','vm3','vm4']
  77         }
  78         :return:
  79         """
  80         super(ImageManager, self).__init__()
  81         cfg = self._check_cfg(cfg)
  82         self.parent_image = cfg['parent_image']
  83         self.image_dir = cfg['dst_location']
  84         self.full_clone = cfg['full_clone']
  85         self.image_type = cfg['type']
  86         self.names = cfg['names']
  87         self.mgr = _ImageManager(self.parent_image, self.image_dir)
  88
  89     @staticmethod
  90     def _check_cfg(cfg):
  91         for key in (
  92             'parent_image',
  93             'dst_location',
  94             'full_clone',
  95             'type',
  96                 'names'):
  97             if key not in cfg:
  98                 raise Exception("does't find %s config" % key)
  99         if cfg['type'] not in ('raw', 'qcow2'):
 100             raise Exception(
 101                 "type:%s not supported, only support 'raw' and 'qcow2'" %
 102                 cfg['type'])
 103         if not cfg['full_clone'] and cfg['type'] == 'raw':
 104             raise Exception(
 105                 "only support 'qcow2' for not full_clone image creation" %
 106                 cfg['type'])
 107         return cfg
 108
 109     def create_all(self):
 110         """
 111         create images by configuration context.
 112
 113         :return: True for success, False for failure.
 114         """
 115         for name in self.names:
 116             image = self.mgr.create_child_image(
 117                 name, self.full_clone, self.image_type)
 118             LOG.info("image: %s created", image)
 119         return True
 120
 121     def clean_all(self):
 122         """
 123         remove all the images created in one go.
 124
 125         :return: True for success. Raise exception otherwise.
 126         """
 127         for name in self.names:
 128             image_path = os.path.join(
 129                 self.image_dir, name + '.' + self.image_type)
 130             try:
 131                 os.unlink(image_path)
 132                 LOG.info("remove:%s successfully", image_path)
 133             except Exception:
 134                 LOG.info("cann't find path:%s", image_path)
 135         return True
 136
 137
 138 if __name__ == '__main__':
 139     import argparse
 140     import json
 141     parser = argparse.ArgumentParser()
 142     parser.add_argument(
 143         'action',
 144         choices=(
 145             'create',
 146             'clean'),
 147         help='action:create|clean')
 148     parser.add_argument('--config', help='config file to parse')
 149     args = parser.parse_args()
 150     logging.basicConfig(level=logging.INFO)
 151     image_cfg = json.load(open(args.config))
 152     mgr = ImageManager(image_cfg)
 153     if args.action == 'create':
 154         mgr.create_all()
 155     if args.action == 'clean':
 156         mgr.clean_all()