Merge pull request #773 from Ostorlab/fix/delete_agent_group_mapping_…

…on_deletion_mutation Delete agent group mappings when deleting agent_group
Ostorlab · Aug 7, 2024 · aa133e4 · aa133e4
2 parents 79e56da + 2bf8a2d
commit aa133e4
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 8 deletions.
diff --git a/src/ostorlab/serve_app/oxo.py b/src/ostorlab/serve_app/oxo.py
@@ -585,7 +585,14 @@ def mutate(
             )
             if agent_group_query.count() == 0:
                 raise graphql.GraphQLError("AgentGroup not found.")
+            mappings = session.query(models.AgentGroupMapping).filter_by(
+                agent_group_id=agent_group_id
+            )
+            if mappings.count() > 0:
+                mappings.delete()
+
             agent_group_query.delete()
+
             session.commit()
             return DeleteAgentGroupMutation(result=True)
 

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -37,6 +37,7 @@
 from ostorlab.scanner.proto.assets import apk_pb2
 from ostorlab.scanner.proto.scan._location import startAgentScan_pb2
 from ostorlab.serve_app import app
+from ostorlab.serve_app import types
 from ostorlab.utils import risk_rating
 
 
@@ -931,10 +932,12 @@ def agent_group(
     """Create dummy agent group."""
     mocker.patch.object(models, "ENGINE_URL", db_engine_path)
     with models.Database() as session:
-        agent_group = models.AgentGroup(
-            name="Agent Group 1",
-            description="Agent Group 1 description",
-            created_time=datetime.datetime.now(),
+        agent = types.OxoAgentGroupAgentCreateInputType()
+        agent.key = "key"
+        agent.args = []
+        agents = [agent]
+        agent_group = models.AgentGroup.create(
+            name="Agent Group 1", description="Agent Group 1 description", agents=agents
         )
         session.add(agent_group)
         session.commit()

diff --git a/tests/serve_app/oxo_test.py b/tests/serve_app/oxo_test.py
@@ -19,7 +19,6 @@
 
 RE_OXO_ENDPOINT = "https://api.ostorlab.co/apis/oxo"
 
-
 INTROSPECT_ENUMS_QUERY = """
     {
         __schema {
@@ -34,7 +33,6 @@
     }    
 """
 
-
 INTROSPECT_INPUTS_QUERY = """
     {
         __schema {
@@ -131,7 +129,6 @@
 }
 """
 
-
 INTROSPECT_UNIONS_QUERY = """
 {
     __schema {
@@ -151,7 +148,6 @@
 }    
 """
 
-
 INTROSPECT_TYPES_QUERY = """
     {
         __schema {
@@ -2251,6 +2247,11 @@ def testDeleteAgentGroupMutation_whenAgentGroupExist_deleteAgentGroup(
     mocker.patch.object(models, "ENGINE_URL", db_engine_path)
     with models.Database() as session:
         nbr_agent_groups_before_delete = session.query(models.AgentGroup).count()
+        nbr_mappings_agent_group_before_delete = (
+            session.query(models.AgentGroupMapping)
+            .filter_by(agent_group_id=agent_group.id)
+            .count()
+        )
 
     query = """
         mutation DeleteAgentGroup ($agentGroupId: Int!){
@@ -2271,6 +2272,12 @@ def testDeleteAgentGroupMutation_whenAgentGroupExist_deleteAgentGroup(
             session.query(models.AgentGroup).count()
             == nbr_agent_groups_before_delete - 1
         )
+        assert (
+            session.query(models.AgentGroupMapping)
+            .filter_by(agent_group_id=agent_group.id)
+            .count()
+            == nbr_mappings_agent_group_before_delete - 1
+        )
 
 
 def testDeleteAgentGroupMutation_whenAgentGroupDoesNotExist_returnErrorMessage(