Feature/octopoes/filter random endpoint (#704)

Co-authored-by: Jesse Lisser <[email protected]> Co-authored-by: Jesse Lisser <[email protected]> Co-authored-by: Patrick <[email protected]>
minvws · Apr 25, 2023 · 478a77c · 478a77c
1 parent af08d69
commit 478a77c
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 18 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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}