Merge pull request #1050 from thunderstore-io/test-data-icons

Add option to reuse PackageVersion icon on test data creation

django/thunderstore/core/management/commands/content/base.py

@@ -1,10 +1,13 @@
 import functools
+import io
 import os
 from abc import ABC
 from dataclasses import dataclass, field
 from typing import Collection, List, Type
 
+from django.core.files.base import File
 from django.db.models import Model
+from PIL import Image
 
 from django_contracts.models import LegalContract
 from thunderstore.community.models import Community
@@ -27,6 +30,7 @@ class ContentPopulatorContext:
     contract_count: int = 0
     contract_version_count: int = 0
     wiki_page_count: int = 0
+    reuse_icon: bool = False
 
 
 class ContentPopulator(ABC):
@@ -75,3 +79,11 @@ def dummy_markdown() -> str:
 Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint
 occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.
 """
+
+
+def dummy_package_icon() -> File:
+    file_obj = io.BytesIO()
+    image = Image.new("RGB", (256, 256), "#231F36")
+    image.save(file_obj, format="PNG")
+    file_obj.seek(0)
+    return File(file_obj, name="dummy.png")

django/thunderstore/core/management/commands/content/package_version.py

@@ -3,8 +3,8 @@
 from thunderstore.core.management.commands.content.base import (
     ContentPopulator,
     ContentPopulatorContext,
+    dummy_package_icon,
 )
-from thunderstore.repository.factories import PackageVersionFactory
 from thunderstore.repository.models import Package, PackageVersion
 from thunderstore.utils.iterators import print_progress
 
@@ -23,21 +23,32 @@ def populate(self, context: ContentPopulatorContext) -> None:
         signals.post_save.disconnect(Package.post_save, sender=Package)
         signals.post_delete.disconnect(Package.post_delete, sender=Package)
 
+        uploaded_icon = None
+
         for i, package in print_progress(
             enumerate(context.packages), len(context.packages)
         ):
             vercount = package.versions.count()
             for vernum in range(context.version_count - vercount):
-                PackageVersionFactory.create(
+                pv = PackageVersion(
                     package=package,
                     name=package.name,
                     version_number=f"{vernum + vercount}.0.0",
                     website_url="https://example.org",
                     description=f"Example mod {i}",
                     readme=f"# This is an example mod number {i}",
                     changelog=f"# Example changelog for mod number {i}",
+                    file_size=5242880,
                 )
 
+                if context.reuse_icon and uploaded_icon:
+                    pv.icon = uploaded_icon
+                else:
+                    pv.icon = dummy_package_icon()
+
+                pv.save()
+                uploaded_icon = pv.icon.name
+
             # Manually calling would-be signals once per package, as it doesn't
             # actually make use of the sender param at all (and can be None)
             package.handle_created_version(None)

django/thunderstore/core/management/commands/create_test_data.py

@@ -72,6 +72,15 @@ def add_arguments(self, parser) -> None:
         parser.add_argument("--contract-version-count", type=int, default=8)
         parser.add_argument("--dependency-count", type=int, default=20)
         parser.add_argument("--clear", default=False, action="store_true")
+        parser.add_argument(
+            "--reuse-icon",
+            default=False,
+            action="store_true",
+            help=(
+                "Reuse the same icon file for all PackageVersions to speed up"
+                "creating large data sets."
+            ),
+        )
         parser.add_argument(
             "--only",
             type=str,
@@ -130,5 +139,6 @@ def handle(self, *args, **kwargs) -> None:
             contract_count=kwargs.get("contract_count", 0),
             contract_version_count=kwargs.get("contract_version_count", 0),
             wiki_page_count=kwargs.get("wiki_page_count", 0),
+            reuse_icon=kwargs.get("reuse_icon", False),
         )
         self.populate(context)

django/thunderstore/core/tests/test_management_commands.py

@@ -198,3 +198,38 @@ def test_create_test_data_only_filter_invalid() -> None:
         CommandError, match="Invalid --only selection provided, options are"
     ):
         call_command("create_test_data", "--only", "badfilter")
+
+
+@override_settings(DEBUG=True)
+@pytest.mark.django_db
+@pytest.mark.parametrize("reuse", (True, False))
+def test_create_test_data_reuse_icon(reuse: bool) -> None:
+    args = [
+        "create_test_data",
+        "--community-count",
+        1,
+        "--team-count",
+        1,
+        "--package-count",
+        1,
+        "--version-count",
+        2,
+        "--dependency-count",
+        0,
+        "--wiki-page-count",
+        0,
+        "--contract-count",
+        0,
+        "--contract-version-count",
+        0,
+    ]
+    if reuse:
+        args.append("--reuse-icon")
+
+    assert not PackageVersion.objects.exists()
+
+    call_command(*args)
+    icon_paths = PackageVersion.objects.values_list("icon", flat=True)
+
+    assert len(icon_paths) == 2
+    assert (icon_paths[0] == icon_paths[1]) == reuse