vcpu_set: Implement plus-minus operators

innogames · Apr 29, 2024 · 0e09185 · 0e09185
1 parent b2326e9
commit 0e09185
Show file tree

Hide file tree

Showing 3 changed files with 17 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,16 @@ 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('-'):
+            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
@@ -476,10 +476,18 @@ def _get_cpus_vm():
             vcpu_set(VM_HOSTNAME, 0, offline=True)
 
         with self.assertRaises(IGVMError):
-            vcpu_set(VM_HOSTNAME, -5)
+            vcpu_set(VM_HOSTNAME, '-5')
 
         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 = (