From 9f7b4b05c343ecb8d9f19d60c7b75b787a74d96b Mon Sep 17 00:00:00 2001
From: adnane <adnane.serrar@ostorlab.dev>
Date: Wed, 7 Aug 2024 13:31:02 +0100
Subject: [PATCH] Delete agent group mappings when deleting agent_group

---
 src/ostorlab/serve_app/oxo.py |  7 +++++++
 tests/conftest.py             | 11 +++++++----
 tests/serve_app/oxo_test.py   | 15 +++++++++++----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/src/ostorlab/serve_app/oxo.py b/src/ostorlab/serve_app/oxo.py
index 387ccf765..9e1183a43 100644
--- a/src/ostorlab/serve_app/oxo.py
+++ b/src/ostorlab/serve_app/oxo.py
@@ -577,7 +577,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
index d69ab695c..ebc7b2193 100644
--- 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
index 5ecbcb1ac..611b1b569 100644
--- 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 {
@@ -2252,6 +2248,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!){
@@ -2272,6 +2273,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(