diff --git a/src/service/service/task_scheduler/TaskExecutor.py b/src/service/service/task_scheduler/TaskExecutor.py index e475f29c65b898895838901ca86636bbd5b9ea88..54aef8c37d1d8aaf3946606d056595e1c7e1538d 100644 --- a/src/service/service/task_scheduler/TaskExecutor.py +++ b/src/service/service/task_scheduler/TaskExecutor.py @@ -105,6 +105,22 @@ class TaskExecutor: self._device_client.ConfigureDevice(device) self._store_grpc_object(CacheableObjectType.DEVICE, device_key, device) + def get_device_manager(self, device : Device) -> Optional[Device]: + json_manager = None + for config_rule in device.device_config.config_rules: + if config_rule.WhichOneof('config_rule') != 'custom': continue + if config_rule.custom.resource_key != '_manager': continue + json_manager = json.loads(config_rule.custom.resource_value) + break + + if json_manager is None: return None + + manager_uuid = json_manager['uuid'] + manager = self.get_device(DeviceId(**json_device_id(manager_uuid))) + manager_uuid = manager.device_id.device_uuid.uuid + if manager is None: raise Exception('Device({:s}) not found'.format(str(manager_uuid))) + return manager + def get_devices_from_connection( self, connection : Connection, exclude_managed : bool = False ) -> Dict[str, Device]: @@ -114,23 +130,13 @@ class TaskExecutor: device_uuid = endpoint_id.device_id.device_uuid.uuid if device is None: raise Exception('Device({:s}) not found'.format(str(device_uuid))) - manager = None - for config_rule in device.device_config.config_rules: - if config_rule.WhichOneof('config_rule') != 'custom': continue - if config_rule.custom.resource_key != '_manager': continue - manager = json.loads(config_rule.custom.resource_value) - break - - if manager is not None: + manager = self.get_device_manager(device) + if manager is None: + devices[device_uuid] = device + else: if not exclude_managed: devices[device_uuid] = device - manager_uuid = manager['uuid'] - manager = self.get_device(DeviceId(**json_device_id(manager_uuid))) - manager_uuid = manager.device_id.device_uuid.uuid - if manager is None: raise Exception('Device({:s}) not found'.format(str(manager_uuid))) - devices[manager_uuid] = manager - else: - devices[device_uuid] = device + devices[manager.device_id.device_uuid.uuid] = manager return devices # ----- Service-related methods ------------------------------------------------------------------------------------