Merge pull request #1308 from caberos/issue1295

vs upgrade disk and add new disk

Author: allmightyspiff

SoftLayer/CLI/virt/upgrade.py

@@ -20,16 +20,20 @@
               help="CPU core will be on a dedicated host server.")
 @click.option('--memory', type=virt.MEM_TYPE, help="Memory in megabytes")
 @click.option('--network', type=click.INT, help="Network port speed in Mbps")
+@click.option('--add-disk', type=click.INT, multiple=True, required=False, help="add Hard disk in GB")
+@click.option('--resize-disk', nargs=2, multiple=True, type=(int, int),
+              help="Update disk number to size in GB. --resize-disk 250 2 ")
 @click.option('--flavor', type=click.STRING,
               help="Flavor keyName\nDo not use --memory, --cpu or --private, if you are using flavors")
 @environment.pass_env
-def cli(env, identifier, cpu, private, memory, network, flavor):
+def cli(env, identifier, cpu, private, memory, network, flavor, add_disk, resize_disk):
     """Upgrade a virtual server."""
 
     vsi = SoftLayer.VSManager(env.client)
 
-    if not any([cpu, memory, network, flavor]):
-        raise exceptions.ArgumentError("Must provide [--cpu], [--memory], [--network], or [--flavor] to upgrade")
+    if not any([cpu, memory, network, flavor, resize_disk, add_disk]):
+        raise exceptions.ArgumentError("Must provide [--cpu],"
+                                       " [--memory], [--network], [--flavor], [--resize-disk], or [--add] to upgrade")
 
     if private and not cpu:
         raise exceptions.ArgumentError("Must specify [--cpu] when using [--private]")
@@ -38,8 +42,19 @@ def cli(env, identifier, cpu, private, memory, network, flavor):
     if not (env.skip_confirmations or formatting.confirm("This action will incur charges on your account. Continue?")):
         raise exceptions.CLIAbort('Aborted')
 
+    disk_json = list()
     if memory:
         memory = int(memory / 1024)
-
-    if not vsi.upgrade(vs_id, cpus=cpu, memory=memory, nic_speed=network, public=not private, preset=flavor):
+    if resize_disk:
+        for guest_disk in resize_disk:
+            disks = {'capacity': guest_disk[0], 'number': guest_disk[1]}
+            disk_json.append(disks)
+
+    elif add_disk:
+        for guest_disk in add_disk:
+            disks = {'capacity': guest_disk, 'number': -1}
+            disk_json.append(disks)
+
+    if not vsi.upgrade(vs_id, cpus=cpu, memory=memory, nic_speed=network, public=not private, preset=flavor,
+                       disk=disk_json):
         raise exceptions.CLIAbort('VS Upgrade Failed')

SoftLayer/fixtures/SoftLayer_Virtual_Guest.py

@@ -8,11 +8,16 @@
         'id': 6327,
         'nextInvoiceTotalRecurringAmount': 1.54,
         'children': [
-            {'nextInvoiceTotalRecurringAmount': 1},
-            {'nextInvoiceTotalRecurringAmount': 1},
-            {'nextInvoiceTotalRecurringAmount': 1},
-            {'nextInvoiceTotalRecurringAmount': 1},
-            {'nextInvoiceTotalRecurringAmount': 1},
+            {'categoryCode': 'port_speed',
+             'nextInvoiceTotalRecurringAmount': 1},
+            {'categoryCode': 'guest_core',
+             'nextInvoiceTotalRecurringAmount': 1},
+            {'categoryCode': 'ram',
+             'nextInvoiceTotalRecurringAmount': 1},
+            {'categoryCode': 'guest_core',
+             'nextInvoiceTotalRecurringAmount': 1},
+            {'categoryCode': 'guest_disk1',
+             'nextInvoiceTotalRecurringAmount': 1},
         ],
         'package': {
             "id": 835,
@@ -622,7 +627,35 @@
             'capacity': '2',
             'description': 'RAM',
         }
-    },
+    }, {
+        "id": 2255,
+        "categories": [
+            {
+                "categoryCode": "guest_disk1",
+                "id": 82,
+                "name": "Second Disk"
+            },
+            {
+                "categoryCode": "guest_disk2",
+                "id": 92,
+                "name": "Third Disk"
+            },
+            {
+                "categoryCode": "guest_disk3",
+                "id": 93,
+                "name": "Fourth Disk"
+            },
+            {
+                "categoryCode": "guest_disk4",
+                "id": 116,
+                "name": "Fifth Disk"
+            }
+        ],
+        "item": {
+            "capacity": "10",
+            "description": "10 GB (SAN)"
+        }
+    }
 ]
 
 DEDICATED_GET_UPGRADE_ITEM_PRICES = [
@@ -641,7 +674,6 @@
 
 getMetricTrackingObjectId = 1000
 
-
 getBandwidthAllotmentDetail = {
     'allocationId': 25465663,
     'bandwidthAllotmentId': 138442,

SoftLayer/managers/vs.py

@@ -18,6 +18,7 @@
 
 LOGGER = logging.getLogger(__name__)
 
+
 # pylint: disable=no-self-use,too-many-lines
 
 
@@ -818,7 +819,7 @@ def capture(self, instance_id, name, additional_disks=False, notes=None):
         return self.guest.createArchiveTransaction(
             name, disks_to_capture, notes, id=instance_id)
 
-    def upgrade(self, instance_id, cpus=None, memory=None, nic_speed=None, public=True, preset=None):
+    def upgrade(self, instance_id, cpus=None, memory=None, nic_speed=None, public=True, preset=None, disk=None):
         """Upgrades a VS instance.
 
         Example::
@@ -862,6 +863,40 @@ def upgrade(self, instance_id, cpus=None, memory=None, nic_speed=None, public=Tr
             'virtualGuests': [{'id': int(instance_id)}],
         }
 
+        if disk:
+            disk_number = 0
+            vsi_disk = self.get_instance(instance_id)
+            for item in vsi_disk.get('billingItem').get('children'):
+                if item.get('categoryCode').__contains__('guest_disk'):
+                    if disk_number < int("".join(filter(str.isdigit, item.get('categoryCode')))):
+                        disk_number = int("".join(filter(str.isdigit, item.get('categoryCode'))))
+            for disk_guest in disk:
+                if disk_guest.get('number') > 0:
+                    price_id = self._get_price_id_for_upgrade_option(upgrade_prices, 'disk',
+                                                                     disk_guest.get('capacity'),
+                                                                     public)
+                    disk_number = disk_guest.get('number')
+
+                else:
+                    price_id = self._get_price_id_for_upgrade_option(upgrade_prices, 'disk',
+                                                                     disk_guest.get('capacity'),
+                                                                     public)
+                    disk_number = disk_number + 1
+
+                if price_id is None:
+                    raise exceptions.SoftLayerAPIError(500,
+                                                       'Unable to find %s option with value %s' % (
+                                                           ('disk', disk_guest.get('capacity'))))
+
+                category = {'categories': [{
+                    'categoryCode': 'guest_disk' + str(disk_number),
+                    'complexType': "SoftLayer_Product_Item_Category"}],
+                            'complexType': 'SoftLayer_Product_Item_Price',
+                            'id': price_id}
+
+                prices.append(category)
+            order['prices'] = prices
+
         for option, value in data.items():
             if not value:
                 continue
@@ -875,8 +910,8 @@ def upgrade(self, instance_id, cpus=None, memory=None, nic_speed=None, public=Tr
                     "Unable to find %s option with value %s" % (option, value))
 
             prices.append({'id': price_id})
-        order['prices'] = prices
 
+            order['prices'] = prices
         if preset is not None:
             vs_object = self.get_instance(instance_id)['billingItem']['package']
             order['presetId'] = self.ordering_manager.get_preset_by_key(vs_object['keyName'], preset)['id']
@@ -994,9 +1029,11 @@ def _get_price_id_for_upgrade_option(self, upgrade_prices, option, value, public
         option_category = {
             'memory': 'ram',
             'cpus': 'guest_core',
-            'nic_speed': 'port_speed'
+            'nic_speed': 'port_speed',
+            'disk': 'guest_disk'
         }
         category_code = option_category.get(option)
+
         for price in upgrade_prices:
             if price.get('categories') is None or price.get('item') is None:
                 continue
@@ -1006,6 +1043,11 @@ def _get_price_id_for_upgrade_option(self, upgrade_prices, option, value, public
                           or product.get('units') == 'DEDICATED_CORE')
 
             for category in price.get('categories'):
+                if option == 'disk':
+                    if (category_code == (''.join([i for i in category.get('categoryCode') if not i.isdigit()]))
+                            and str(product.get('capacity')) == str(value)):
+                        return price.get('id')
+
                 if not (category.get('categoryCode') == category_code
                         and str(product.get('capacity')) == str(value)):
                     continue

tests/CLI/modules/vs/vs_tests.py

@@ -542,6 +542,27 @@ def test_upgrade(self, confirm_mock):
         self.assertIn({'id': 1122}, order_container['prices'])
         self.assertEqual(order_container['virtualGuests'], [{'id': 100}])
 
+    @mock.patch('SoftLayer.CLI.formatting.confirm')
+    def test_upgrade_disk(self, confirm_mock):
+        confirm_mock.return_value = True
+        result = self.run_command(['vs', 'upgrade', '100', '--flavor=M1_64X512X100',
+                                   '--resize-disk=10', '1', '--resize-disk=10', '2'])
+        self.assert_no_fail(result)
+        self.assert_called_with('SoftLayer_Product_Order', 'placeOrder')
+        call = self.calls('SoftLayer_Product_Order', 'placeOrder')[0]
+        order_container = call.args[0]
+        self.assertEqual(799, order_container['presetId'])
+        self.assertIn({'id': 100}, order_container['virtualGuests'])
+        self.assertEqual(order_container['virtualGuests'], [{'id': 100}])
+
+    @mock.patch('SoftLayer.CLI.formatting.confirm')
+    def test_upgrade_disk_error(self, confirm_mock):
+        confirm_mock.return_value = True
+        result = self.run_command(['vs', 'upgrade', '100', '--flavor=M1_64X512X100',
+                                   '--resize-disk=1000', '1', '--resize-disk=10', '2'])
+        self.assertEqual(result.exit_code, 1)
+        self.assertIsInstance(result.exception, SoftLayerAPIError)
+
     @mock.patch('SoftLayer.CLI.formatting.confirm')
     def test_upgrade_with_flavor(self, confirm_mock):
         confirm_mock.return_value = True
@@ -554,6 +575,17 @@ def test_upgrade_with_flavor(self, confirm_mock):
         self.assertIn({'id': 100}, order_container['virtualGuests'])
         self.assertEqual(order_container['virtualGuests'], [{'id': 100}])
 
+    @mock.patch('SoftLayer.CLI.formatting.confirm')
+    def test_upgrade_with_add_disk(self, confirm_mock):
+        confirm_mock.return_value = True
+        result = self.run_command(['vs', 'upgrade', '100', '--add-disk=10', '--add-disk=10'])
+        self.assert_no_fail(result)
+        self.assert_called_with('SoftLayer_Product_Order', 'placeOrder')
+        call = self.calls('SoftLayer_Product_Order', 'placeOrder')[0]
+        order_container = call.args[0]
+        self.assertIn({'id': 100}, order_container['virtualGuests'])
+        self.assertEqual(order_container['virtualGuests'], [{'id': 100}])
+
     @mock.patch('SoftLayer.CLI.formatting.confirm')
     def test_upgrade_with_cpu_memory_and_flavor(self, confirm_mock):
         confirm_mock.return_value = True