vcpu_set: Implement plus-minus operators

innogames · Apr 5, 2024 · 0aff980 · 0aff980
1 parent b2326e9
commit 0aff980
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 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,7 +141,11 @@ 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 count == vm.dataset_obj['num_cpu']:
             raise Warning('CPU count is the same.')
 
         if offline:

diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -481,6 +481,14 @@ def _get_cpus_vm():
         with self.assertRaises(IGVMError):
             vcpu_set(VM_HOSTNAME, -5, offline=True)
 
+        vcpu_set(VM_HOSTNAME, -1, offline=True)
+        self.assertEqual(_get_cpus_hv(), 1)
+        self.assertEqual(_get_cpus_vm(), 1)
+
+        vcpu_set(VM_HOSTNAME, +1)
+        self.assertEqual(_get_cpus_hv(), 2)
+        self.assertEqual(_get_cpus_vm(), 2)
+
     def test_sync(self):
         obj = (
             Query({'hostname': VM_HOSTNAME}, ['disk_size_gib', 'memory']).get()