diff --git a/mula/scheduler/connectors/services/octopoes.py b/mula/scheduler/connectors/services/octopoes.py index b6f5f6b844b..4e9ea22ff50 100644 --- a/mula/scheduler/connectors/services/octopoes.py +++ b/mula/scheduler/connectors/services/octopoes.py @@ -16,9 +16,12 @@ def __init__(self, host: str, source: str, orgs: List[Organisation]): @exception_handler def get_objects_by_object_types( - self, organisation_id: str, object_types: List[str], scan_level: List[int] = [] + self, organisation_id: str, object_types: List[str], scan_level: List[int] ) -> List[OOI]: """Get all oois from octopoes""" + if scan_level is None: + scan_level = [] + url = f"{self.host}/{organisation_id}/objects" params = { @@ -46,10 +49,20 @@ def get_objects_by_object_types( return oois @exception_handler - def get_random_objects(self, organisation_id: str, n: int) -> List[OOI]: + def get_random_objects(self, organisation_id: str, n: int, scan_level: List[int]) -> List[OOI]: """Get `n` random oois from octopoes""" + if scan_level is None: + scan_level = [] + url = f"{self.host}/{organisation_id}/objects/random" - response = self.get(url, params={"amount": str(n)}) + + params = { + "amount": str(n), + "scan_level": {s for s in scan_level}, + } + + response = self.get(url, params=params) + return [OOI(**ooi) for ooi in response.json()] @exception_handler diff --git a/mula/scheduler/schedulers/boefje.py b/mula/scheduler/schedulers/boefje.py index 14b119d151f..ca76bb91233 100644 --- a/mula/scheduler/schedulers/boefje.py +++ b/mula/scheduler/schedulers/boefje.py @@ -425,6 +425,7 @@ def push_tasks_for_random_objects(self) -> None: random_oois = self.ctx.services.octopoes.get_random_objects( organisation_id=self.organisation.id, n=self.ctx.config.pq_populate_max_random_objects, + scan_level=[1, 2, 3, 4], ) except (requests.exceptions.RetryError, requests.exceptions.ConnectionError): self.logger.warning( diff --git a/octopoes/octopoes/api/router.py b/octopoes/octopoes/api/router.py index 65240fb2091..d87f90c797d 100644 --- a/octopoes/octopoes/api/router.py +++ b/octopoes/octopoes/api/router.py @@ -139,8 +139,9 @@ def list_random_objects( octopoes: OctopoesService = Depends(octopoes_service), valid_time: datetime = Depends(extract_valid_time), amount: int = 1, + scan_level: Set[ScanLevel] = Query(DEFAULT_SCAN_LEVEL_FILTER), ) -> List[OOI]: - return octopoes.list_random_ooi(amount, valid_time) + return octopoes.list_random_ooi(valid_time, amount, scan_level) @router.delete("/") diff --git a/octopoes/octopoes/core/service.py b/octopoes/octopoes/core/service.py index 8220fbfd77a..f2e38383567 100644 --- a/octopoes/octopoes/core/service.py +++ b/octopoes/octopoes/core/service.py @@ -426,8 +426,10 @@ def _on_update_scan_profile(self, event: ScanProfileDBEvent) -> None: def _on_delete_scan_profile(self, event: ScanProfileDBEvent) -> None: self._run_inferences(event) - def list_random_ooi(self, amount: int, valid_time: datetime) -> List[OOI]: - oois = self.ooi_repository.list_random(amount, valid_time) + def list_random_ooi( + self, valid_time: datetime, amount: int = 1, scan_levels: Set[ScanLevel] = DEFAULT_SCAN_LEVEL_FILTER + ) -> List[OOI]: + oois = self.ooi_repository.list_random(valid_time, amount, scan_levels) self._populate_scan_profiles(oois, valid_time) return oois diff --git a/octopoes/octopoes/repositories/ooi_repository.py b/octopoes/octopoes/repositories/ooi_repository.py index 14c8cda649d..6d47de5cda2 100644 --- a/octopoes/octopoes/repositories/ooi_repository.py +++ b/octopoes/octopoes/repositories/ooi_repository.py @@ -77,7 +77,9 @@ def list( ) -> Paginated[OOI]: raise NotImplementedError - def list_random(self, amount: int, valid_time: datetime) -> List[OOI]: + def list_random( + self, valid_time: datetime, amount: int = 1, scan_levels: Set[ScanLevel] = DEFAULT_SCAN_LEVEL_FILTER + ) -> List[OOI]: raise NotImplementedError def list_neighbours(self, references: Set[Reference], paths: Set[Path], valid_time: datetime) -> Set[OOI]: @@ -261,20 +263,29 @@ def list( items=oois, ) - def list_random(self, amount: int, valid_time: datetime) -> List[OOI]: + def list_random( + self, valid_time: datetime, amount: int = 1, scan_levels: Set[ScanLevel] = DEFAULT_SCAN_LEVEL_FILTER + ) -> List[OOI]: query = """ - {{ - :query {{ - :find [(rand {amount} ?id)] - :where [ - [?e :crux.db/id ?id] - [?e :object_type] - ] + {{ + :query {{ + :find [(rand {amount} ?id)] + :in [[_scan_level ...]] + :where [ + [?e :crux.db/id ?id] + [?e :object_type] + [?scan_profile :type "ScanProfile"] + [?scan_profile :reference ?e] + [?scan_profile :level _scan_level] + ] + }} + :in-args [[{scan_levels}]] }} - }} - """.format( - amount=amount + """.format( + amount=amount, + scan_levels=" ".join([str(scan_level.value) for scan_level in scan_levels]), ) + res = self.session.client.query(query, valid_time) if not res: return [] diff --git a/octopoes/tests/robot/02_list_objects.robot b/octopoes/tests/robot/02_list_objects.robot index 853088e8ffe..aeb74b25ff5 100644 --- a/octopoes/tests/robot/02_list_objects.robot +++ b/octopoes/tests/robot/02_list_objects.robot @@ -18,6 +18,15 @@ List Objects With Filter Verify Object List With Filter +List Random Objects With Filter + Insert Normalizer Output + Declare Scan Profile ${REF_HOSTNAME} ${1} + Await Sync + Length Of Random Object List With Filter Should Be ${1} ${5} + Length Of Random Object List With Filter Should Be ${0} ${1} + Length Of Random Object List With Filter Should Be ${{ [1,0] }} ${6} + Length Of Random Object List With Filter Should Be ${{ [2,3] }} ${0} + *** Keywords *** Setup Test Start Monitoring ${QUEUE_URI} @@ -35,3 +44,9 @@ Get Objects With ScanLevel 0 ${response} Get ${OCTOPOES_URI}/objects params=scan_level=0 ${response_data} Set Variable ${response.json()} RETURN ${response_data} + +Length Of Random Object List With Filter Should Be + [Arguments] ${scan_levels} ${expected_length} + ${params} = Create Dictionary scan_level=${scan_levels} amount=10 + ${response} Get ${OCTOPOES_URI}/objects/random params=${params} + Length Should Be ${response.json()} ${expected_length}