JIRA: BOTTLENECKS-29
[bottlenecks.git] / vstf / 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     def __init__(self, parent_image_path, child_image_dir):
23         """
24         :param parent_image_path    str: the parent image path.
25         :param child_image_dir      str: the destination path to put child images.
26         """
27         self._create_child_str = 'qemu-img create -f %(image_type)s %(child_path)s -o backing_file=%(parent_path)s'
28         self._convert_str = "qemu-img convert -O %(image_type)s %(parent_path)s %(child_path)s"
29         self.child_image_dir = child_image_dir
30         self.parent_image_path = parent_image_path
31         assert os.path.isfile(self.parent_image_path)
32         assert os.path.isdir(self.child_image_dir)
33
34     def create_child_image(self, child_name, full_clone=False, image_type='qcow2'):
35         """
36         create a child image and put it in self.child_image_dir.
37
38         :param child_name:  the image name to be created..
39         :return: return the path of child image.
40         """
41
42         image_path = os.path.join(self.child_image_dir, child_name) + '.' + image_type
43         if full_clone:
44             cmd = self._convert_str % {'image_type': image_type, 'child_path': image_path, 'parent_path': self.parent_image_path}
45         else:
46             cmd = self._create_child_str % {'child_path': image_path, 'parent_path': self.parent_image_path, 'image_type':image_type}
47         check_call(cmd.split())
48         return image_path
49
50
51 class ImageManager(object):
52     def __init__(self, cfg):
53         """
54         ImageManager creates images from configuration context.
55
56         :param cfg: dict, example:
57         {
58             'parent_image': "/mnt/sdb/ubuntu_salt_master.img",
59             'dst_location': "/mnt/sdb",
60             'full_clone':False,
61             'type': "qcow2",
62             'names': ['vm1','vm2','vm3','vm4']
63         }
64         :return:
65         """
66         super(ImageManager, self).__init__()
67         cfg = self._check_cfg(cfg)
68         self.parent_image = cfg['parent_image']
69         self.image_dir = cfg['dst_location']
70         self.full_clone = cfg['full_clone']
71         self.image_type = cfg['type']
72         self.names = cfg['names']
73         self.mgr = _ImageManager(self.parent_image, self.image_dir)
74
75     @staticmethod
76     def _check_cfg(cfg):
77         for key in ('parent_image', 'dst_location', 'full_clone', 'type', 'names'):
78             if key not in cfg:
79                 raise Exception("does't find %s config" % key)
80         if cfg['type'] not in ('raw', 'qcow2'):
81             raise Exception("type:%s not supported, only support 'raw' and 'qcow2'" % cfg['type'])
82         if not cfg['full_clone'] and cfg['type'] == 'raw':
83             raise Exception("only support 'qcow2' for not full_clone image creation" % cfg['type'])
84         return cfg
85
86     def create_all(self):
87         """
88         create images by configuration context.
89
90         :return: True for success, False for failure.
91         """
92         for name in self.names:
93             image = self.mgr.create_child_image(name, self.full_clone, self.image_type)
94             LOG.info("image: %s created", image)
95         return True
96
97     def clean_all(self):
98         """
99         remove all the images created in one go.
100
101         :return: True for success. Raise exception otherwise.
102         """
103         for name in self.names:
104             image_path = os.path.join(self.image_dir, name + '.' + self.image_type)
105             try:
106                 os.unlink(image_path)
107                 LOG.info("remove:%s successfully", image_path)
108             except Exception:
109                 LOG.info("cann't find path:%s", image_path)
110         return True
111
112
113 if __name__ == '__main__':
114     import argparse
115     import json
116     parser = argparse.ArgumentParser()
117     parser.add_argument('action', choices = ('create','clean'), help='action:create|clean')
118     parser.add_argument('--config', help='config file to parse')
119     args = parser.parse_args()
120     logging.basicConfig(level=logging.INFO)
121     image_cfg = json.load(open(args.config))
122     mgr = ImageManager(image_cfg)
123     if args.action == 'create':
124         mgr.create_all()
125     if args.action == 'clean':
126         mgr.clean_all()
127
128