From 517276c6ff8af554196863044176e55abe9e887d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Antti=20M=C3=A4ki?= <antti@anttimaki.com>
Date: Fri, 6 Oct 2023 11:03:17 +0300
Subject: [PATCH] Add Celery task for updating CommunityAggregatedFields

Refs TS-1850
---
 django/thunderstore/community/tasks.py        | 18 +++++++
 .../tests/test_community_aggregated_fields.py | 53 +++++++++++++++++++
 django/thunderstore/core/tests/test_celery.py |  1 +
 3 files changed, 72 insertions(+)
 create mode 100644 django/thunderstore/community/tasks.py

diff --git a/django/thunderstore/community/tasks.py b/django/thunderstore/community/tasks.py
new file mode 100644
index 000000000..c0f03c3bc
--- /dev/null
+++ b/django/thunderstore/community/tasks.py
@@ -0,0 +1,18 @@
+import logging
+from typing import Optional
+
+from celery import shared_task
+
+from thunderstore.community.models import Community, CommunityAggregatedFields
+from thunderstore.core.settings import CeleryQueues
+
+logger = logging.getLogger(__name__)
+
+
+@shared_task(queue=CeleryQueues.BackgroundTask)
+def update_community_aggregated_fields(batch_size: Optional[int] = None) -> None:
+    communities = Community.objects.all()
+    logger.info(f"Updating aggregated fields for {communities.count()} communities")
+
+    CommunityAggregatedFields.update_for_communities(communities, batch_size)
+    logger.info("CommunityAggregatedFields updated")
diff --git a/django/thunderstore/community/tests/test_community_aggregated_fields.py b/django/thunderstore/community/tests/test_community_aggregated_fields.py
index 383253546..60b72032b 100644
--- a/django/thunderstore/community/tests/test_community_aggregated_fields.py
+++ b/django/thunderstore/community/tests/test_community_aggregated_fields.py
@@ -5,6 +5,8 @@
 from thunderstore.community.consts import PackageListingReviewStatus
 from thunderstore.community.factories import CommunityFactory, PackageListingFactory
 from thunderstore.community.models import Community, CommunityAggregatedFields
+from thunderstore.community.tasks import update_community_aggregated_fields
+from thunderstore.repository.factories import PackageVersionFactory
 
 
 @pytest.mark.django_db
@@ -219,3 +221,54 @@ def test_community_aggregated_fields__update_for_communities__skips_inactive_ver
 
     assert listing.community.aggregated.package_count == 1
     assert listing.community.aggregated.download_count == 2
+
+
+@pytest.mark.django_db
+def test_community_aggregated_fields__celery_tasks__handles_mixed_situations():
+    # Community 1 has existing packages and downloads.
+    caf1 = CommunityAggregatedFields.objects.create(package_count=2, download_count=5)
+    c1 = CommunityFactory(aggregated_fields=caf1)
+    PackageListingFactory(
+        community_=c1,
+        package_version_kwargs={"downloads": 1},
+    )
+    c1_l2 = PackageListingFactory(
+        community_=c1,
+        package_version_kwargs={"downloads": 2, "version_number": "1.0.1"},
+    )
+    c1_l2_v2 = PackageVersionFactory(package=c1_l2.package, downloads=2)
+
+    # Community 1 has changes not reflected in the initial aggregated values.
+    c1_l2_v2.is_active = False
+    c1_l2_v2.save()
+    PackageListingFactory(community_=c1)
+
+    # Community 2 has existing CommunityAggregatedFields but no packages.
+    caf2 = CommunityAggregatedFields.objects.create()
+    CommunityFactory(aggregated_fields=caf2)
+
+    # Community 3 has no pre-existing CommunityAggregatedFields.
+    c3_l1 = PackageListingFactory(
+        package_version_kwargs={"downloads": 3},
+    )
+
+    assert Community.objects.count() == 3
+    assert CommunityAggregatedFields.objects.count() == 2
+    assert caf1.package_count == 2
+    assert caf1.download_count == 5
+    assert caf2.package_count == 0
+    assert caf2.download_count == 0
+
+    update_community_aggregated_fields()
+    caf1.refresh_from_db()
+    caf2.refresh_from_db()
+    caf3 = Community.objects.get(pk=c3_l1.community.pk).aggregated_fields
+
+    assert Community.objects.count() == 3
+    assert CommunityAggregatedFields.objects.count() == 3
+    assert caf1.package_count == 3
+    assert caf1.download_count == 3
+    assert caf2.package_count == 0
+    assert caf2.download_count == 0
+    assert caf3.package_count == 1
+    assert caf3.download_count == 3
diff --git a/django/thunderstore/core/tests/test_celery.py b/django/thunderstore/core/tests/test_celery.py
index 11a227ccc..d93aa96d7 100644
--- a/django/thunderstore/core/tests/test_celery.py
+++ b/django/thunderstore/core/tests/test_celery.py
@@ -28,6 +28,7 @@ def test_task():
     "celery.chain",
     "celery.starmap",
     "celery.backend_cleanup",
+    "thunderstore.community.tasks.update_community_aggregated_fields",
     "thunderstore.core.tasks.celery_post",
     "thunderstore.cache.tasks.invalidate_cache",
     "thunderstore.repository.tasks.update_api_caches",