diff --git a/django/thunderstore/core/management/commands/content/base.py b/django/thunderstore/core/management/commands/content/base.py index c4fe1dda5..da149052f 100644 --- a/django/thunderstore/core/management/commands/content/base.py +++ b/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") diff --git a/django/thunderstore/core/management/commands/content/package_version.py b/django/thunderstore/core/management/commands/content/package_version.py index fe498f3fe..5d1975229 100644 --- a/django/thunderstore/core/management/commands/content/package_version.py +++ b/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,12 +23,14 @@ 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", @@ -36,8 +38,17 @@ def populate(self, context: ContentPopulatorContext) -> None: 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) diff --git a/django/thunderstore/core/management/commands/create_test_data.py b/django/thunderstore/core/management/commands/create_test_data.py index 745ccfa0d..f0955098c 100644 --- a/django/thunderstore/core/management/commands/create_test_data.py +++ b/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) diff --git a/django/thunderstore/core/tests/test_management_commands.py b/django/thunderstore/core/tests/test_management_commands.py index a3703c543..52909803d 100644 --- a/django/thunderstore/core/tests/test_management_commands.py +++ b/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