From ac4c5f598513509482077e6944020ca87e94188f Mon Sep 17 00:00:00 2001
From: grossmj <grossmj@gns3.net>
Date: Fri, 3 Jan 2025 19:56:41 +0700
Subject: [PATCH] Improve tests for pruning images

---
 gns3server/db/repositories/templates.py       |  2 +-
 tests/api/routes/controller/test_images.py    |  5 ++++-
 tests/api/routes/controller/test_templates.py | 22 ++++++++++++++-----
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/gns3server/db/repositories/templates.py b/gns3server/db/repositories/templates.py
index ef5eb525..ec8215af 100644
--- a/gns3server/db/repositories/templates.py
+++ b/gns3server/db/repositories/templates.py
@@ -171,7 +171,7 @@ class TemplatesRepository(BaseRepository):
             await self._db_session.refresh(template_in_db)
         return template_in_db
 
-    async def get_template_images(self, template_id: UUID) -> int:
+    async def get_template_images(self, template_id: UUID) -> List[models.Image]:
         """
         Return all images attached to a template.
         """
diff --git a/tests/api/routes/controller/test_images.py b/tests/api/routes/controller/test_images.py
index be000d97..7a8ac4d7 100644
--- a/tests/api/routes/controller/test_images.py
+++ b/tests/api/routes/controller/test_images.py
@@ -261,10 +261,13 @@ class TestImageRoutes:
 
     async def test_prune_images(self, app: FastAPI, client: AsyncClient, db_session: AsyncSession) -> None:
 
+        images_repo = ImagesRepository(db_session)
+        images_in_db = await images_repo.get_images()
+        assert len(images_in_db) != 0
+
         response = await client.delete(app.url_path_for("prune_images"))
         assert response.status_code == status.HTTP_204_NO_CONTENT
 
-        images_repo = ImagesRepository(db_session)
         images_in_db = await images_repo.get_images()
         assert len(images_in_db) == 0
 
diff --git a/tests/api/routes/controller/test_templates.py b/tests/api/routes/controller/test_templates.py
index e66e5df6..25ce4922 100644
--- a/tests/api/routes/controller/test_templates.py
+++ b/tests/api/routes/controller/test_templates.py
@@ -16,6 +16,8 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import os
+import shutil
+
 import pytest
 import uuid
 import unittest.mock
@@ -145,29 +147,39 @@ class TestTemplateRoutes:
             tmpdir: str,
     ) -> None:
 
-        path = os.path.join(tmpdir, "test.qcow2")
-        with open(path, "wb+") as f:
+        image1 = os.path.join(tmpdir, "image1.qcow2")
+        with open(image1, "wb+") as f:
             f.write(b'\x42\x42\x42\x42')
+
+        image2 = os.path.join(tmpdir, "image2.qcow2")
+        with open(image2, "wb+") as f:
+            f.write(b'\x42\x42\x42\x42')
+
         images_repo = ImagesRepository(db_session)
-        await images_repo.add_image("test.qcow2", "qemu", 42, path, "e342eb86c1229b6c154367a5476969b5", "md5")
+        await images_repo.add_image("image1.qcow2", "qemu", 42, image1, "e342eb86c1229b6c154367a5476969b5", "md5")
+        await images_repo.add_image("image2.qcow2", "qemu", 42, image2, "e342eb86c1229b6c154367a5476969b5", "md5")
 
         template_id = str(uuid.uuid4())
         params = {"template_id": template_id,
                   "name": "QEMU_TEMPLATE",
                   "compute_id": "local",
-                  "hda_disk_image": "test.qcow2",
+                  "hda_disk_image": "image1.qcow2",
+                  "hdb_disk_image": "image2.qcow2",
                   "template_type": "qemu"}
 
         response = await client.post(app.url_path_for("create_template"), json=params)
         assert response.status_code == status.HTTP_201_CREATED
 
+        templates_repo = TemplatesRepository(db_session)
+        images = await templates_repo.get_template_images(response.json().get("template_id"))
+        assert len(images) == 2
+
         response = await client.delete(
             app.url_path_for("delete_template", template_id=template_id),
             params={"prune_images": True}
         )
         assert response.status_code == status.HTTP_204_NO_CONTENT
 
-        images_repo = ImagesRepository(db_session)
         images = await images_repo.get_images()
         assert len(images) == 0