From 715021f7010ca4a007f4860502dbde05dbeb2c0e Mon Sep 17 00:00:00 2001
From: Rowan Seymour <rowanseymour@gmail.com>
Date: Thu, 9 Jan 2025 19:36:42 +0000
Subject: [PATCH] Data migration to backfill new category counts

---
 .../0355_backfill_new_cat_counts.py           | 32 +++++++++
 temba/flows/tests/test_migrations.py          | 65 +++++++++++++++++++
 2 files changed, 97 insertions(+)
 create mode 100644 temba/flows/migrations/0355_backfill_new_cat_counts.py
 create mode 100644 temba/flows/tests/test_migrations.py

diff --git a/temba/flows/migrations/0355_backfill_new_cat_counts.py b/temba/flows/migrations/0355_backfill_new_cat_counts.py
new file mode 100644
index 0000000000..45877b6234
--- /dev/null
+++ b/temba/flows/migrations/0355_backfill_new_cat_counts.py
@@ -0,0 +1,32 @@
+# Generated by Django 5.1.4 on 2025-01-09 18:35
+
+from django.db import migrations, transaction
+from django.db.models import Sum
+
+
+def backfill_new_cat_counts(apps, schema_editor):
+    Flow = apps.get_model("flows", "Flow")
+    FlowResultCount = apps.get_model("flows", "FlowResultCount")
+
+    for flow in Flow.objects.only("id"):
+        to_create = []
+
+        counts = flow.category_counts.values("result_key", "category_name").annotate(total=Sum("count"))
+        for c in counts:
+            if c["category_name"] and c["total"] > 0:
+                to_create.append(
+                    FlowResultCount(
+                        flow=flow, result=c["result_key"][:64], category=c["category_name"][:64], count=c["total"]
+                    )
+                )
+
+        with transaction.atomic():
+            flow.result_counts.all().delete()
+            FlowResultCount.objects.bulk_create(to_create)
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [("flows", "0354_flowresultcount")]
+
+    operations = [migrations.RunPython(backfill_new_cat_counts, migrations.RunPython.noop)]
diff --git a/temba/flows/tests/test_migrations.py b/temba/flows/tests/test_migrations.py
new file mode 100644
index 0000000000..5d8d7ca998
--- /dev/null
+++ b/temba/flows/tests/test_migrations.py
@@ -0,0 +1,65 @@
+from temba.tests import MigrationTest
+
+
+class BackfillNewCategoryCountsTest(MigrationTest):
+    app = "flows"
+    migrate_from = "0354_flowresultcount"
+    migrate_to = "0355_backfill_new_cat_counts"
+
+    def setUpBeforeMigration(self, apps):
+        self.flow1 = self.create_flow("Flow 1")
+        self.flow2 = self.create_flow("Flow 2")
+
+        self.flow1.category_counts.create(
+            node_uuid="0043bf6f-f385-4ba3-80d7-4313771efa86",
+            result_key="color",
+            result_name="Color",
+            category_name="Red",
+            count=1,
+        )
+        self.flow1.category_counts.create(
+            node_uuid="2c230f54-a7f6-4a71-9f24-b3ba81734552",
+            result_key="color",
+            result_name="Color",
+            category_name="Red",
+            count=2,
+        )
+        self.flow1.category_counts.create(
+            node_uuid="0043bf6f-f385-4ba3-80d7-4313771efa86",
+            result_key="color",
+            result_name="Color",
+            category_name="Blue",
+            count=4,
+        )
+        self.flow1.category_counts.create(
+            node_uuid="f27d77e9-4ccb-4c36-8277-f5708822cf26",
+            result_key="name",
+            result_name="Name",
+            category_name="All Responses",
+            count=6,
+        )
+        self.flow2.category_counts.create(
+            node_uuid="daf35e03-b82e-4c12-8ff7-60f7d990eb05",
+            result_key="thing",
+            result_name="Thing",
+            category_name="Looooonnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnggggggggggggnnggggggggggggggggg",
+            count=2,
+        )
+        self.flow2.category_counts.create(
+            node_uuid="daf35e03-b82e-4c12-8ff7-60f7d990eb05",
+            result_key="thing",
+            result_name="Thing",
+            category_name="Zero",
+            count=0,
+        )
+
+    def test_migration(self):
+        def assert_counts(flow, expected):
+            actual = {}
+            for count in flow.result_counts.all():
+                actual[f"{count.result}/{count.category}"] = count.count
+
+            self.assertEqual(actual, expected)
+
+        assert_counts(self.flow1, {"color/Red": 3, "color/Blue": 4, "name/All Responses": 6})
+        assert_counts(self.flow2, {"thing/Looooonnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnnggggggggggggnngggggggggg": 2})