vcpu_set: Implement plus-minus operators

innogames · Aug 12, 2024 · 7d5b5a9 · 7d5b5a9
1 parent e5eeab8
commit 7d5b5a9
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 6 deletions.
diff --git a/igvm/cli.py b/igvm/cli.py
@@ -319,8 +319,7 @@ def parse_args():
     )
     subparser.add_argument(
         'count',
-        type=int,
-        help='New number of CPUs',
+        help='New number of CPUs, integers with optional prefix + or -',
     )
     subparser.add_argument(
         '--offline',

diff --git a/igvm/commands.py b/igvm/commands.py
@@ -141,12 +141,20 @@ def vcpu_set(vm_hostname, count, offline=False):
             )
             offline = False
 
-        if count == vm.dataset_obj['num_cpu']:
+        if str(count).startswith('+'):
+            count = vm.dataset_obj['num_cpu'] + int(str(count)[1:])
+        elif str(count).startswith('-'):
+            if not offline:
+                raise IGVMError(
+                    'Decreasing CPU count is only allowed offline.'
+                )
+            count = vm.dataset_obj['num_cpu'] - int(str(count)[1:])
+        elif int(count) == vm.dataset_obj['num_cpu']:
             raise Warning('CPU count is the same.')
 
         if offline:
             vm.shutdown()
-        vm.set_num_cpu(count)
+        vm.set_num_cpu(int(count))
         if offline:
             vm.start()
 

diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -475,11 +475,21 @@ def _get_cpus_vm():
         with self.assertRaises(IGVMError):
             vcpu_set(VM_HOSTNAME, 0, offline=True)
 
+        # Has to be offline
         with self.assertRaises(IGVMError):
-            vcpu_set(VM_HOSTNAME, -5)
+            vcpu_set(VM_HOSTNAME, '-1')
 
+        # Not enough CPUs to remove
         with self.assertRaises(IGVMError):
-            vcpu_set(VM_HOSTNAME, -5, offline=True)
+            vcpu_set(VM_HOSTNAME, '-5', offline=True)
+
+        vcpu_set(VM_HOSTNAME, '+2')
+        self.assertEqual(_get_cpus_hv(), 4)
+        self.assertEqual(_get_cpus_vm(), 4)
+
+        vcpu_set(VM_HOSTNAME, '-2', offline=True)
+        self.assertEqual(_get_cpus_hv(), 2)
+        self.assertEqual(_get_cpus_vm(), 2)
 
     def test_sync(self):
         obj = (