Fixing up generics, adding support for random & keep on error to fixed, some minor fixes

server/src/uds/core/services/generics/dynamic/userservice.py

@@ -191,6 +191,7 @@ class DynamicUserService(services.UserService, autoserializable.AutoSerializable
         """
         return self.name_generator().get(self.service().get_basename(), self.service().get_lenname())
 
+    @typing.final
     def _error(self, reason: typing.Union[str, Exception]) -> types.states.TaskState:
         """
         Internal method to set object as error state
@@ -209,11 +210,12 @@ class DynamicUserService(services.UserService, autoserializable.AutoSerializable
                     self.service().remove(self, self._vmid)
                     self._vmid = ''
                 except Exception as e:
-                    logger.exception('Exception removing machine: %s', e)
+                    logger.exception('Exception removing machine %s: %s', self._vmid, e)
+                    self._vmid = ''
+                    self.do_log(log.LogLevel.ERROR, f'Error removing machine: {e}')
             else:
                 logger.debug('Keep on error is enabled, not removing machine')
-                if self.keep_state_sets_error is False:
-                    self._set_queue([Operation.FINISH])
+                self._set_queue([Operation.FINISH] if self.keep_state_sets_error else [Operation.ERROR])
                 return types.states.TaskState.FINISHED
 
         self._set_queue([Operation.ERROR])
@@ -258,7 +260,9 @@ class DynamicUserService(services.UserService, autoserializable.AutoSerializable
                     # Provide self to the service, so it can use some of our methods for whaterever it needs
                     self._ip = self.service().get_ip(self, self._vmid)
             except Exception:
-                logger.warning('Error obtaining IP for %s: %s', self.__class__.__name__, self._vmid, exc_info=True)
+                logger.warning(
+                    'Error obtaining IP for %s: %s', self.__class__.__name__, self._vmid, exc_info=True
+                )
         return self._ip
 
     @typing.final
@@ -293,8 +297,11 @@ class DynamicUserService(services.UserService, autoserializable.AutoSerializable
     def set_ready(self) -> types.states.TaskState:
         # If already ready, return finished
         try:
-            if self.cache.get('ready') == '1' or self.service().is_running(self, self._vmid):
-                self._set_queue([Operation.START_COMPLETED, Operation.FINISH])
+            if self.cache.get('ready', '0') == '1':
+                self._set_queue([Operation.FINISH])
+            elif self.service().is_running(self, self._vmid):
+                self.cache.put('ready', '1', consts.cache.SHORT_CACHE_TIMEOUT // 2)  # short cache timeout
+                self._set_queue([Operation.FINISH])
             else:
                 self._set_queue([Operation.START, Operation.START_COMPLETED, Operation.FINISH])
         except Exception as e:
@@ -411,7 +418,7 @@ class DynamicUserService(services.UserService, autoserializable.AutoSerializable
             self._set_queue([op] + destroy_operations)
             # Do not execute anything.here, just continue normally
         return types.states.TaskState.RUNNING
-    
+
     def error_reason(self) -> str:
         return self._reason
 
@@ -667,7 +674,7 @@ class DynamicUserService(services.UserService, autoserializable.AutoSerializable
         This method is called to check if the service is reset
         """
         return types.states.TaskState.FINISHED
-    
+
     def op_reset_completed_checker(self) -> types.states.TaskState:
         """
         This method is called to check if the service reset is completed

server/src/uds/core/services/generics/fixed/service.py

@@ -38,6 +38,7 @@ import collections.abc
 from django.utils.translation import gettext_noop as _, gettext
 from uds.core import services, types, exceptions
 from uds.core.ui import gui
+from uds.core.util import fields
 
 # Not imported at runtime, just for type checking
 if typing.TYPE_CHECKING:
@@ -124,10 +125,14 @@ class FixedService(services.Service, abc.ABC):  # pylint: disable=too-many-publi
     # Randomize machine assignation isntead of linear
     randomize = gui.CheckBoxField(
         label=_('Randomize machine assignation'),
-        order=33,
+        order=100,
         default=True,
         tooltip=_('If active, UDS will assign machines in a random way, instead of linear'),
-        tab=types.ui.Tab.MACHINE,
+        tab=types.ui.Tab.ADVANCED,
+    )
+    maintain_on_error = fields.maintain_on_error_field(
+        order=101,
+        tab=types.ui.Tab.ADVANCED,
     )
 
     def initialize(self, values: 'types.core.ValuesType') -> None:
@@ -191,17 +196,24 @@ class FixedService(services.Service, abc.ABC):  # pylint: disable=too-many-publi
         raise NotImplementedError()
 
     # default implementation, should be sufficient for most cases
-    def remove_and_free(self, vmid: str) -> str:
+    def remove_and_free(self, vmid: str) -> types.states.TaskState:
         try:
             with self._assigned_access() as assigned:
                 # In error situations, due to the "process_snapshot" post runasign, the element could be already removed
                 # So we need to check if it's there
                 if vmid in assigned:
                     assigned.remove(vmid)
-            return types.states.State.FINISHED
+            return types.states.TaskState.FINISHED
         except Exception as e:
             logger.error('Error processing remove and free: %s', e)
             raise Exception(f'Error processing remove and free: {e} on {vmid}') from e
+        
+    def is_ready(self, vmid: str) -> bool:
+        """
+        Returns if the machine is ready for usage
+        Defaults to True
+        """
+        return True
 
     @abc.abstractmethod
     def get_first_network_mac(self, vmid: str) -> str:
@@ -231,7 +243,6 @@ class FixedService(services.Service, abc.ABC):  # pylint: disable=too-many-publi
         """
         raise NotImplementedError()
 
-    @typing.final
     def sorted_assignables_list(self, alternate_field_name: typing.Optional[str] = None) -> list[str]:
         """
         Randomizes the assignation of machines if needed
@@ -249,3 +260,8 @@ class FixedService(services.Service, abc.ABC):  # pylint: disable=too-many-publi
             return random.sample(machines.as_list(), len(machines.as_list()))
 
         return machines.as_list()
+
+    def should_maintain_on_error(self) -> bool:
+        if self.has_field('maintain_on_error'):  # If has been defined on own class...
+            return self.maintain_on_error.value
+        return False

server/src/uds/core/services/generics/fixed/userservice.py

@@ -35,7 +35,7 @@ import logging
 import typing
 import collections.abc
 
-from uds.core import services, types
+from uds.core import consts, services, types
 from uds.core.types.services import Operation
 from uds.core.util import log, autoserializable
 
@@ -110,6 +110,7 @@ class FixedUserService(services.UserService, autoserializable.AutoSerializable,
         self._queue.insert(0, Operation.NOP)
         return types.states.TaskState.RUNNING
 
+    @typing.final
     def _error(self, reason: typing.Union[str, Exception]) -> types.states.TaskState:
         """
         Internal method to set object as error state
@@ -122,12 +123,17 @@ class FixedUserService(services.UserService, autoserializable.AutoSerializable,
         self.do_log(log.LogLevel.ERROR, reason)
 
         if self._vmid:
-            try:
-                self.service().remove_and_free(self._vmid)
-                self.service().process_snapshot(remove=True, userservice_instance=self)
-                self._vmid = ''
-            except Exception as e:
-                logger.exception('Exception removing machine: %s', e)
+            if self.service().should_maintain_on_error() is False:
+                try:
+                    self.service().remove_and_free(self._vmid)
+                    self.service().process_snapshot(remove=True, userservice_instance=self)
+                    self._vmid = ''
+                except Exception as e:
+                    logger.exception('Exception removing machine: %s', e)
+            else:
+                logger.debug('Keep on error is enabled, not removing machine')
+                self._queue = [Operation.FINISH]
+                return types.states.TaskState.FINISHED
 
         self._queue = [Operation.ERROR]
         self._reason = reason
@@ -177,6 +183,20 @@ class FixedUserService(services.UserService, autoserializable.AutoSerializable,
         """
         return self._error('Cache for fixed userservices not supported')
 
+    def set_ready(self) -> types.states.TaskState:
+        # If already ready, return finished
+        try:
+            if self.cache.get('ready', '0') == '1':
+                self._queue = [Operation.FINISH]
+            elif self.service().is_ready(self._vmid):
+                self.cache.put('ready', '1', consts.cache.SHORT_CACHE_TIMEOUT // 2)  # short cache timeout
+                self._queue = [Operation.FINISH]
+            else:
+                self._queue = [Operation.START, Operation.START_COMPLETED, Operation.FINISH]
+        except Exception as e:
+            return self._error(f'Error on setReady: {e}')
+        return self._execute_queue()
+
     @typing.final
     def assign(self, vmid: str) -> types.states.TaskState:
         logger.debug('Assigning from VM {}'.format(vmid))

server/tests/core/services/generics/fixtures.py

@@ -78,6 +78,8 @@ class FixedTestingService(fixed_service.FixedService):
     token = fixed_service.FixedService.token
     snapshot_type = fixed_service.FixedService.snapshot_type
     machines = fixed_service.FixedService.machines
+    randomize = fixed_service.FixedService.randomize
+    maintain_on_error = fixed_service.FixedService.maintain_on_error
 
     user_service_type = FixedTestingUserService
     first_process_called = False
@@ -108,10 +110,14 @@ class FixedTestingService(fixed_service.FixedService):
         self.assigned_machine = 'assigned'
         return self.assigned_machine
 
-    def remove_and_free(self, vmid: str) -> str:
+    def remove_and_free(self, vmid: str) -> types.states.TaskState:
         self.mock.remove_and_free_machine(vmid)
         self.assigned_machine = ''
         return types.states.TaskState.FINISHED
+    
+    def is_ready(self, vmid: str) -> bool:
+        self.mock.is_ready(vmid)
+        return True
 
     def get_first_network_mac(self, vmid: str) -> str:
         self.mock.get_first_network_mac(vmid)
@@ -511,7 +517,7 @@ class DynamicTestingService(dynamic_service.DynamicService):
     ) -> None:
         self.mock.suspend(caller_instance, vmid)
         self.machine_running_flag = False
-
+        
 
 class DynamicTestingPublication(dynamic_publication.DynamicPublication):
     mock: 'mock.Mock' = mock.MagicMock()

server/tests/core/services/generics/test_dynamic_service.py

@@ -238,6 +238,21 @@ class DynamicServiceTest(UDSTestCase):
         
         service.mock.shutdown.assert_called_once_with(userservice, userservice._vmid)
         
+    def test_service_set_ready(self) -> None:
+        service = fixtures.create_dynamic_service()
+        userservice = fixtures.create_dynamic_userservice(service)
+        
+        # full deploy
+        state = userservice.deploy_for_user(models.User())
+        self.assertEqual(state, types.states.TaskState.RUNNING)
+        for _ in limited_iterator(lambda: state != types.states.TaskState.FINISHED, limit=128):
+            state = userservice.check_state()
+
+        # Call for set_ready
+        service.mock.reset_mock()
+        self.assertEqual(userservice.set_ready(), types.states.TaskState.FINISHED)
+        # is_ready should have been called
+        service.mock.is_running.assert_called_once()
         
 
 

server/tests/core/services/generics/test_fixed_service.py

@@ -246,3 +246,41 @@ class FixedServiceTest(UDSTestCase):
 
         # Userservice is in deployed state, so we can remove it
         self.check_iterations(service, userservice, EXPECTED_REMOVAL_ITERATIONS_INFO, removal=True)
+        
+    def test_service_set_ready(self) -> None:
+        _prov, service, userservice = self.create_elements()
+        self.deploy_service(service, userservice)
+        # Call for set_ready
+        self.assertEqual(userservice.set_ready(), types.states.TaskState.FINISHED)
+        # is_ready should have been called
+        service.mock.is_ready.assert_called_once()
+
+    def test_service_random_machine_list(self) -> None:
+        _prov, service, _userservice = self.create_elements()
+        service.machines.value = [f'machine{i}' for i in range(10)]
+        for randomized in (True, False):
+            service.randomize.value = randomized
+            machines_list = service.sorted_assignables_list()
+            if randomized:
+                self.assertNotEqual(machines_list, service.machines.value)
+                self.assertEqual(len(service.machines.value), len(machines_list))
+                self.assertEqual(sorted(service.machines.value), sorted(machines_list))
+            else:
+                self.assertEqual(machines_list, service.machines.value)
+
+    def test_service_keep_on_error(self) -> None:
+        for maintain_on_error in (True, False):
+            _prov, service, userservice = self.create_elements()
+            service.machines.value = [f'machine{i}' for i in range(10)]
+            service.maintain_on_error.value = maintain_on_error
+            self.deploy_service(service, userservice)
+            self.assertEqual(userservice.check_state(), types.states.TaskState.FINISHED)
+            
+            # Now, ensure that we will raise an exception, overriding is_ready of service
+            service.is_ready = mock.MagicMock(side_effect=Exception('Error'))
+            if maintain_on_error is False:
+                self.assertEqual(userservice.set_ready(), types.states.TaskState.ERROR)
+            else:
+                self.assertEqual(userservice.set_ready(), types.states.TaskState.FINISHED)
+            
+

server/tests/services/proxmox/test_userservice_linked.py

@@ -202,6 +202,8 @@ class TestProxmoxLinkedUserService(UDSTransactionTestCase):
             api.get_machine_pool_info.assert_called_with(vmid, service.pool.value, force=True)
             api.start_machine.assert_called_with(vmid)
 
+            # Ensure vm is stopped
+            fixtures.replace_vm_info(vmid, status='stopped')
             # Set ready state with the valid machine
             state = userservice.set_ready()
             # Machine is stopped, so task must be RUNNING (opossed to FINISHED)