From 7d5b5a99bb6f412725b28e22467b93133c6d9f9b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=BCrkan=20G=C3=BCr?= <gurkan.gur@innogames.com>
Date: Fri, 5 Apr 2024 20:48:58 +0200
Subject: [PATCH] vcpu_set: Implement plus-minus operators

---
 igvm/cli.py               |  3 +--
 igvm/commands.py          | 12 ++++++++++--
 tests/test_integration.py | 14 ++++++++++++--
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/igvm/cli.py b/igvm/cli.py
index 0173f29e..e9f64366 100644
--- 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
index 6b1906b2..bea8f0c5 100644
--- 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
index 211f01e5..481850c8 100644
--- 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 = (