fix and improve the new code

Author: caberos

SoftLayer/CLI/virt/upgrade.py

@@ -1,8 +1,6 @@
 """Upgrade a virtual server."""
 # :license: MIT, see LICENSE for more details.
 
-import json
-
 import click
 
 import SoftLayer
@@ -22,18 +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', type=click.INT, required=False, help="add Hard disk in GB")
-@click.option('--disk', nargs=1, help="update the number and capacity in GB Hard disk, E.G {'number':2,'capacity':100}")
+@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, disk, add):
+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, disk, add]):
-        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]")
@@ -44,9 +44,18 @@ def cli(env, identifier, cpu, private, memory, network, flavor, disk, add):
 
     if memory:
         memory = int(memory / 1024)
-    if disk is not None:
-        disk = json.loads(disk)
+    if resize_disk:
+        disk_json = list()
+        for guest_disk in resize_disk:
+            disks = {'capacity': guest_disk[0], 'number': guest_disk[1]}
+            disk_json.append(disks)
+
+    elif add_disk:
+        disk_json = list()
+        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, add=add):
+                       disk=disk_json):
         raise exceptions.CLIAbort('VS Upgrade Failed')

SoftLayer/managers/vs.py

@@ -819,8 +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,
-                disk=None, add=None):
+    def upgrade(self, instance_id, cpus=None, memory=None, nic_speed=None, public=True, preset=None, disk=None):
         """Upgrades a VS instance.
 
         Example::
@@ -853,10 +852,7 @@ def upgrade(self, instance_id, cpus=None, memory=None, nic_speed=None, public=Tr
         if memory is not None and preset is not None:
             raise ValueError("Do not use memory, private or cpu if you are using flavors")
         data['memory'] = memory
-        if disk is not None:
-            data['disk'] = disk.get('capacity')
-        elif add is not None:
-            data['disk'] = add
+
         maintenance_window = datetime.datetime.now(utils.UTC())
         order = {
             'complexType': 'SoftLayer_Container_Product_Order_Virtual_Guest_Upgrade',
@@ -867,6 +863,35 @@ 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
+
+                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
@@ -879,28 +904,7 @@ def upgrade(self, instance_id, cpus=None, memory=None, nic_speed=None, public=Tr
                 raise exceptions.SoftLayerError(
                     "Unable to find %s option with value %s" % (option, value))
 
-            if disk is not None:
-                category = {'categories': [{
-                    'categoryCode': 'guest_disk' + str(disk.get('number')),
-                    'complexType': "SoftLayer_Product_Item_Category"
-                }], 'complexType': 'SoftLayer_Product_Item_Price'}
-                prices.append(category)
-                prices[0]['id'] = price_id
-            elif add:
-                vsi_disk = self.get_instance(instance_id)
-                disk_number = 0
-                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'))))
-                category = {'categories': [{
-                    'categoryCode': 'guest_disk' + str(disk_number + 1),
-                    'complexType': "SoftLayer_Product_Item_Category"
-                }], 'complexType': 'SoftLayer_Product_Item_Price'}
-                prices.append(category)
-                prices[0]['id'] = price_id
-            else:
-                prices.append({'id': price_id})
+            prices.append({'id': price_id})
 
             order['prices'] = prices
         if preset is not None:
@@ -1024,6 +1028,7 @@ def _get_price_id_for_upgrade_option(self, upgrade_prices, option, value, public
             '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
@@ -1034,7 +1039,7 @@ def _get_price_id_for_upgrade_option(self, upgrade_prices, option, value, public
 
             for category in price.get('categories'):
                 if option == 'disk':
-                    if not (category_code == (''.join([i for i in category.get('categoryCode') if not i.isdigit()]))
+                    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')
 
@@ -1052,8 +1057,6 @@ def _get_price_id_for_upgrade_option(self, upgrade_prices, option, value, public
                 elif option == 'nic_speed':
                     if 'Public' in product.get('description'):
                         return price.get('id')
-                elif option == 'disk':
-                    return price.get('id')
                 else:
                     return price.get('id')
 

tests/CLI/modules/vs/vs_tests.py

@@ -370,19 +370,19 @@ def test_dns_sync_both(self, confirm_mock):
                                            'getResourceRecords')
         getResourceRecords.return_value = []
         createAargs = ({
-            'type': 'a',
-            'host': 'vs-test1',
-            'domainId': 12345,  # from SoftLayer_Account::getDomains
-            'data': '172.16.240.2',
-            'ttl': 7200
-        },)
+                           'type': 'a',
+                           'host': 'vs-test1',
+                           'domainId': 12345,  # from SoftLayer_Account::getDomains
+                           'data': '172.16.240.2',
+                           'ttl': 7200
+                       },)
         createPTRargs = ({
-            'type': 'ptr',
-            'host': '2',
-            'domainId': 123456,
-            'data': 'vs-test1.test.sftlyr.ws',
-            'ttl': 7200
-        },)
+                             'type': 'ptr',
+                             'host': '2',
+                             'domainId': 123456,
+                             'data': 'vs-test1.test.sftlyr.ws',
+                             'ttl': 7200
+                         },)
 
         result = self.run_command(['vs', 'dns-sync', '100'])
 
@@ -425,12 +425,12 @@ def test_dns_sync_v6(self, confirm_mock):
             }
         }
         createV6args = ({
-            'type': 'aaaa',
-            'host': 'vs-test1',
-            'domainId': 12345,
-            'data': '2607:f0d0:1b01:0023:0000:0000:0000:0004',
-            'ttl': 7200
-        },)
+                            'type': 'aaaa',
+                            'host': 'vs-test1',
+                            'domainId': 12345,
+                            'data': '2607:f0d0:1b01:0023:0000:0000:0000:0004',
+                            'ttl': 7200
+                        },)
         guest.return_value = test_guest
         result = self.run_command(['vs', 'dns-sync', '--aaaa-record', '100'])
         self.assert_no_fail(result)
@@ -574,7 +574,7 @@ def test_upgrade(self, confirm_mock):
     def test_upgrade_disk(self, confirm_mock):
         confirm_mock.return_value = True
         result = self.run_command(['vs', 'upgrade', '100', '--flavor=M1_64X512X100',
-                                   '--disk={"number":1,"capacity":10}'])
+                                   '--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]
@@ -598,7 +598,7 @@ def test_upgrade_with_flavor(self, confirm_mock):
     @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=10'])
+        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]